Backprop Without the Memory

Ask what limits the depth of a network and the natural guess is compute: more layers, more arithmetic. The actual binding constraint, on real hardware, is usually something less glamorous. To compute a gradient, backpropagation must remember the forward pass, every activation of every layer of every example in the batch, held until the backward sweep arrives to consume it. Depth is a memory bill before it is a compute bill. But a residual network that carries a velocity alongside its state can be run backward, and a forward pass reconstructible from its own endpoint is a forward pass nobody has to remember. The memory bill for depth goes to a constant. Here is that measured, and then the fine print measured too.

What the gradient actually holds hostage#

The chain rule looks memoryless on paper. It is not. To push the loss gradient backward through layer $l$, you need the Jacobian of layer $l$ evaluated at the input that layer saw, which means you need that input. Standard backprop gets it the obvious way: during the forward pass it banks every layer’s activation, and the backward pass withdraws them in reverse order. The bank is the cost. Our network is a modest 784-wide MLP classifier; at depth 8 the compiled training step keeps 13.4 MB of activations, at depth 32 it keeps 44.8 MB, at 128 it keeps 170.7, at 512 it keeps 674.0. The line is straight and its slope is the batch: roughly 1.3 MB per layer, every layer alive simultaneously, because the very first activation is needed by the very last backward step.

Nothing about the gradient wants this. The gradient at the end is a single small object. The hostage-taking comes from one assumption so default it is invisible: that the only way to have the past is to have stored it.

There is another way to have the past. Recompute it. If layer $l+1$‘s output determines layer $l$‘s output, exactly, by algebra, then the backward sweep can walk the states backward as it goes, deriving each layer’s input the moment it is needed and discarding it after use. The bank closes. Memory stops depending on depth at all.

The catch is the if. An ordinary residual block $x_{l+1} = x_l + f(x_l)$ is not invertible in general; two inputs can land on one output and the past is genuinely gone. Which is where the velocity ledger turns into an engineering property.

The block you can run backward#

The momentum residual network updates a state and a velocity together:

Read it backward and every step undoes:

Nothing is approximate. Given the endpoint $(x_L, v_L)$, the entire trajectory is recoverable, layer by layer, using only the same $f$ the forward pass used.

So walk the gradient backward alongside the states, and the bookkeeping writes itself. Suppose the loss gradients with respect to layer $l{+}1$‘s outputs are $(g_x, g_v)$, and ask what they are one layer earlier. The forward equations say $v_{l+1}$ influences the loss twice, once directly and once through $x_{l+1} = x_l + v_{l+1}$, so its total gradient is $g_v + g_x$. From there each forward operation transposes: the $\mu$ that scaled $v_l$ scales its gradient, and the $f$ that fed $x_l$ into the velocity returns a Jacobian-vector product,

Every quantity on the right is available at the moment it is needed: the gradients arrive from the layer above, and $x_l$ has just been re-derived by the inverse. That pair of lines, plus the accumulation of $f$‘s parameter gradients, is the entire backward pass. The companion implements it as a jax.custom_vjp whose saved state is the endpoint and nothing else, and the claim it must pass is severe: not similar gradients, the same gradient.

On the run’s memory instrument, XLA’s own static analysis of the compiled step, the payoff is as flat as the theory says. Standard: 13.4, 44.8, 170.7, 674.0 MB at depths 8, 32, 128, 512. Reversible: 3.2 MB at depth 8, and 3.2 MB at depth 512. The runtime allocator agrees, 681 MB against 90 at depth 512. The price appears where it should: the backward pass re-runs $f$ once per layer, and the measured step time goes from 83 ms to 103, about 24% slower. That is the whole trade, stated in two numbers: at depth 512, 200 times less activation memory for a quarter more time. Gomez et al. (2017) built this trade into RevNets by architectural partition; Sander et al. (2021) showed the momentum form; what this run adds is the ledger of what it costs on one network you have already met.

Friction, and why the rewind cannot be free#

Before trusting that inverse, look at it once more. The backward velocity update divides by $\mu$. On paper, harmless. In floating point, a division by $\mu = 0.9$ is a multiplication of everything, signal and rounding error alike, by $1.11$, once per layer. Two hundred layers of that is a factor of $10^{10}$. The forward pass shrank the velocity’s history with every $\mu$; the rewind must unshrink it, and float32 kept only about seven digits of what there was to unshrink.

Chase that amplification to its source and you arrive at the oldest law in the building. A pendulum without friction retraces its swing if you reverse time; a damped pendulum does not, because friction converts its motion into heat, and heat does not remember which swing it came from. Dissipation destroys information, and $\mu < 1$ is dissipation, the same friction the skip-connections post identified in the ledger. Watch it in the physics where the intuition lives:

