不用Adam,初始化时调学习率就能让SGD轻松训练Transformer。
No More Adam: Learning Rate Scaling at Initialization is All You Need
- 在初始化时按梯度信噪比调整学习率,避免训练失衡。
- 训练ViT和GPT-2时性能媲美甚至超越AdamW,且对超参不敏感。
- 内存占用减半,适合大模型训练与低资源场景。
本文质疑了深度神经网络训练中自适应梯度方法的必要性。SGD-SaI 是一种对随机梯度下降带动量(SGDM)的简单而有效的改进,通过在初始化时对不同参数组进行学习率缩放(SaI),依据其梯度信噪比(g-SNR)进行调整。该方法无需依赖自适应二阶动量,在首个迭代起即防止训练不平衡,并将优化器内存消耗降低一半,相比AdamW。尽管结构简单高效,SGD-SaI在多种基于Transformer的任务中表现稳定,性能可匹敌或超越AdamW,有效解决了长期存在的使用SGD训练Transformer的难题。在ImageNet-1K图像分类(ViT)和大语言模型(LLM,Transformer解码器仅用)的GPT-2预训练任务中表现出色,对超参数变化具有鲁棒性,具备广泛适用性。进一步测试显示,在LoRA微调和扩散模型等任务中,其性能持续优于当前最优优化器。从内存效率看,全精度训练下,相较于AdamW,SGD-SaI在GPT-2(15亿参数)上节省5.93 GB内存,在Llama2-7B上节省25.15 GB内存。
原文摘要 · Abstract (English)
In this work, we question the necessity of adaptive gradient methods for training deep neural networks. SGD-SaI is a simple yet effective enhancement to stochastic gradient descent with momentum (SGDM). SGD-SaI performs learning rate Scaling at Initialization (SaI) to distinct parameter groups, guided by their respective gradient signal-to-noise ratios (g-SNR). By adjusting learning rates without relying on adaptive second-order momentum, SGD-SaI helps prevent training imbalances from the very first iteration and cuts the optimizer's memory usage by half compared to AdamW. Despite its simplicity and efficiency, SGD-SaI consistently matches or outperforms AdamW in training a variety of Transformer-based tasks, effectively overcoming a long-standing challenge of using SGD for training Transformers. SGD-SaI excels in ImageNet-1K classification with Vision Transformers(ViT) and GPT-2 pretraining for large language models (LLMs, transformer decoder-only), demonstrating robustness to hyperparameter variations and practicality for diverse applications. We further tested its robustness on tasks like LoRA fine-tuning for LLMs and diffusion models, where it consistently outperforms state-of-the-art optimizers. From a memory efficiency perspective, SGD-SaI achieves substantial memory savings for optimizer states, reducing memory usage by 5.93 GB for GPT-2 (1.5B parameters) and 25.15 GB for Llama2-7B compared to AdamW in full-precision training settings.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。