STEP: Staged Parameter-Efficient Pre-training for Large Language Models

📄 arXiv: 2504.04151v1 📥 PDF

作者: Kazuki Yano, Takumi Ito, Jun Suzuki

分类: cs.CL

发布日期: 2025-04-05

备注: Accepted to NAACL 2025 Main

💡 一句话要点

提出STEP：一种用于大规模语言模型的阶段式参数高效预训练方法

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 参数高效预训练 大规模语言模型 模型增长 内存优化 指令微调

📋 核心要点

  1. 大规模语言模型预训练面临着巨大的内存挑战，因为模型参数量非常庞大。
  2. STEP方法的核心思想是将参数高效微调技术与模型增长相结合，从而降低内存需求。
  3. 实验结果表明，STEP在保持性能的同时，显著降低了内存需求，并在下游任务上表现良好。

📝 摘要（中文）

本文提出了一种名为STEP（STaged parameter-Efficient Pre-training）的阶段式参数高效预训练方法，旨在解决大规模语言模型（LLM）预训练过程中因模型参数量巨大而带来的内存挑战。STEP将参数高效微调技术与模型增长相结合。在不同规模的LLM预训练实验中，结果表明，与传统的预训练方法相比，STEP在保持同等性能的同时，最多可减少53.9%的最大内存需求。此外，经过指令微调后，STEP训练的模型在下游任务上的表现与传统预训练模型相当。

🔬 方法详解

问题定义：大规模语言模型（LLM）的预训练由于其庞大的参数量，对计算资源提出了极高的要求，特别是内存需求。传统的预训练方法需要消耗大量的GPU内存，限制了模型规模的扩展和训练效率的提升。现有方法在降低内存消耗方面存在不足，例如无法在不损失性能的前提下显著减少内存占用。

核心思路：STEP的核心思路是通过阶段性的参数高效预训练，逐步增加模型容量，并结合参数高效微调技术，在每个阶段只训练少量参数，从而显著降低内存需求。这种方法允许模型在有限的资源下进行预训练，并最终达到与传统预训练方法相当的性能。

技术框架：STEP方法包含多个阶段。在每个阶段，首先使用参数高效微调技术（如LoRA、Adapter等）对模型进行训练，只更新少量参数。然后，逐步增加模型的容量，例如增加Transformer层的数量或扩大词汇表。每个阶段的训练都基于前一个阶段的模型，并使用参数高效微调技术进行优化。最终，得到一个具有完整容量的预训练模型。

关键创新：STEP的关键创新在于将参数高效微调技术与模型增长相结合，实现了一种阶段性的预训练方法。与传统的预训练方法相比，STEP能够在显著降低内存需求的同时，保持模型的性能。此外，STEP还允许在有限的资源下训练更大规模的模型。

关键设计：STEP的关键设计包括：1) 选择合适的参数高效微调技术，例如LoRA或Adapter，以最小化内存占用；2) 设计合理的模型增长策略，例如逐步增加Transformer层的数量或扩大词汇表；3) 确定每个阶段的训练时长和学习率等超参数，以保证模型的收敛性和性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，STEP方法在预训练LLM时，与传统的预训练方法相比，最多可减少53.9%的最大内存需求，同时保持了相当的性能。经过指令微调后，STEP训练的模型在下游任务上的表现与传统预训练模型相当。这些结果验证了STEP方法的有效性和实用性。

🎯 应用场景

STEP方法可应用于各种需要大规模语言模型的场景，例如自然语言处理、机器翻译、文本生成等。它尤其适用于资源受限的环境，例如在低端GPU或云服务器上进行模型预训练。通过降低内存需求，STEP可以加速LLM的开发和部署，并促进其在更广泛领域的应用。未来，该方法有望进一步扩展到其他类型的深度学习模型。

📄 摘要（原文）

Pre-training large language models (LLMs) faces significant memory challenges due to the large size of model parameters. We introduce STaged parameter-Efficient Pre-training (STEP), which integrates parameter-efficient tuning techniques with model growth. We conduct experiments on pre-training LLMs of various sizes and demonstrate that STEP achieves up to a 53.9% reduction in maximum memory requirements compared to vanilla pre-training while maintaining equivalent performance. Furthermore, we show that the model by STEP performs comparably to vanilla pre-trained models on downstream tasks after instruction tuning.