Why can activation checkpointing around a stochastic layer corrupt a training run?

## The assumption checkpointing rests on Activation checkpointing discards a segment's interior and rebuilds it later by running the segment forward a second time from its retained input. That is only sound if the segment is a pure function: same input, same parameters, same output. Every deterministic layer satisfies this. Anything that consumes randomness or mutates state as a side effect of running forward does not. ## What goes wrong with a stochastic layer A layer that samples random values during training draws from a generator whose state advances with every draw. On the first forward pass through the segment it draws one set of values. That set shapes the intermediates, which shape the loss the optimiser is about to differentiate. Later, when the backward pass reaches the segment and triggers the recompute, the generator has moved on — every other draw in the step has consumed from the same stream — so the second pass draws a different set. The rebuilt intermediates are the activations of a *different* sampled network. The gradients then formed are not the gradients of the loss that was actually computed. They are a mismatched pair: a loss from one sample, derivatives from another. Nothing raises an error, because every tensor has a valid shape and a finite value. ## How it shows up The symptoms are quiet and easy to misattribute: - The checkpointed run no longer matches an un-checkpointed run with the same configuration, even though checkpointing is supposed to be mathematically neutral. - Training is noisier than the baseline, and convergence is slower or reaches a worse plateau, because every step adds a mismatch term that does not average to a useful direction. - Effects scale with how much randomness the segment contains: a model with stochastic behaviour in every block is hit far harder than one with a single stochastic layer near the head. - Small models often survive it, which is what makes it dangerous — the bug is discovered only after it is scaled up. A candidate who says the run will crash has not thought it through. Silent quality loss is the whole problem. ## The fix Capture and restore the generator state. On entry to a checkpointed segment, record the state of every generator that segment will consume from. Before recomputing that segment, restore the recorded state so the second pass draws exactly the values the first one drew. The recompute then reproduces the original intermediates bit for bit, and the gradients match the loss again. This is per-segment bookkeeping, not a global seed. Setting one seed at the start of training does not help at all: both forwards still consume from the same advancing stream, just from different positions in it. Nor is disabling stochastic behaviour during the recompute a fix — that reproduces a different network again, this time a deterministic one, and biases the gradients in the opposite direction. The requirement is exact reproduction of the first pass, and only saving and restoring the state achieves it. Note what this bookkeeping costs: a little state per segment and one restore per recompute. It is negligible against the recompute itself, which is why correct implementations simply always do it. ## The wider class: side effects in the forward pass Randomness is the loudest member of a family. Anything a layer *does* while running forward, beyond producing its output, will happen twice under checkpointing: - A layer that maintains a running statistic updated on each forward pass will update it twice per step, so the statistic moves at double the intended rate and drifts from the value an un-checkpointed run would hold. - A counter, a cache, a queue written during the forward pass gets two writes per step. - Any per-sample statistic computed inside the segment must be reproducible from the segment input alone; if it is, the recompute regenerates it exactly and nothing is lost, which is the good case. The general rule for reviewing a checkpointed model: for each segment, ask what the forward pass does other than return a value. If the answer is anything at all, that side effect now happens twice, and either it must be made idempotent, suppressed on the recompute pass, or the boundary must be redrawn so the offending layer sits outside the checkpointed region. ## Interview framing This is a favourite senior question because it separates people who have read the flag description from people who have debugged a training run. The tell of the second group is that they describe the symptom as a silent quality regression against a baseline, name generator state as the mechanism, and then generalise to forward-pass side effects without being prompted.