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

TL;DR

GradLite is a novel optimizer that enables memory-efficient training of large language models by using approximate gradients with error correction, maintaining convergence and performance without changing model architecture.

Contribution

Introduces GradLite, a backward-friendly optimizer that reduces memory usage through low-rank Jacobian approximation and error-feedback, without altering model architecture.

Findings

Reduces optimizer and activation memory by up to 50%.

Maintains convergence guarantees comparable to Adam.

Achieves competitive or superior performance on multiple benchmarks.

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).