DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation

Their summary:

For over a decade, we’ve accepted that end-to-end backprop is the only way to train deep networks. But holding the entire network in memory all at once is why AI training is hitting a resource wall.
.
We found a new way to break the network into blocks and train them independently. The trick? Treating the network’s forward pass like a diffusion model denoising a signal.
.
This reinterpretation slashes the memory needed to train deep models. In this paper presented at ICLR 2026, we matched end-to-end performance across ViTs, DiTs, and LLMs. We did this while training just one isolated block at a time.

Their technical blog: https://pub.sakana.ai/diffusionblocks/

From what I understand, if this works at a scale, it will allow training with 3-4x memory savings all around (weights, gradients, optimizer, activations).

Also makes training more parallelizable with less communication, can improve Looped Transformer training, and can improve diffusion model inference efficiency. Really hope it scales well.