用于加速语言模型预训练的Transformer中锐度差异原理
The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Training
AI总结：
本研究发现Transformer各模块间存在持续的锐度差异，据此提出分块学习率策略，适配各模块锐度以加速LLM预训练，结合AdamW等优化器实现训练提速与内存节省，为优化大模型训练提供新思路。
摘要：
Transformer由多种不同的构建模块组成，例如嵌入层、归一化层、自注意力机制以及逐点前馈网络。因此，理解这些模块之间的差异与相互作用十分重要。本文中，我们发现这些模块间存在明显的锐度差异（Sharpness Disparity），该差异在训练早期就会出现，且有趣的是会在整个训练过程中持续存在。受这一发现启发，我们提出了分块学习率（Blockwise LR）策略，该策略根据每个模块的锐度定制学习率，以加速大语言模型（LLM）的预训练。将分块学习率整合到AdamW中后，与普通AdamW相比，我们始终能获得更低的最终损失，且训练速度提升近2倍。我们在GPT-2和LLaMA模型上验证了该加速效果，模型规模从0.12B到2B不等，所用数据集包括OpenWebText、MiniPile和C4。最后，我们将分块学习率整合到Adam-mini（Zhang等人，2024）——一种近期提出的内存高效型Adam变体中，实现了2倍速度提升与2倍内存节省的双重效果。这些结果凸显了利用锐度差异改进LLM训练的潜力。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
