Optimal low-rank stochastic gradient estimation for LLM training

• URL: http://arxiv.org/abs/2603.20632v1
• Date: Sat, 21 Mar 2026 04:03:32 GMT
• Title: Optimal low-rank stochastic gradient estimation for LLM training
• Abstract summary: Large language model (LLM) training is often bottlenecked by memory constraints and gradient noise.<n>We present an unbiased, memory-efficient, low-rank matrix estimator with the lowest variance.
• Score: 8.22982038180272
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Large language model (LLM) training is often bottlenecked by memory constraints and stochastic gradient noise in extremely high-dimensional parameter spaces. Motivated by empirical evidence that many LLM gradient matrices are effectively low-rank during training, we present an unbiased, memory-efficient, low-rank matrix estimator with the lowest variance that is applicable across common stochastic gradient estimation paradigms. The core idea is to project a high-dimensional stochastic gradient estimator onto a random low-dimensional subspace and lift it back, reducing memory while keeping the estimator unbiased and controlling mean-squared error via an optimally designed projection distribution, including Haar--Stiefel projections. The projection distribution is derived by solving a constrained functional optimization problem, yielding an optimal random projector that guides algorithm design. Empirically, the resulting low-rank gradient estimators deliver both practical memory savings and improved training behavior. In RoBERTa-large fine-tuning, our method attains the lowest peak GPU memory among compared methods (e.g., 3.83GB versus 16.7GB for full BP) while remaining competitive in accuracy; in autoregressive LLM pretraining (LLaMA-20M/60M/100M), our method outperforms the traditional methods, supporting the benefit of the proposed optimal projection strategy.

Related papers

• Prior-Informed Zeroth-Order Optimization with Adaptive Direction Alignment for Memory-Efficient LLM Fine-Tuning [4.278794376089146]
  We propose a plug-and-play method that incorporates prior-informed perturbations to refine gradient estimation.<n>Our method significantly accelerates convergence compared to standard ZO approaches.<n>We prove that our gradient estimator achieves stronger alignment with the true gradient direction.
  arXiv  Detail & Related papers  (2026-01-08T08:27:15Z)
• Low-Rank Curvature for Zeroth-Order Optimization in LLM Fine-Tuning [8.349781300731225]
  We introduce LOREN, a curvature-aware zeroth-order (ZO) optimization method for fine-tuning large language models (LLMs)<n>Existing ZO methods, which estimate gradients via finite differences using random perturbations, often suffer from high variance and suboptimal search directions.<n>Our approach addresses these challenges by: (i) adaptively estimating an anisotropic perturbation distribution for gradient estimation, (ii) capturing curvature through a low-rank block diagonal preconditioner, and (iii) applying a REINFORCE leave-one-out (RLOO) gradient estimator to reduce variance.
  arXiv  Detail & Related papers  (2025-11-11T08:34:09Z)
• Towards Fast LLM Fine-tuning through Zeroth-Order Optimization with Projected Gradient-Aligned Perturbations [23.409093103129706]
  Fine-tuning large language models (LLMs) using zeroth-order (ZO) optimization has emerged as a promising alternative to traditional gradient-based methods.<n>Existing ZO methods suffer from high variance in gradient estimation, leading to slow convergence and suboptimal performance on large-scale models.<n>We propose P-GAP, a fast LLM fine-tuning approach through zeroth-order optimization with Projected Gradient-Aligned Perturbations.
  arXiv  Detail & Related papers  (2025-10-21T02:19:11Z)
• Unbiased Gradient Low-Rank Projection [32.57081286181632]
  A popular strategy involves gradient low-rank projection, storing only the projected states, with GaLore being a representative example.<n>This paper investigates the layerwise sampling technique for debiasing low-rank projection mechanisms.<n>An instantiation of the paradigm gives rise to a novel and unbiased low-rank optimization method built upon GaLore's mechanism and the Muon algorithm.
  arXiv  Detail & Related papers  (2025-10-20T17:59:25Z)
• From Noisy Traces to Stable Gradients: Bias-Variance Optimized Preference Optimization for Aligning Large Reasoning Models [90.45197506653341]
  Large reasoning models generate intermediate reasoning traces before producing final answers.<n> aligning LRMs with human preferences, a crucial prerequisite for model deployment, remains underexplored.<n>A common workaround optimized a single sampled trajectory, which introduces substantial gradient variance from trace sampling.
  arXiv  Detail & Related papers  (2025-10-06T17:58:01Z)
• Memory-Efficient LLM Training by Various-Grained Low-Rank Projection of Gradients [36.31328244962629]
  Low-rank gradient projection (LoRP) has emerged as a promising solution for memory-efficient fine-tuning.<n>Existing LoRP methods treat each row of the gradient matrix as the default projection unit, leaving the role of projection granularity underexplored.<n>We propose a novel framework, VLoRP, that extends low-rank gradient projection by introducing an additional degree of freedom for controlling the trade-off between memory efficiency and performance.
  arXiv  Detail & Related papers  (2025-05-03T08:50:14Z)
• Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining [43.39070237124361]
  Low-rank optimization is a promising approach to enabling memory-efficient training of large language models.<n>A key challenge in these methods is selecting suitable subspaces to ensure an effective trajectory.<n>We propose importance sampling for low-rank optimization in LLM pretraining with a provable convergence guarantee.
  arXiv  Detail & Related papers  (2025-02-09T06:30:19Z)
• Zeroth-Order Fine-Tuning of LLMs in Random Subspaces [63.10833446782114]
  As language models grow in size, memory demands for backpropagation increase.<n>Zeroth-order (ZO) optimization methods offer a memory-efficient alternative.<n>In this paper, we propose Subspace Zero-order optimization to address the challenges posed by posed by high dimensionality perturbations.
  arXiv  Detail & Related papers  (2024-10-11T17:01:43Z)
• Differentially Private Optimization with Sparse Gradients [60.853074897282625]
  We study differentially private (DP) optimization problems under sparsity of individual gradients. Building on this, we obtain pure- and approximate-DP algorithms with almost optimal rates for convex optimization with sparse gradients.
  arXiv  Detail & Related papers  (2024-04-16T20:01:10Z)
• Scaling Forward Gradient With Local Losses [117.22685584919756]
  Forward learning is a biologically plausible alternative to backprop for learning deep neural networks. We show that it is possible to substantially reduce the variance of the forward gradient by applying perturbations to activations rather than weights. Our approach matches backprop on MNIST and CIFAR-10 and significantly outperforms previously proposed backprop-free algorithms on ImageNet.
  arXiv  Detail & Related papers  (2022-10-07T03:52:27Z)
• Zeroth-Order Hybrid Gradient Descent: Towards A Principled Black-Box Optimization Framework [100.36569795440889]
  This work is on the iteration of zero-th-order (ZO) optimization which does not require first-order information. We show that with a graceful design in coordinate importance sampling, the proposed ZO optimization method is efficient both in terms of complexity as well as as function query cost.
  arXiv  Detail & Related papers  (2020-12-21T17:29:58Z)
• Extrapolation for Large-batch Training in Deep Learning [72.61259487233214]
  We show that a host of variations can be covered in a unified framework that we propose. We prove the convergence of this novel scheme and rigorously evaluate its empirical performance on ResNet, LSTM, and Transformer.
  arXiv  Detail & Related papers  (2020-06-10T08:22:41Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.