The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
作者: Jinbo Wang, Mingze Wang, Zhanpeng Zhou, Junchi Yan, Weinan E, Lei Wu
发布时间: 2025-02-27
来源: arxiv
研究方向: 自然语言处理（NLP）
主要内容
该论文研究了Transformer架构中不同模块的锐度差异，并提出了一种名为Blockwise Learning Rate（BLR）的策略，旨在通过调整每个模块的学习率来加速大型语言模型（LLM）的预训练过程。
主要贡献
1. 揭示了Transformer中不同模块之间的锐度差异，并提出了Blockwise LR策略。
2. 将Blockwise LR集成到AdamW和Adam-mini优化器中，实现了显著的加速效果。
3. 在GPT-2和LLaMA模型上展示了BLR的有效性，模型大小从0.12B到1.1B，数据集包括OpenWebText和MiniPile。
4. 将BLR集成到Adam-mini中，实现了加速和内存节省的双重效果。
研究方法
1. Hessian矩阵分析
2. Blockwise LR策略
3. AdamW和Adam-mini优化器
4. 实验评估
实验结果
实验结果表明，使用BLR策略的AdamW和Adam-mini在LLM预训练任务中均取得了显著的加速效果，终端损失更低，训练速度提升了近2倍。
未来工作
将BLR策略应用于非LLM任务，例如计算机视觉，并与其他优化器结合使用，如Muon和Mamba。此外，研究其他自适应优化策略，如Blockwise weight decay和gradient clipping，以进一步提高训练效率和性能。