Training large neural networks is a VRAM nightmare. The standard approach requires keeping the entire model's activations in memory during backpropagation, and memory needs skyrocket as layers pile up. Now, Sakana AI and the University of Tokyo are taking a different tack. At ICLR 2026, they're unveiling DiffusionBlocks, a training framework that chops a neural network into independent blocks and retrains each one by reframing the update as a reverse-denoising process.
The result? Memory consumption that used to scale linearly with depth now depends only on a single block. During training, the algorithm randomly selects one block to load and update; the rest never touch VRAM or even get computed. That slashes memory usage to 1/B of the original — where B is the number of blocks. Experiments show that block training matches or even beats traditional end-to-end methods across Vision Transformers, DiT image generation, and text generation tasks, all while sipping GPU memory.
DiffusionBlocks also shines on looped transformers, which typically rely on expensive backpropagation-through-time (BPTT) across steps. Instead, by simulating a gradual denoising trajectory, the framework requires only a single forward pass per update during training. At inference, it keeps the original K-iteration mechanism intact — so training gets dramatically cheaper without sacrificing the looped model's capabilities.