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

2025-10-25

AI Summary

Large language models (LLMs) suffer from prohibitive memory overhead during full fine-tuning, primarily due to storing complete activations for exact gradient computation. This work introduces the first optimizer-centric lightweight training framework that eliminates the need for exact gradients. Our method employs efficient low-rank Jacobian approximation to estimate gradients and incorporates an error-feedback correction mechanism to ensure unbiased estimation with bounded variance. Crucially, it is architecture-agnostic—requiring no model structural modifications—and compatible with standard LLMs. Theoretical analysis guarantees convergence properties, while empirical evaluation across diverse downstream tasks demonstrates stable convergence and performance on par with or exceeding conventional baselines. Memory consumption for optimizer states and activations is reduced by up to 50%. Our core contribution lies in decoupling optimization from exact gradient computation, establishing a novel paradigm for memory-efficient full fine-tuning of LLMs.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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

Problem

• Reduces memory consumption in large language model training
• Enables efficient training with approximate gradients under constraints
• Maintains convergence guarantees while discarding intermediate activations

Innovation

• Uses low-rank Jacobian approximation for gradients
• Employs error-feedback correction for convergence
• Reduces memory usage without architectural changes