Training large neural networks has always been a memory hog. Traditional methods keep the entire model’s activations alive during backpropagation, and memory demands skyrocket as layer count grows. To lower the barrier, Sakana AI, in collaboration with the University of Tokyo, unveiled DiffusionBlocks at ICLR 2026. The new framework splits a neural network into multiple modules and reframes block updates as a reverse denoising process, allowing each block to be trained independently.
This block training shatters the linear scaling of memory with depth. During training, the algorithm randomly picks one block to load and update; the others sit idle in storage, not touching GPU memory. The result: memory consumption drops to just 1/B of the original, where B is the number of blocks. Experiments show that this approach not only sips memory but also matches—or even beats—traditional end-to-end training on vision transformers, DiT image generation, and text generation tasks.
For looped transformers (which repeatedly iterate over the same layers), DiffusionBlocks offers a unique optimization. These models typically rely on costly backpropagation-through-time (BPTT). DiffusionBlocks simulates a gradual approximation process, needing only a single forward pass per parameter update during training. At inference, it retains the original K-iteration mechanism, dramatically cutting training compute.