Backward-Friendly Optimization: Training Large Language Models with Approximate Gradients under Memory Constraints

TL;DR

This work tackles the memory bottleneck of full fine-tuning large language models by designing GradLite, a backward-friendly optimizer that tolerates approximate gradients. It achieves memory efficiency through a dual mechanism: low-rank Jacobian approximation to compress backpropagated signals and error-feedback correction to offset approximation residuals, with theoretical guarantees of unbiased updates and $O(1/\sqrt{T})$ convergence. Empirically, GradLite reduces optimizer-state and activation memory by up to ~50% on large models while matching or surpassing checkpointing and optimizer-based baselines on reasoning, multilingual, and dialogue benchmarks. The results indicate that optimizer-level innovations can meaningfully complement architectural and system-level approaches for memory-efficient LLM training, enabling full-parameter updates under tighter hardware constraints.

Abstract

Full fine-tuning of Large Language Models (LLMs) is notoriously memory-intensive, primarily because conventional optimizers such as SGD or Adam assume access to exact gradients derived from cached activations. Existing solutions either alter the model architecture (e.g., reversible networks) or trade memory for computation (e.g., activation checkpointing), but the optimizer itself remains untouched. In this work, we introduce GradLite, a backward-friendly optimizer that relaxes the requirement of exact gradients, enabling efficient training even when intermediate activations are aggressively discarded or approximated. GradLite leverages two key techniques: (i) low-rank Jacobian approximation, which reduces the dimensionality of backpropagated error signals, and (ii) error-feedback correction, which accumulates and compensates approximation errors across iterations to preserve convergence guarantees. We provide a theoretical analysis showing that GradLite maintains unbiased gradient estimates with bounded variance, ensuring convergence rates comparable to Adam. Empirically, GradLite reduces optimizer-state and activation memory consumption by up to 50\% without architectural changes, and achieves on-par or superior downstream performance on reasoning (MMLU, GSM8K), multilingual, and dialogue benchmarks compared to checkpointing and optimizer-centric baselines (LoMo, GaLore).

Figures (2)

• Figure 1: Comparison of update mechanisms. Left: Standard RevFFN model. Top Right (Ours): GradLite uses approximate Jacobian with residuals to cut memory. Bottom Right: Standard optimizer requires full activations.
• Figure 2: GradLite workflow. Solid path computes $\hat{g}_t$ from $\tilde{g}_t$ with $\bm{r}_t$; dashed path shows exact $g_t$.