# Optimal low-rank stochastic gradient estimation for LLM training

arXiv: https://arxiv.org/abs/2603.20632 · HTML: https://mlsanity.com/paper/2603.20632/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 76 references that are in the map.
- First posted on arXiv 21 Mar 2026.

## Builds on (most cited first, 30 of 35)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming](https://mlsanity.com/paper/1309.5549/) (2013)
- [AdaLoRA: Adaptive Budget Allocation for Parameter-Efficient Fine-Tuning](https://mlsanity.com/paper/2303.10512/) (2023)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [Optimal rates for zero-order convex optimization: the power of two function evaluations](https://mlsanity.com/paper/1312.2139/) (2013)
- [ORLM: A Customizable Framework in Training Large Models for Automated Optimization Modeling](https://mlsanity.com/paper/2405.17743/) (2024)
- [Uncertainty Estimation and Quantification for LLMs: A Simple Supervised Approach](https://mlsanity.com/paper/2404.15993/) (2024)
- [Enhancing Zeroth-order Fine-tuning for Language Models with Low-rank Structures](https://mlsanity.com/paper/2410.07698/) (2024)
- [Improved Zeroth-Order Variance Reduced Algorithms and Analysis for Nonconvex Optimization](https://mlsanity.com/paper/1910.12166/) (2019)
- [Zeroth-Order Regularized Optimization (ZORO): Approximately Sparse Gradients and Adaptive Sampling](https://mlsanity.com/paper/2003.13001/) (2020)
- [Non-asymptotic Analysis of Biased Stochastic Approximation Scheme](https://mlsanity.com/paper/1902.00629/) (2019)
- [Subspace Optimization for Large Language Models with Convergence Guarantees](https://mlsanity.com/paper/2410.11289/) (2024)
- [Zeroth order optimization with orthogonal random directions](https://mlsanity.com/paper/2107.03941/) (2021)
- [LOLA: LLM-Assisted Online Learning Algorithm for Content Experiments](https://mlsanity.com/paper/2406.02611/) (2024)
- [A stochastic subspace approach to gradient-free optimization in high dimensions](https://mlsanity.com/paper/2003.02684/) (2020)
- [AgentGit: A Version Control Framework for Reliable and Scalable LLM-Powered Multi-Agent Systems](https://mlsanity.com/paper/2511.00628/) (2025)
- [Coordinate Descent Without Coordinates: Tangent Subspace Descent on Riemannian Manifolds](https://mlsanity.com/paper/1912.10627/) (2019)
- [Random projections for linear programming](https://mlsanity.com/paper/1706.02768/) (2017)
- [Gradient-based Algorithms for Convex Discrete Optimization via Simulation](https://mlsanity.com/paper/2010.16250/) (2020)
- [Reusing Historical Trajectories in Natural Policy Gradient via Importance Sampling: Convergence and Convergence Rate](https://mlsanity.com/paper/2403.00675/) (2024)
- [Understanding the Training and Generalization of Pretrained Transformer for Sequential Decision Making](https://mlsanity.com/paper/2405.14219/) (2024)
- [Assured autonomy: How operations research powers and orchestrates generative AI systems](https://mlsanity.com/paper/2512.23978/) (2025)
- [Accelerating Greedy Coordinate Descent Methods](https://mlsanity.com/paper/1806.02476/) (2018)
- [Adaptive Importance Sampling for Efficient Stochastic Root Finding and Quantile Estimation](https://mlsanity.com/paper/2102.10631/) (2021)
- [Proximal Gradient methods with Adaptive Subspace Sampling](https://mlsanity.com/paper/2004.13356/) (2020)
- [Distributionally Constrained Black-Box Stochastic Gradient Estimation and Optimization](https://mlsanity.com/paper/2105.09177/) (2021)
- [Large-Scale Inventory Optimization: A Recurrent-Neural-Networks-Inspired Simulation Approach](https://mlsanity.com/paper/2201.05868/) (2022)
- [Stochastic Gradient Descent with Adaptive Data](https://mlsanity.com/paper/2410.01195/) (2024)
- [Stochastic Optimization with Optimal Importance Sampling](https://mlsanity.com/paper/2504.03560/) (2025)

## Built on by (most cited first, 1 of 1)

- [Scaling Self-Evolving Agents via Parametric Memory](https://mlsanity.com/paper/2606.04536/) (2026)

## Abstract (arXiv)

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.