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

AI summary

Overview

Research area: Memory-efficient reinforcement learning (RL) post-training for large language models, combining gradient compression (low-rank sketches) with step-size control.

Technical level: Advanced. The paper assumes familiarity with RL fine-tuning loops (PPO, policy/rollout separation), optimizer states (Adam/SGD), and KL-divergence trust-region arguments. Appendix A derives the KL estimator analytically.

Scope: One sentence: the paper proposes and evaluates LoGRA, which replaces full gradient buffers in RL post-training with low-rank gradient sketches reused for both weight updates and rollout-policy synchronization, combined with a predicted-KL controller that scales each update to a KL budget.

What This Paper Is About

RL post-training is memory-hungry because gradients and optimizer states can consume several times more memory than the model weights. The paper gives the example of a 7B-parameter Qwen model whose BF16 weights require approximately 14 GB while Adam's two FP32 moment buffers add another 56 GB, before gradients and activations. The goal is to reduce that memory footprint without giving up reasoning accuracy, so that RL training becomes feasible on hardware that can already serve the model.

Key Contributions

1. Low-rank gradient sketches accumulated during backpropagation. LoGRA picks a rank r ≪ min(d, k) and a fixed random projection A ∈ ℝ^(r×k), representing the gradient of a weight matrix W ∈ ℝ^(d×k) as a sketch S = GA^T ∈ ℝ^(d×r). The sketch is accumulated directly from each microbatch in fp32 — the full gradient matrix is never formed or stored. The update is applied as W ← W − αηUA, where U is the (optionally optimizer-adjusted) sketch. The accumulator shrinks from d·k to d·r values; for d = k = 4096 and r = 64, fp32 storage falls from 64 MiB to 1 MiB.

2. Reuse of the compressed representation for policy synchronization. The rollout engine regenerates A from its seed and applies the received factor αηU, so the synchronization payload is O(dr) values plus a seed and metadata instead of a full d×k update. Compression therefore extends from training-state storage to policy communication, with no persistent adapter branch — updates are merged directly into the weights.

3. Predicted-KL step control. Because an update small in parameter space can still shift output probabilities, the method estimates the average KL between the current and proposed policies over a subset of contexts, giving q(D) = (1 / (2|P|T²)) Σ_s Var_{v∼p(s)}[ż_v(s)], then sets α = min{α_max, √(δ/q(D))} for KL budget δ. Steps are rescaled, not redirected; predicted KL scales with α², so halving α cuts predicted KL by a factor of four.

4. An empirical demonstration that the combination reduces memory and enables otherwise infeasible training. The abstract reports reductions of up to 45.7% in average training memory across reasoning tasks "without compromising performance," and 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. Code is released in the Molt library. (The paper reports that a measured memory-reduction ratio for 27B is precluded because Dense fails before completing an update.)

Main Findings

• Memory reduction at 1.5B and 7B. Over complete runs, average per-GPU memory fell from 9.18 to 7.18 GiB at 1.5B (21.8%) and from 31.82 to 17.29 GiB at 7B (45.7%). Under the seed-averaged peak definition, peak memory fell from 9.19 to 8.66 GiB at 1.5B and from 31.85 to 22.26 GiB at 7B.
• Accuracy at 1.5B improved. Peak Pass@1 rose from 63.77% (Dense) to 67.87% (LoGRA) and peak Pass@4 from 80.33% to 82.47%. At 7B the methods were near-matched: 72.33% versus 72.48% Pass@1 and 85.13% versus 85.33% Pass@4 for LoGRA and Dense respectively.
• 27B becomes trainable where Dense does not. Dense ran out of memory at its first Adam-state allocation (n = 0 runs). LoGRA trained Qwen3.8-27B at 51.54 GiB average memory, 53.88 GiB peak, 6.99 updates/h, reaching 71.52% peak Pass@1 and 81.87% peak Pass@4.
• Throughput is roughly comparable, not the main selling point. End-to-end update rates were 55.44 versus 54.28 updates/h at 1.5B and 25.73 versus 25.28 updates/h at 7B for LoGRA and Dense.
• Sustained 27B training over 1,100 steps. On the Reasoning-Gym Hard mixture (38,948 training prompts, 193 hard held-out problems across eight tasks), held-out macro score rose from 39.69% to 62.94%, peaking at 65.52% at step 1,060; the final five observed evaluations ranged from 59.11% to 64.16%, with no sustained collapse. Mean response length decreased from 8,312 to 4,705 tokens and truncation at the 16,384-token limit fell from 34.72% to 12.95%.
• Rank matters more than projection design. Increasing rank from 4 to 256 improved pass@1 by 3.02–3.95 percentage points across configurations. At rank 256, fixed Rademacher reached 68.50%, refreshed Rademacher 68.31%, and refreshed Gaussian 67.42%, with ordering varying across ranks and no consistent benefit from refreshing. Gradient-reconstruction cosine similarity rose from approximately 0.046 at rank 4 to 0.348 at rank 256, while the largest difference between corresponding layer–rank cells was only 0.0014 (fixed vs. refreshed Rademacher) and 0.0008 (refreshed Rademacher vs. Gaussian).
• Against LoRA, LoGRA wins on memory only. Over the first 300 training steps on Qwen2.5-Math-1.5B, LoGRA used 7.18 GiB average and 8.66 GiB peak versus 13.21 and 13.38 GiB for LoRA. LoRA attained higher mean Pass@1 (70.45% vs. 68.47%) and a shorter policy-training call (10.63 vs. 11.31 seconds/update), while LoGRA attained a 0.40-percentage-point higher mean Pass@4 (81.80% vs. 81.40%). The paper states these results motivate gradient compression when memory is the limiting resource but do not establish an accuracy or speed advantage.
• Conservative KL budgets help long-horizon accuracy. With pass@1 averaged over steps 700–800: 70.30% without KL control, 63.43% with the large constant budget (1.5×10⁻³), 71.03% with the small constant budget (1.5×10⁻⁴), and 71.47% with the cosine-annealed budget (2×10⁻⁴ to 2×10⁻⁵). Corresponding response lengths were 774, 976, 858, and 832 tokens. The paper notes these are two-seed results and that they suggest a benefit from well-calibrated KL control without establishing that it is necessary for stability or that annealing alone explains the gain.

