# PLUMAGE: Probabilistic Low rank Unbiased Min Variance Gradient Estimator for Efficient Large Model Training

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 28 of its 38 references that are in the map.
- First posted on arXiv 23 May 2025.

## Builds on (most cited first, 28 of 28)

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [QLoRA: Efficient Finetuning of Quantized LLMs](https://mlsanity.com/paper/2305.14314/) (2023)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [8-bit Optimizers via Block-wise Quantization](https://mlsanity.com/paper/2110.02861/) (2021)
- [Sparse Communication for Distributed Gradient Descent](https://mlsanity.com/paper/1704.05021/) (2017)
- [PowerSGD: Practical Low-Rank Gradient Compression for Distributed Optimization](https://mlsanity.com/paper/1905.13727/) (2019)
- [Scalable Methods for 8-bit Training of Neural Networks](https://mlsanity.com/paper/1805.11046/) (2018)
- [ATOMO: Communication-efficient Learning via Atomic Sparsification](https://mlsanity.com/paper/1806.04090/) (2018)
- [Flora: Low-Rank Adapters Are Secretly Gradient Compressors](https://mlsanity.com/paper/2402.03293/) (2024)
- [Accelerated Sparse Neural Training: A Provable and Efficient Method to Find N:M Transposable Masks](https://mlsanity.com/paper/2102.08124/) (2021)
- [AI and Memory Wall](https://mlsanity.com/paper/2403.14123/) (2024)
- [Variance Reduction in SGD by Distributed Importance Sampling](https://mlsanity.com/paper/1511.06481/) (2015)
- [Memory-Efficient LLM Training with Online Subspace Descent](https://mlsanity.com/paper/2408.12857/) (2024)
- [Fisher Information and Natural Gradient Learning of Random Deep Networks](https://mlsanity.com/paper/1808.07172/) (2018)
- [Grass: Compute Efficient Low-Memory LLM Training with Structured Sparse Gradients](https://mlsanity.com/paper/2406.17660/) (2024)
- [CompAct: Compressed Activations for Memory-Efficient LLM Training](https://mlsanity.com/paper/2410.15352/) (2024)
- [Towards Cheaper Inference in Deep Networks with Lower Bit-Width Accumulators](https://mlsanity.com/paper/2401.14110/) (2024)

## Abstract (arXiv)

Accelerator memory and networking constraints have emerged as dominant bottlenecks when training large language models LLMs with billions of parameters. Existing low rank gradient estimators such as GaLoRE and FLORA compress gradients and optimizer tensors by projecting weight gradients onto a rank r subspace, enabling LLM training on consumer hardware. Yet, these methods are either biased or subject to high estimator variance. Moreover, the optimizer state based on the first and second moments estimates expressed in the previous subspace becomes misaligned whenever the projection is updated, leading to instabilities during training. We propose PLUMAGE: Probabilistic Low rank Unbiased Minimum vAriance Gradient Estimator. PLUMAGE is a drop in replacement for existing low rank gradient estimators. It does not introduce new hyperparameters beyond the chosen rank r and the update interval. In addition, we resolve optimizer state misalignment issues to prevent spurious weight updates and enhance training stability. We empirically demonstrate that PLUMAGE shrinks the full rank optimization's gap over the pre training evaluation loss by 33% on average across models and the average training loss across the GLUE benchmark by 28% within a similar computational and memory footprint as GaloRE.