Low-rank gradient sketches nearly halve RL post-training memory at 7B scale

The memory wall for RL post-training is real and specific: a 7B model's BF16 weights cost roughly 14 GB, but Adam's two FP32 moment buffers add another 56 GB on top, before a single activation is counted. LoGRA attacks this by never materializing full gradients for the large attention and MLP projection matrices in the first place. Instead, it accumulates a low-rank sketch directly during backpropagation by projecting the input activations before multiplying by the backward derivatives. The sketch stores a compressed representation rather than the full gradient matrix, and the resulting weight update is expressed as a product of two thin matrices that gets merged directly into the weights. No persistent adapter branch, no frozen backbone — the full weight matrices remain trainable.

The compression ratio is controlled by rank. Ablations show that rank has by far the strongest effect on gradient reconstruction quality: cosine similarity between the dense gradient and its reconstruction rises from roughly 0.1 at rank 4 to roughly 0.6 at rank 256. Whether the projection basis is fixed or refreshed each step matters far less — the largest difference between fixed and refreshed Rademacher projections across all layer-rank cells is under 0.1. This is a practically useful finding: the simpler fixed-basis variant loses little.

The second contribution is predicted-KL step control. Compressing gradients doesn't automatically prevent destabilizing updates; a small change in parameter space can still shift next-token probabilities substantially. LoGRA estimates the KL divergence a proposed update would induce — using a quadratic approximation derived from how token scores change along the update direction — and rescales the update to fit a prescribed budget before applying it. The math is clean: predicted KL scales with the square of the step multiplier, so the required rescaling is a square root. Annealing the budget from a looser to a tighter value over training achieves the best tail accuracy in the ablation, though the authors are careful to note that two-seed results don't firmly establish that annealing alone explains the gain.

The headline numbers are credible and specific. Average per-update memory drops 21.8% at 1.5B parameters and 45.7% at 7B, falling from 31.82 GiB to 17.29 GiB in the latter case. Throughput is essentially unchanged: 25.73 versus 25.28 updates per hour at 7B. The 27B result is the most striking — dense Adam runs out of memory before completing a single update on eight H100s, while LoGRA trains that model stably for over 1,100 steps, with macro-averaged verifier score rising from its initial value to a peak at step 1,060 and no sustained collapse afterward.

Compared with LoRA at rank 256 on the 1.5B model, LoGRA uses roughly half the average memory (7.18 GiB versus 13.21 GiB) but is slower per update and trails on mean Pass@1. The authors frame this honestly: gradient compression is an alternative when memory is the binding constraint, not a universal improvement. The code ships in the Molt library.