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

Wooin Lee; Hyun-Tae Kim

cs.LGApr 9, 2026

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

Authors: Wooin Lee, Hyun-Tae Kim

Organizations: London, United Kingdom · Seoul, Republic of Korea

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.

Figures & tables

Figure 1: Test Perplexity for SinkGD-Pure vs. SinkGD-Hybrid. Even with tuned learning rates, the pure-stateless optimizer fails to learn effective embedding representations, resulting in significantly worse final perplexity compared to the hybrid design.

Model Component	Optimizer	State Size
Embedding Layer	SAGE	$O(Vd)+O(d)$
Biases/Norms (1D)	SAGE	$2\times O(d)$
Dense Weights (2D)	SinkGD	$O(1)$

Table 1: Composition of the proposed SAGE-Hybrid optimizer. By using SAGE for embeddings instead of AdamW (which requires

2\times O(Vd)

states), we reduce the dominant memory cost by

\approx 50\%

while maintaining stability.

Figure 2: Convergence and efficiency on the 1.3B model. (a) Test perplexity - lower the better; SAGE achieves the lowest final perplexity and faster convergence. (b) Throughput - higher the better; Effective throughput normalizes processing speed by the steps required to match the lowest (Lion-Hybrid) performance. (c) Peak memory usage breakdown by component.

Group	Method	270M		0.6B		1.3B
Group	Method	PPL	Mem (GB)	PPL	Mem (GB)	PPL	Mem (GB)
Baselines	AdamW	37.35	2.1	31.74	4.4	27.81	9.8
	APOLLO	52.83	1.2	44.84	2.1	39.54	3.6
	Lion	30.24	1.0	26.58	2.2	28.37	4.9
	SinkGD-Hybrid	34.30	0.9	32.85	1.5	28.71	1.9
Ablations	SinkGD-Pure	192.7	0.0	117.8	0.0	98.84	0.0

Table 2: Main results comparing final test perplexity (PPL) and peak optimizer memory usage (Mem) across three model scales. SAGE consistently achieves the best perplexity while maintaining a low memory footprint, outperforming both pure and hybrid baselines. Lower PPL is better.

Figure 3: Visualization of the SAGE adaptive scale $\mathbf{H}_{t}$ (for $d=1024$ dimensions) over the first 500 (left) and 5k (right) training steps. The raw $\mathbf{H}_{t}\in(0,1.0]$ values are log-transformed and normalized. The white background (value 1.0) corresponds to the majority of "quiet" dimensions where the gradient variance is below the layer average, causing SAGE to clip the scale to 1.0 (Lion’s default). The colored stripes (low values $\ll 1.0$ ) correspond to specific "loud" high-variance dimensions where SAGE actively applies strong damping. This confirms SAGE acts as a sparse damper, intervening only on specific unstable dimensions while letting the majority update at full scale.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Configuration	270M	0.6B	1.3B
Hidden Size	1024	1536	2048
Intermediate Size	2816	4224	5632
Max Position Embeddings	8192	8192	8192
Num Attention Heads	16	24	32
Num Hidden Layers	13	16	24
Num Key-Value Heads	2	3	4

Appendix

Table 3: Model configurations for the different scales used in the experiments.

	270M		0.6B		1.3B
	(Emb. Params: 48%)		(Emb. Params: 33%)		(Emb. Params: 20%)
Optimizer	Opt. (GB)	Total (GB)	Opt. (GB)	Total (GB)	Opt. (GB)	Total (GB)
AdamW	2.045	4.091	4.421	8.843	9.833	19.67
APOLLO	1.208	3.254	2.102	6.524	3.648	13.48
Lion	1.023	3.067	2.211	6.631	4.916	14.75
SinkGD-Hybrid	0.979	3.023	1.468	5.890	1.958	11.79

Appendix

Table 4: Peak memory footprint comparison across all model scales. "Opt." refers to optimizer state memory, while "Total" is the sum of model weights, gradients and the optimizer states. SAGE is comparable to the low total memory of a pure stateless optimizer SinkGD-Pure , offering a substantial reduction compared to the AdamW-based hybrid, which translates to significant savings in total memory.

Figure 4: PCA of the adaptive scale vector

\mathbf{H}_{t}

. (a) Trajectories of the top 3 principal components. PC1 and PC2 show a sharp, distinct spike in activity during the early training phase (initial 1k steps), reflecting rapid adaptation to the embedding layer’s geometry. This is followed by a gradual stabilization phase where components settle into steady states. (b) The trajectory in PC1-PC2 space reveals a clear "hook" pattern: a rapid initial excursion away from the origin (adaptation) followed by a smooth return path that converges toward a stable operating point (stabilization).

LLM Input Optimization

See all 22 papers

Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers
May 9, 2026Aditya RanganathDeep Learning Optimization LLM Input Optimization
Center for Applied Scientific Computing Lawrence Livermore National Laboratory Livermore, CA, USA
Revisiting the Adam-SGD Gap in LLM Pre-Training: The Role of Large Effective Learning Rates
May 18, 2026Athanasios Glentis, Dawei Li, Chung-Yiu Yau +1Language Model Pretraining Gradient Clipping
University of Minnesota
AdamPM: Memory-Efficient LLM Optimization via Moment Pooling
Oct 17, 2026Xiaowei Ji, Lishan Xu, Jianing An +3Efficient Language Model Training Memory-Efficient Fine-Tuning
SKLCCSE, School of Artificial Intelligence, Beihang University, Beijing, China · Hangzhou International Innovation Institute, Beihang University, Hangzhou, China

May 9, 2026J/K move · Enter open · S save

Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers

Aditya Ranganath

Center for Applied Scientific Computing Lawrence Livermore National Laboratory Livermore, CA, USA

Deep Learning Optimization LLM Input Optimization Memory-Efficient Optimization Matrix Optimization LLM Training

Open paper AskPDF arXiv

Training large language models requires optimization algorithms that are not only statistically effective, but also computationally and memory efficient at extreme scale. Although Adam remains the dominant optimizer for large-scale language-model pretraining and fine-tuning, recent work has revisited nearly every component of the optimization stack: adaptive moment estimation, decoupled weight decay, memory footprint, curvature approximation, sign-based updates, large-batch stability, low-rank gradient structure, and matrix-wise orthogonalized updates. This survey reviews optimizer design for large language models through a systems-and-optimization lens. We organize the literature into classical first-order optimizers, adaptive optimizers, memory-efficient variants, second-order and curvature-aware methods, sign-based and discovered optimizers, low-rank and projection-based methods, and matrix-based optimizers such as Muon. We also discuss benchmarking methodology, including hyperparameter fairness, scale dependence, wall-clock efficiency, token efficiency, memory overhead, and downstream evaluation. We argue that optimizer research for LLMs is entering a new phase: moving from single-algorithm speedup claims toward rigorous, scale-aware comparisons that jointly evaluate convergence, stability, memory, and implementation complexity.

Abstract

LLM Input Optimization

Explore similar work

Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers

Revisiting the Adam-SGD Gap in LLM Pre-Training: The Role of Large Effective Learning Rates

AdamPM: Memory-Efficient LLM Optimization via Moment Pooling