Where RL training memory actually goes
The paper frames the memory problem as bigger than storing weights. For a 7B-parameter Qwen model, BF16 weights take roughly 14 GB, while Adam's two FP32 moment buffers add another 56 GB, before gradients and activations are counted. That gap is why hardware able to run a model may still be unable to improve it through reinforcement learning.
LoGRA targets the gradient and optimizer side rather than the weights. It keeps the original weight matrices trainable, unlike LoRA-style adapters that freeze pretrained weights and learn small low-rank branches. The paper's stated intuition is that binary outcome rewards give each generated response only a single correctness signal, so learning from rewards may need far less information than a full gradient represents.
How a sketch replaces the full gradient buffer
For a weight matrix W of size d×k, LoGRA picks a rank r much smaller than min(d,k) and a random projection matrix A of size r×k. Instead of accumulating a full d×k gradient, it accumulates a sketch S = G Aᵀ of size d×r. The projection A is sampled independently of the current batch, held fixed during accumulation, and not learned; by default its entries are independent random signs scaled by 1/√r.
The implementation accumulates projected contributions from each microbatch directly into S in fp32, projecting the inputs used in the weight-gradient calculation before multiplying by backward derivatives. No full gradient buffer is formed for the compressed matrices. The update is then W ← W − αηUA, where U equals S for the basic SGD variant and α is a multiplier chosen by the step controller. The update is merged directly into the weights, with no persistent adapter branch.
The memory arithmetic is concrete: for d = k = 4096 and r = 64, fp32 storage for the gradient accumulator falls from 64 MiB to 1 MiB. The paper is explicit that this is not the whole training footprint — model weights, cached projections, activations, and other buffers still consume memory, and stateful sketch optimizers add their own storage.
The same sketch also shrinks policy synchronization
Compression extends beyond training memory to communication. For synchronization, the rollout engine regenerates A from its seed and applies the received factor αηU. The payload is O(dr) values plus a seed and metadata, instead of a full d×k update. The paper presents this as reusing one compact representation across gradient accumulation, weight updates, and policy synchronization.
Predicted-KL control keeps large updates from breaking learning
Compression alone is not enough. The paper notes that an update small in parameter space can still substantially change the policy's output probabilities, and an overly large update can abruptly shift the policy and disrupt previously learned behavior. LoGRA therefore adds predicted-KL step control, which estimates how much a proposed update would change next-token probabilities before applying it.
The estimate samples a small subset of responses from the current batch and compares the current model's vocabulary distribution with that of the proposed updated model at the same temperature. The paper's approximation predicts a quadratic dependence on step size: the average predicted KL scales as α²q(D), so halving α reduces predicted KL by a factor of four. Given a KL budget δ, the controller sets α = min{α_max, √(δ/q(D))}. A proposal with q(D) = 4δ is scaled to at most half its original size, and the α_max cap prevents excessive enlargement when predicted KL is small.
The paper is careful about what this controls: it regulates the predicted change, while the realized KL can differ because the estimate is a local approximation. Budgets can be constant or annealed, with an annealed budget gradually lowering δ to permit less policy change later in training.
Memory savings grow with model scale
Across reasoning tasks, LoGRA reduces average training memory by 21.8% at 1.5B and 45.7% at 7B without sacrificing task performance. The reported 7B figures are 17.29 GiB average and 22.26 GiB peak per training GPU for LoGRA, against 31.82 GiB average and 31.85 GiB peak for dense Adam. At 1.5B the comparison is 7.18 GiB average and 8.66 GiB peak for LoGRA versus 9.18 GiB average and 9.19 GiB peak for dense Adam.
The 27B result is a feasibility claim rather than a measured memory number. LoGRA sustains stable training of a 27B model for over 1,100 steps on a single eight-GPU node, where the dense Adam baseline runs out of memory under the same hardware allocation. The paper marks the 27B dense Adam bar as OOM at GPU capacity, not a measured value, so the comparison there is feasibility, not a memory ratio.
Projection rank matters more than basis refreshing
The design-choice experiments use Qwen2.5-Math-1.5B on DAPO-Math-7.5k with ranks r ∈ {4, 16, 64, 256}. Increasing rank from 4 to 256 improves MATH-500 pass@1 by 3.02 to 3.95 percentage points across configurations. At rank 256, fixed Rademacher, refreshed Rademacher, and refreshed Gaussian projections reach 68.50%, 68.31%, and 67.42% respectively, and their ordering varies across ranks.
The paper reports no consistent benefit from refreshing the projection basis, possibly because a fixed low-rank subspace captures sufficient task-relevant signal over this short training horizon. Gradient-reconstruction diagnostics point the same way: mean cosine similarity between the dense gradient and its reconstruction rises from about 0.046 at rank 4 to 0.348 at rank 256, while the largest difference between corresponding layer–rank cells is only 0.0014 for fixed versus refreshed Rademacher and 0.0008 for refreshed Rademacher versus Gaussian. The Gaussian decrease at ranks 64 and 256 stays below 2×10⁻⁴. Rank, not refreshing or entry distribution, dominates single-minibatch reconstruction quality.
What these savings do and do not settle
The comparison with LoRA is bounded in the paper's own terms: LoRA learns low-rank weight adapters while LoGRA compresses gradients and updates to the weights, and the reported comparison uses rank-256 Qwen2.5-Math-1.5B on DAPO-Math-7.5k over the first 300 training steps, with LoRA on AdamW at a constant 2×10⁻⁴ learning rate and LoGRA on a predicted-KL budget annealed from 2×10⁻⁴ to 2×10⁻⁵.
The design-choice evidence is also scoped: the rank and refreshing experiments run on a 1.5B model over a short horizon, and the paper itself flags that the fixed-subspace result may reflect that short horizon rather than a general property. The 45.7% figure is an average-memory saving at 7B, not a peak-memory guarantee, and the 27B claim is about completing 1,100-plus steps where dense Adam cannot start, not about a measured memory reduction at that scale.
The practical priority this paper establishes is a pairing, not a single trick: the sketch is what removes the full gradient buffer, but predicted-KL step control is what keeps the compressed update from destabilizing the policy. Readers weighing LoGRA against adapter-based or system-level memory savings should compare on the axis the paper actually measures — average training memory per GPU at a given model scale and rank — and treat the 27B result as evidence of feasibility under one eight-GPU allocation rather than as a general scaling law.
Sources used in this answer
LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches
LoGRA replaces full gradients with low-rank sketches for selected weight matrices and adds predicted-KL step control, reducing average training memory by up to 45.7% at 7B without sacrificing reasoning performance and enabling 27B training for over 1,100 steps on a single eight-GPU node where dense Adam runs out of memory.