发现变压器模型各模块尖锐度差异，按模块调整学习率可提速近2倍。
The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
- 按模块尖锐度差异动态调节学习率，提升训练效率。
- 相比原版AdamW，终端损失更低，训练速度提升近2倍。
- 适用于0.12B至2B模型，适合追求高效训练的开发者。
Transformer由嵌入层、归一化层、自注意力机制和前馈网络等不同模块构成，理解其差异与交互至关重要。本文首次揭示这些模块间存在显著的尖锐度差异，该现象在训练初期即出现并贯穿全程。受此启发，提出分块学习率（Blockwise LR）策略，根据各模块的尖锐度定制学习率，显著加速大语言模型预训练。将该策略集成至AdamW后，在GPT-2与LLaMA系列模型上（0.12B至2B规模），使用OpenWebText、MiniPile和C4数据集，均实现终端损失更低、训练速度接近2倍提升。进一步将其应用于内存高效的Adam-mini（Zhang et al., 2024），实现2倍加速与2倍内存节省。结果表明，利用尖锐度差异可有效提升语言模型训练效率。
原文摘要 · Abstract (English)
Transformers consist of diverse building blocks, such as embedding layers, normalization layers, self-attention mechanisms, and point-wise feedforward networks. Thus, understanding the differences and interactions among these blocks is important. In this paper, we uncover a clear Sharpness Disparity across these blocks, which emerges early in training and intriguingly persists throughout the training process. Motivated by this finding, we propose Blockwise Learning Rate (LR), a strategy that tailors the LR to each block's sharpness, accelerating large language model (LLM) pre-training. By integrating Blockwise LR into AdamW, we consistently achieve lower terminal loss and nearly $2\times$ speedup compared to vanilla AdamW. We demonstrate this acceleration across GPT-2 and LLaMA, with model sizes ranging from 0.12B to 2B and datasets of OpenWebText, MiniPile, and C4. Finally, we incorporate Blockwise LR into Adam-mini (Zhang et al., 2024), a recently proposed memory-efficient variant of Adam, achieving a combined $2\times$ speedup and $2\times$ memory saving. These results underscore the potential of exploiting the sharpness disparity to improve LLM training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。