SAGE优化器让大模型训练更省内存，同时提升收敛效果。
SAGE: Sign-Adaptive Gradient for Memory-Efficient LLM Optimization
- 用符号自适应机制替代AdamW，降低内存占用
- 在1.3B参数Llama模型上达到新最佳困惑度
- 适合追求高效训练的大模型研究者
AdamW优化器虽为大模型预训练标准，但其优化器状态占用相当于模型大小两倍，构成关键内存瓶颈。尽管轻量级优化器如SinkGD试图缓解此问题，我们发现嵌入层存在梯度稀疏且方差高的困境：现有方法无法有效处理，迫使采用混合设计回退至AdamW，部分抵消内存优势。为此提出SAGE（Sign Adaptive GradiEnt），通过结合Lion式更新方向与新型$O(d)$自适应尺度，实现内存高效的优化。该尺度作为“安全阻尼器”，理论上上限为1.0，能更有效地抑制高方差维度，提升稳定性。实验表明，基于SAGE的混合方案在高达1.3B参数的Llama模型上，实现了新的最优困惑度，超越所有基线（包括SinkGD混合方案），同时显著减少优化器状态内存消耗。
原文摘要 · Abstract (English)
The AdamW optimizer, while standard for LLM pretraining, is a critical memory bottleneck, consuming optimizer states equivalent to twice the model's size. Although light-state optimizers like SinkGD attempt to address this issue, we identify the embedding layer dilemma: these methods fail to handle the sparse, high-variance gradients inherent to embeddings, forcing a hybrid design that reverts to AdamW and partially negates the memory gains. We propose SAGE (Sign Adaptive GradiEnt), a novel optimizer that resolves this dilemma by replacing AdamW in this hybrid structure. SAGE combines a Lion-style update direction with a new, memory-efficient $O(d)$ adaptive scale. This scale acts as a "safe damper," provably bounded by 1.0, which tames high-variance dimensions more effectively than existing methods. This superior stability allows SAGE to achieve better convergence. On Llama models up to 1.3B parameters, our SAGE-based hybrid achieves new state-of-the-art perplexity, outperforming all baselines, including SinkGD hybrid, while significantly reducing optimizer state memory.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。