LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches

Shaokun Zhang; Yifan Zhang; Jian Hu; Yueying Li; Hao Zhang; Binfeng Xu; Jan Kautz; Yi Dong

cs.CLOct 5, 2026

LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches

Authors: Shaokun Zhang, Yifan Zhang, Jian Hu, Yueying Li, Hao Zhang, Binfeng Xu, Jan Kautz, Yi Dong

Abstract

Reinforcement learning (RL) has greatly advanced the capabilities of large language models (LLMs), but its memory demands remain a barrier to broader adoption. We introduce LoGRA, an approach to RL post-training that reduces memory by retaining useful learning signals in low-rank gradient sketches. These compact representations support both model updates and efficient policy synchronization. To prevent overly large updates from disrupting learning, we complement gradient compression with predicted-KL step control, which estimates policy changes before applying each update and adjusts its magnitude accordingly. Across reasoning tasks, LoGRA reduces average training memory by up to 45.7% without sacrificing performance. It also enables stable training of a 27B-parameter model for over 1,100 steps on a single eight-GPU node, where dense Adam runs out of memory, making previously memory-infeasible RL training practical. Code is available in the Molt library.

Figures & tables

Figure 1: Lower training memory across model scales. Colors denote model sizes; hatching distinguishes Dense Adam from LoGRA. Dark bars show average memory and pale extensions reach peak memory (Table 1 ); labels report average / peak in GiB per training GPU. Arrows show average-memory savings over Dense Adam. The broken axis omits 34–48 GiB. The 27B Dense Adam bar marks OOM at GPU capacity, not a measured value.

Figure 2: Learning performance across model scales. MATH-500 accuracy over the first 200 updates. Curves and bands show means and

\pm

one sample standard deviation over three seeds.

Figure 3: Training memory across model scales. Per-update peak allocated memory, averaged over three seeds, with

\pm

one sample standard deviation. Arrows compare temporal means over updates 1–200 at 1.5B/7B; at 27B, the arrow shows headroom to device capacity.

Model	Method	$n$	Average memory $\downarrow$ (GiB/GPU)	Peak memory $\downarrow$ (GiB/GPU)	Updates/h $\uparrow$	Peak Pass@1 $\uparrow$ (%)	Peak Pass@4 $\uparrow$ (%)
1.5B	Dense	3	$9.18\pm 0.00$	$9.19\pm 0.00$	$54.28\pm 0.40$	$63.77\pm 1.35$	$80.33\pm 1.47$
1.5B	LoGRA	3	$\mathbf{7.18\pm 0.10}$	$\mathbf{8.66\pm 1.62}$	$\mathbf{55.44\pm 0.41}$	$\mathbf{67.87\pm 1.23}$	$\mathbf{82.47\pm 1.10}$
7B	Dense	3	$31.82\pm 0.00$	$31.85\pm 0.00$	$25.28\pm 0.68$	$72.48\pm 0.35$	$85.33\pm 0.83$
7B	LoGRA	3	$\mathbf{17.29\pm 0.34}$	$\mathbf{22.26\pm 4.98}$	$\mathbf{25.73\pm 0.28}$	$72.33\pm 1.33$	$85.13\pm 1.14$
27B	Dense	0	OOM
27B	LoGRA	3	$\mathbf{51.54\pm 0.17}$	$\mathbf{53.88\pm 0.00}$	$\mathbf{6.99\pm 0.02}$	$\mathbf{71.52\pm 0.94}$	$\mathbf{81.87\pm 0.81}$

Table 1: Accuracy, memory, and throughput of the main experiments. Memory and throughput cover full runs; accuracy peaks use updates 0–200. Peak memory is the maximum of the three-seed mean curve.

Figure 4: Sustained 27B training on a single eight-GPU node. Qwen3.8-27B trained with LoGRA on the Reasoning-Gym Hard mixture. (a) Macro-averaged verifier score across held-out tasks. (b) Mean evaluation response length. (c) Fraction of evaluation responses truncated at the 16,384-token context limit. All panels describe one training trajectory.

Figure 5: Components of gradient compression. (a) Training performance as a function of projection rank; (b) Gradient-reconstruction cosine for fixed and refreshed Rademacher bases. (c) The same diagnostic for refreshed Rademacher and Gaussian bases. Heatmap rows denote Transformer layers and columns denote ranks, with a shared $[0,0.4]$ color scale; higher values indicate closer alignment with the dense gradient.

