Backpropagation through a deep network requires storing every activation, which is why training memory grows with depth. Sakana AI's DiffusionBlocks finds a way around it: train each block independently, reinterpreting the block-wise updates as the reverse process of a diffusion model. The memory savings are a factor of B, and the quality matches end-to-end backprop.
The core idea
Partition a network into B blocks. Assign each block a "noise range" along a diffusion trajectory, where the noise level represents how close to the target the block's output should be. Train each block independently, only storing activations for one block at a time. The diffusion framing makes the block-wise training consistent, so the blocks compose into a network that works end to end.
The residual-diffusion connection
The insight builds on the Neural ODE observation: residual connections are discretized ODEs. A residual network is a dynamical system, and block-wise updates are steps along a trajectory. Diffusion is exactly that: a process that goes from noise to signal in steps. So training a residual network block-wise is equivalent to learning the reverse process of a diffusion model, with each block responsible for a segment of the trajectory.
How it works
- Partition the L layers into B blocks.
- Assign each block a noise range along the diffusion trajectory.
- Condition each block on its range (e.g. via adaptive layer norm), so it knows where it sits.
- Train one randomly-sampled block per iteration, storing activations for that block only.
The memory saving is a factor of B: you never hold the whole network's activations at once.
Why it matters for inference
This is primarily a training innovation, but it has inference implications worth watching:
- Block-level parallelism. Models trained block-wise could enable block-level parallelism during inference.
- Progressive generation. The diffusion framing suggests early exit at different "noise levels" - a principled way to trade quality for speed.
- Democratized training. Bx memory reduction means larger models can train on smaller clusters, which changes what hardware you need.
DiffusionBlocks is like building a house by having each contractor work independently on their floor, guided by a blueprint that says "your floor is between the 3rd and 4th stories." They never see the whole house, but the blueprint keeps the floors aligned.
The takeaway
Training memory doesn't have to grow with depth. If you can frame the network as a process, you can train it block by block and still get a whole that works.
Sources
- Sakana AI: DiffusionBlocks: the project page.
- DiffusionBlocks paper (arXiv 2506.14202): the full method and experiments.
- OpenReview: the ICLR 2026 version with reviews.
- Neural ODEs: the residual-as-ODE connection that grounds the method.