LIT-004: Training Deep Nets with Sublinear Memory Cost

Chen et al. (2016) — ARXIV-1604.06174

Key takeaways

• O(sqrt(n)) memory to train an n-layer network, for one extra forward pass per minibatch. That is the headline and the whole shape of the result: the memory saving is asymptotic, the compute cost is a constant factor
• The trade is a curve, not a switch. The same analysis gives O(log n) memory for O(n log n) extra forward computation. So the design question is where to sit on it, which is what makes this a rule rather than a trick
• Measured: a 1,000-layer deep residual network, 48G to 7G, +30% running time, on ImageNet. Significant reductions are also reported for recurrent networks on very long sequences — the case that carries over to long-context language models
• Computation-graph analysis does the rest, automating in-place operation and memory sharing. The contribution is a systematic method over the graph, not a hand-tuned schedule for one architecture

Standing in the anthology

The general memory/compute trade the record recommended only for attention. SOTA-249 is now filed from it. Before that, SOTA-087 recommended recomputation for attention — better evidenced than this, since LIT-074 argues it from HBM traffic with a mechanism this paper does not have — and nothing stated the trade itself, despite the memory accounting underneath SOTA-017, SOTA-019 and SOTA-031 assuming activations are recomputed.

That is DP-007 in its ordinary form rather than its dramatic one: recomputation is universal enough that nobody writes the rule down.

Why it took three passes to get here, which is the part worth keeping. This note carried four bullets naming the category — “introduces gradient checkpointing”, “trade-off between memory and computation” — and survived the pass that cleared the rest of them, because that pass selected on a bullet signature and a document can be equally empty without matching one. #121 then audited it and deliberately filed nothing, on the grounds that filing a practice from four bullets describing a category is the failure two passes had been spent correcting. That was right. What changed is not the judgement but the input: the paper’s own abstract states the algorithm, both asymptotics and a measured before-and-after, so the practice is filed from what the paper says rather than from what the note said about it.

Read — NOTE-184. The reading was the thing #85 actually asked for, and it earns its keep on the limits rather than the claim: the allocation half of the paper (liveness, in-place, memory sharing) deliberately carries no practice, because it is what a framework does for you and this record’s distributed-optimization entries are about decisions a person makes.