LoGRA: Low-Rank Gradient Sketches

Why People Wanted It

In RL post-training, a selected weight matrix can have a gradient as large as the matrix itself. Keeping that full gradient while training costs memory, and sending a full update to the rollout engine costs communication. So the real question was: can you accumulate and apply a useful gradient representation without ever building the full gradient buffer?

The Core Move

LoGRA projects each weight-gradient contribution into a small fixed random basis, accumulates that sketch during backpropagation, then uses the sketch factors for a low-rank update and policy synchronization.

For a weight matrix $W\in\mathbb{R}^{d\times k}$, choose a random projection $A\in\mathbb{R}^{r\times k}$, where $r\ll\min(d,k)$. The paper defines the sketch as $S=GA^\top$, with $G=\nabla_W\mathcal{L}$. This notation states the desired sketch; the implementation accumulates projected contributions directly rather than first forming $G$.

Quiz

What does LoGRA accumulate during a microbatch, instead of first materializing the full weight gradient?

A learned projection matrix updated by the batch

The projected gradient sketch

Only the scalar loss, with no gradient-related state

A full-sized copy of the weight matrix

$ check

Why the Sketch Can Be Accumulated Without a Full Gradient

Here is the first deeper layer: the paper’s key implementation detail is to move the projection into the weight-gradient calculation. For each microbatch contribution $\mathcal{L}_m$, the algorithm adds $(\nabla_W\mathcal{L}_m)A^\top$ to $S$. Since $A$ is held fixed during accumulation, this gives the projected sum of microbatch gradients without allocating the full $d\times k$ gradient matrix.

The projection is sampled independently of the current batch and is not learned. Its entries are random signs scaled by $1/\sqrt{r}$ by default. The sketch has $dr$ values instead of $dk$; this saves gradient-accumulator storage, not the memory for model weights, activations, cached projections, or other buffers. The ordinary forward computation is preserved.

Why the Update Is Low-Rank but Not an Exact Recovery

Now the version for readers who want more of the machinery: the factors reconstruct a full-shaped approximation, not the full gradient itself. The paper uses $SA$ in place of $G$; each column of $SA$ is a combination of the $r$ columns of $S$, so the update has rank at most $r$.

An optimizer may adjust the sketch before applying it. Calling the adjusted sketch $U$, the update is $W\leftarrow W-\alpha\eta UA$. In the basic SGD variant, $U=S$; $\eta$ is the nominal step size and $\alpha$ scales it. The equation is the update written as a product of compact factors, while $W$ itself remains a full-sized matrix. It is an approximation to the gradient, not an exact reconstruction.

Why Refreshing the Projection Enables New Update Directions

Here is the next technical layer: the projection is held fixed across one accumulation, then refreshed each training step. The rollout engine can regenerate the same $A$ from its seed and apply the received factor $\alpha\eta U$.

Because successive steps can use different projection subspaces, their summed updates need not have rank at most $r$. Synchronization sends the compact factor plus a seed and metadata instead of a full $d\times k$ update. The mechanism’s stated scope is selected attention and MLP projection matrices, with other parameters held fixed. What the sketch saves is clear; what it leaves to the optimizer is how faithfully the low-rank approximation serves a particular training objective.

Minimal Causal Code

This plain loop shows the update state: each microbatch contributes directly to the sketch, and the full gradient matrix is never created.

# W: [d, k]; A: [r, k]; S: [d, r]
A = sample_projection(seed, r, k)
S = torch.zeros(d, r, dtype=torch.float32, device=W.device)

for microbatch in microbatches:
    # Conceptual projected weight-gradient contribution.
    # A real implementation projects the inputs in the weight-gradient calculation.
    S += projected_weight_gradient(microbatch, W, A)

U = S  # basic SGD variant
W -= alpha * eta * (U @ A)

# Send alpha * eta * U and the seed; rollout regenerates A and applies the update.

What You Gain, and What You Lose

You gain a smaller gradient accumulator and a synchronization payload based on the sketch factors rather than a full matrix update. For the paper’s stated example with $d=k=4096$ and $r=64$, fp32 gradient-accumulator storage falls from 64 MiB to 1 MiB; other model and training buffers still consume memory. Basic SGD needs no optimizer moments, while stateful sketch optimizers add storage of their own.

You give up exact full-gradient updates: $SA$ is a low-rank approximation. The weights remain full-sized, and the savings do not erase activation or other memory costs. Refreshing the projection changes the subspace from step to step, but each individual update still uses a rank-limited sketch. Policy-update magnitude control is a separate concern from this module’s compact representation.

Quiz

Why can successive LoGRA updates have a sum whose rank exceeds the per-step sketch rank?

The weight matrix becomes smaller after every update

The projection is learned until its rank grows

Each step may use a freshly sampled projection and therefore a different subspace

The rollout engine sends back extra gradient columns

$ check