The dictionary makes a quantitative prediction, the same shape that post measured for trajectories and now applied to gradients: reconstruction noise grows like $(1/\mu)^L$. Here is the same experiment at the scale of the panel, twelve particles through 64 momentum layers in genuine float32, every operation rounded the way a GPU rounds:

The napkin that locates the cliff#

So the gradient computed by rewind should be exact while $(1/\mu)^L \varepsilon_{32}$ is small, and garbage once it reaches order one. That is an arithmetic you can do before running anything: with float32’s $\varepsilon \approx 1.2 \times 10^{-7}$, the budget crosses 1 near $L^* = \ln(1/\varepsilon)\,/\ln(1/\mu)$, which is 151 layers at $\mu = 0.9$, 31 at $\mu = 0.6$, 13 at $\mu = 0.3$.

The run measured the actual gradient agreement, cosine between the rewind-computed gradient and the stored-activation gradient, on a 3-by-3 grid of $\mu$ and $L$:

Read the middle column against the napkin. At $\mu = 0.6$ the prediction said the cliff sits at 31 layers; the measurement at 32 layers is mid-fall, cosine 0.30. At $\mu = 0.3$ the prediction said 13; by the next measured depth the gradient is already rubble. At $\mu = 0.9$ the prediction said 151, and at 128 the agreement has collapsed to 0.026, the compounding having caught up somewhat ahead of the single-step estimate, which is what an estimate that ignores accumulation across layers should do. Inside the budget the rewind is not approximately right, it is exact to float precision, cosine 1.000000. Outside it there is no gradient at all.

Maclaurin, Duvenaud, and Adams hit this same wall in 2015, rewinding training itself to differentiate through learning, and their repair is telling: they stored the low-order bits that division was about to lose, paying a little memory to make the rewind lossless. The bits you refuse to store, you eventually pay for in noise.

Does a rewound gradient train a network?#

Inside the budget it must, since it is the same gradient, but the run checks anyway rather than trusting the syllogism. At depth 64, $\mu = 0.9$, budget comfortably inside float32’s range, both training modes run to the same schedule: the standard net reaches 81.6% test accuracy, the reversible one 79.7%, from one seed each, with the reversible step 17% slower on the wall clock. The 1.9-point gap is a single-seed reading and we make no claim about its persistence; the structural claims, identical gradients inside the budget and constant memory at every depth, are the measured ones.

And the budget points somewhere specific as it tightens. The cliff exists because $\mu < 1$; push $\mu \to 1$ and $(1/\mu)^L \to 1$, the amplification disappears at every depth. That limit is not a thought experiment, because the $\mu = 1$ machine already exists: the leapfrog classifier steps kick, drift, kick, and each of those undoes by a subtraction, add back the half kick, step back the drift, add back the other half kick. Write its inverse out and there is no division anywhere in it, hence no amplification factor, hence no budget and no cliff: the conservative block rewinds at float precision to any depth, for the same reason it conserves its energy, its dynamics dissipate nothing. The architectures that conserve are the architectures that rewind for free, which are the architectures that train without the memory bill at any depth. Conservation, reversibility, and constant-memory training are not three features. They are one property wearing three coats.

Scope. One run on Kaggle (bundle kgl_blog-revmem-v1) produced every number here: a 784-16-wide momentum MLP on a 10-class problem, activation memory read from XLA’s compiled-step analysis with each depth measured in a fresh process, gradient agreement measured against stored-activation backprop on identical weights and batch, and the training comparison run once per mode at depth 64. The memory constant depends on architecture width and batch, and the 24% time overhead on the cost of $f$; the flatness in depth and the $(1/\mu)^L$ budget do not.

Structure keeps paying#

The trade here has a clean shape, and it is the shape every borrowed integrator has had: structure in the block becomes a guarantee about the network, and the guarantee turns out to be spendable. A block you can invert is a block whose activations nobody has to keep, so training memory stops knowing how deep the network is. The whole cost of that is a quarter more wall clock and a ceiling on depth set by $\mu$. Take the friction out and the ceiling leaves with it, which is why the conservative block of the previous post rewinds at any depth for free. There is one more tool waiting in the same body of numerical analysis, an integrator that picks its own step size. It has nothing to say about memory, and a great deal to say about whether a stack of layers has any right to be uniform.

@misc{bouhsine2026backpropwithoutthememory,
  author       = {Bouhsine, Taha},
  title        = {Backprop Without the Memory},
  year         = {2026},
  month        = {jul},
  howpublished = {\url{https://tahabouhsine.com/blog/backprop-without-the-memory/}},
  note         = {Blog post, Records of the !mmortal Data Scientist}
}