Methodology in Plain English

The authors start from the observation that RL with binary outcome rewards gives each generated response only a single correctness signal, so the learning signal should be far lower-dimensional than a full gradient.

Instead of accumulating a full d×k gradient across microbatches, they multiply each microbatch's gradient contribution by a fixed random projection matrix A^T, producing a much narrower sketch S with only r columns. Because the projection is applied to the inputs of the weight-gradient calculation before the backward derivatives are multiplied in, the full gradient is never materialized. The same A is reused across microbatches during accumulation, then redrawn each step so that successive updates can occupy different subspaces.

To apply an update, the sketch is multiplied back by A to give a rank-r approximation of the full gradient, which is merged directly into the weight matrix — no adapter stays attached to the model. An optimizer can modify the sketch (the paper's RowAdam, used in the main experiments) before that multiplication.

Because a compressed update can still move the policy too far, they add a controller. Before committing an update, they take a sample of contexts from the current batch and estimate how much the proposed step would change next-token probabilities, using the variance of logit-change rates weighted by current token probabilities. That estimate predicts a quadratic dependence on step size, so they rescale the update by the square root of the ratio of the KL budget to the estimate, capped by a maximum multiplier. For the 1,000-step 27B run they replace the analytic estimator with a finite-difference probe based on direct KL evaluations.

The experimental setup is deliberately simple to isolate the method: PPO with verifiable rewards, one response per prompt, a global running reward baseline, batches of 128 prompts, a 4,096-token context limit, and a single node with eight H100 80GB GPUs and 128 CPU cores (four GPUs for FSDP training, four for generation). Main comparisons run on DAPO-Math-7.5K with three seeds, evaluated on MATH-500.

Why This Matters

Impact on research. The paper argues that gradient representation, not just parameter representation, is an underused lever for lowering the hardware barrier to RL post-training. It extends prior compact-gradient work such as GaLore and Flora from optimizer-state storage into gradient computation and policy communication, and it connects low-rank compression to the trust-region logic of TRPO in a way that is specific to RL's training–generation loop.

Real-world applications:

• Academic and small-lab RL fine-tuning, where a 27B model is otherwise out of reach on a single eight-GPU node.
• Reinforcement learning from verifiable rewards for mathematical and multi-step reasoning, where accuracy must be preserved while memory is cut.
• Long-context reasoning training, where the reported drop in truncation rate (34.72% to 12.95%) and shorter responses (8,312 to 4,705 tokens) matter for cost.
• Rollout-heavy distributed serving setups, where synchronizing compressed O(dr) updates instead of full d×k updates reduces communication payload.

Industry relevance. Memory is often the binding constraint on whether a model can be improved after deployment, not whether it can be served. A training method that fits a 27B model on the same node type — with average memory of 51.54 GiB against 80GB GPUs — changes infrastructure planning. The paper also positions the work as complementary to existing system-level savings (FSDP, gradient checkpointing, CPU offloading) and to asynchronous RL overlap.

Future Directions

• How far does rank scaling go? Rank had the strongest single effect on both accuracy and gradient reconstruction, but the paper only tests r ∈ {4, 16, 64, 256} on Qwen2.5-Math-1.5B. Whether the trend continues, or where it saturates, is not established.
• Is predicted-KL control necessary, or only helpful? The no-control baseline in the KL ablation remained stable, the annealed runs had a different initial budget, and the paper reports only two seeds. A cleaner separation of annealing from budget calibration is left open.
• Better KL estimators and their cost. The 27B long-horizon run used a finite-difference probe rather than the JVP-based estimator of Equation 6 for efficiency. Comparing estimator variants at scale is an open question.
• Scaling beyond 27B and beyond a single node. The largest demonstration is one 27B model on one eight-GPU node; multi-node behavior, larger models, and interaction with asynchronous RL overlap and other memory-saving techniques are untested here.

Target Audience

This paper is most useful to RL and post-training engineers and systems researchers who train or fine-tune large language models under fixed GPU budgets, and who already understand PPO-style pipelines and optimizer-state accounting. It also suits researchers working on gradient compression and memory-efficient optimization who want to see low-rank methods moved into the RL training–generation loop. Readers looking for a drop-in accuracy improvement over dense Adam or LoRA should note the paper's own framing: the demonstrated gains are in memory and feasibility, with throughput roughly comparable and no established accuracy or speed advantage.

Authors’ 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 \href{https://github.com/skzhang1/labs-molt/tree/logra/examples/scripts/logra}{Molt library}.