SAGE: Sign-Adaptive Gradient for Memory-Efficient LLM Optimization

2026-04-08

AI Summary

This work addresses the high memory overhead of the AdamW optimizer in large language model training, where existing lightweight methods struggle to effectively handle the sparse and high-variance gradients of embedding layers, often necessitating hybrid use with AdamW and limiting memory savings. To overcome this, the authors propose SAGE, the first optimizer capable of fully replacing AdamW. SAGE combines Lion-style sign-based momentum updates with a novel “safe damping” adaptive scaling mechanism that provably bounds gradient amplification (with an upper bound ≤1.0), effectively stabilizing training under high gradient variance while maintaining O(d) space complexity. Evaluated on a 1.3B-parameter Llama model, SAGE substantially reduces optimizer memory usage and achieves state-of-the-art perplexity, outperforming all baselines—including SinkGD hybrid approaches.

Technology Category

Search and Optimization: Learning to SearchMachine Learning: OptimizationNatural Language Processing: Learning & Optimization for NLP

Application Category

Search and Retrieval-Augmented AI: Large language models for searchSemantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMsGraph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• memory bottleneck
• optimizer state
• embedding layer
• high-variance gradients
• large language models

Innovation

• SAGE
• memory-efficient optimization
• adaptive gradient
• embedding layer dilemma
• LLM pretraining