Activation Compression
Recomputation discards activations and regenerates them. The third option is to keep them in two or four bits — the backward pass turns out to tolerate a surprising amount of noise in what it reads.
Standing
PromisingPromising and actively moving. The results are real but narrow — one lab, one model family, or one benchmark suite — and the picture may look different in a year.
Sound in principle, with convergence guarantees, and rarely used — recomputation is simpler and FlashAttention removed the largest tensor it targeted.
judged as of 2026-09 · what the labels mean
Theory
The saved activations are read exactly once, by the backward pass, to compute a gradient that is already a noisy estimate. That framing suggests they do not need to be stored exactly.
At two bits that is a sixteen-fold reduction against bf16, for the cost of a quantise on the forward and a dequantise on the backward — both elementwise, both cheap.
Why the gradient tolerates it
Because it is already stochastic. The minibatch gradient has variance from sampling; quantisation adds a second, independent source. If the second is small relative to the first, convergence is essentially unaffected.
Chen et al. make this precise with stochastic rounding rather than round-to-nearest. Rounding up or down with probability proportional to the remainder makes the quantised value unbiased, so and SGD’s convergence guarantees carry over with a modified constant. Deterministic rounding introduces a bias that accumulates over layers and over steps, and it does not.
Against recomputation
They target the same memory and the currency differs.
| Compression | Recomputation | |
|---|---|---|
| costs | ~20% throughput | 4–36% FLOPs |
| memory | ~12× less | ~10× less |
| exact | no | yes |
| implementation | custom kernels | one wrapper |
Recomputation is exact, is one function call, and is in every framework. Compression needs quantise/dequantise kernels fused into the autograd graph, and gives up exactness for a guarantee about the expectation.
The comparison used to be closer. What changed is FlashAttention: the quadratic score matrix was the biggest single target for compression, and it is no longer stored by anyone. What remains is the linear term, where selective recomputation already gets most of the saving for about 4% compute — a bar that a 20% throughput cost does not clear.
Implementation
import torch
from torch import Tensor
class Compressed(torch.autograd.Function):
"""Store the activation in 4 bits; dequantise when the backward reads it."""
@staticmethod
def forward(ctx, x: Tensor, group: int = 128):
flat = x.reshape(-1, group)
lo = flat.amin(1, keepdim=True)
scale = (flat.amax(1, keepdim=True) - lo).clamp(min=1e-8) / 15
# Stochastic rounding: unbiased, which is what preserves convergence.
q = (flat - lo) / scale
q = (q.floor() + (torch.rand_like(q) < q.frac())).to(torch.uint8)
ctx.save_for_backward(q, scale, lo) # 4 bits, not 16
ctx.shape, ctx.group = x.shape, group
return x
@staticmethod
def backward(ctx, grad: Tensor):
q, scale, lo = ctx.saved_tensors
_ = (q.float() * scale + lo).reshape(ctx.shape) # the recovered input
return grad, NoneTwo things that implementation glosses over. uint8 stores four bits in eight —
a real version packs two values per byte, or the saving is halved. And
save_for_backward must receive the quantised tensor only; keeping a reference
to x anywhere in the closure keeps the full-precision tensor alive and the
whole exercise achieves nothing, which is a mistake that produces correct results
and no memory saving.