Method	Average memory $\downarrow$ (GiB)	Peak memory $\downarrow$ (GiB)	Pass@1 $\uparrow$ (%)	Pass@4 $\uparrow$ (%)	Seconds/update $\downarrow$ (s)
LoRA	$13.21\pm 0.07$	$13.38\pm 0.00$	$70.45\pm 3.09$	$81.40\pm 2.09$	$10.63\pm 0.10$
LoGRA	$\mathbf{7.18\pm 0.10}$	$\mathbf{8.66\pm 1.62}$	$68.47\pm 0.95$	$\mathbf{81.80\pm 1.06}$	$11.31\pm 0.09$

Table 2: LoRA and LoGRA on Qwen2.5-Math-1.5B. Values are mean

\pm

sample standard deviation over three seeds.

Figure 6: Long-horizon accuracy and response length under different KL controls. Thick solid lines show two-seed means, thin lines individual runs, and shading their range. Curves use five-evaluation moving means.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Training problems	Evaluation problems
codeio	4,868	25
emoji_mystery	4,752	22
rectangle_count	5,000	25
word_ladder	4,994	25
pool_matrix	5,000	25
tsumego	5,000	25

Appendix

Table 3: Reasoning-Gym Hard mixture.

Setting	Dense	LoGRA
Training seeds	$1,2,42$	$1,2,42$
Prompts / responses per update	$128$ / $128$	$128$ / $128$
Microbatch per training GPU	$1$	$1$
Optimizer	Adam (zero weight decay)	RowAdam
Moment decay	$(\beta_{1},\beta_{2})=(0.9,0.95)$	$\beta_{2}=0.95$ ; no first moment
Epsilon	$10^{-8}$ , absolute	$10^{-8}$ , relative

Appendix

Table 4: Main-experiment settings. Schedule horizons count actual optimizer updates.

Explore similar work

Memory-Efficient Policy Libraries with Low-Rank Adaptation in Reinforcement Learning
Jun 24, 2026Samuel Valland Lyngset, Tor Viljen Raanaas, Gard Sveipe +4Reinforcement Fine-Tuning Model Fine-Tuning
Department of Informatics, University of Oslo
Just-In-Time Reinforcement Learning: Continual Learning in LLM Agents Without Gradient Updates
Jan 26, 2026Yibo Li, Zijie Lin, Ailin Deng +5Large Language Model Reinforcement Learning Offline Reinforcement Learning
National University of Singapore, Singapore.
How LoRA Remembers? A Parametric Memory Law for LLM Finetuning
May 28, 2026Ziwen Xu, Haiwen Hong, Linsong Yu +4Parametric Memory Large Language Model Fine-Tuning
Zhejiang University · Alibaba Group

Jun 24, 2026, cs.LGJ/K move · Enter open · S save

Memory-Efficient Policy Libraries with Low-Rank Adaptation in Reinforcement Learning

Samuel Valland Lyngset, Tor Viljen Raanaas, Gard Sveipe, Eirik Møller Nilsen+3

Department of Informatics, University of Oslo

Reinforcement Fine-Tuning Model Fine-Tuning Peak Memory

Open paper AskPDF arXiv

When fine-tuning Large Language Models (LLMs), there has been success in minimizing both memory usage and computation with Parameter-Efficient Fine-Tuning (PEFT), like Low Rank Adaptation (LoRA). In this article, we have explored whether this approach is transferable to the world of robotics and Reinforcement Learning (RL), allowing learning with reduced memory usage and improved computational performance. Specifically, we focused on a version of multi-task robotics, where a library of specialist policies are created. In such a library memory efficiency is especially important. We used a Proximal Policy Optimization (PPO) algorithm and fine-tuned a baseline model to different tasks using LoRA. Our results demonstrate that, depending on the hyperparameters, LoRA can minimize memory usage by a factor of 20-160 compared to full fine-tuning of all layers. This implies a 90-95% storage saving when deploying a library of many (10-50) specialized policies, which can be the differentiating factor between being able to store the entire library in memory or having to use swap-memory in an applied robotics setting. At the same time, our results indicate that there is no significant difference in the success-rate between full fine-tuning and LoRA fine-tuning for the selected tasks.

Abstract

Explore similar work

Memory-Efficient Policy Libraries with Low-Rank Adaptation in Reinforcement Learning

Just-In-Time Reinforcement Learning: Continual Learning in LLM Agents Without Gradient Updates

How LoRA Remembers? A Parametric Memory Law for LLM Finetuning