Sakana AI releases DiffusionBlocks to train neural networks one block at a time, cutting training memory up to 8x
It treats the network's forward pass as a diffusion-like denoising process.
#18hardmaru@HARDMARUFor 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 our #ICLR2026 paper (https://arxiv.org/abs/2506.14202), we matched end-to-end performance across ViTs, DiTs, and LLMs. We did this while training just one isolated block at a time.
Sakana AI@SakanaAILabsIntroducing DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation http://pub.sakana.ai/diffusionblocks What if we didn’t have to hold an entire neural network in memory to train it? Standard neural net training optimizes all parameters jointly. As a result, the memory required during training grows linearly with the depth of the network. In our #ICLR2026 paper, we propose DiffusionBlocks, a principled framework to train networks one block at a time, drastically reducing memory requirements while matching end-to-end performance. With DiffusionBlocks, we split the network into blocks and train them one at a time, so you only need memory for a single block. How? We explicitly assign each block a role: to move the representation a little closer to the target than the block before it did. That role turns out to be precisely what a diffusion model does, step by step. Each block only needs to optimize its own objective and can be trained independently. We validated this across five different architectures: • ViT • DiT • Masked diffusion • Autoregressive transformers • Recurrent-depth transformers In each case, performance is competitive with end-to-end training while using a fraction of the memory. This perspective also extends naturally to recurrent-depth (Looped) transformers, which apply the same network iteratively and normally require expensive backpropagation through time (BPTT). Viewed through DiffusionBlocks, we can replace those multiple iterations with a single forward pass during training. Read our paper and code, to learn more. Paper: https://arxiv.org/abs/2506.14202 GitHub: https://github.com/SakanaAI/DiffusionBlocks 🐟
@hardmaru This is essentially "the cat paper" in the age of LLM. @quocleix 😁
hardmaru@hardmaruFor 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 our #ICLR2026 paper (https://arxiv.org/abs/2506.14202), we matched end-to-end performance across ViTs, DiTs, and LLMs. We did this while training just one isolated block at a time.
#465Furong Huang@FURONGHResNets walked so DiffusionBlocks could denoise 😄
Block-wise training did not disappear — it just went off, got a diffusion-era glow-up, and came back cooler.
This is especially fun to see because in our ancient ICML 2018 BoostResNet paper, we studied a related idea: learning deep ResNet blocks sequentially, one block at a time, through the lens of boosting theory. Each block is not just another layer; it is a stepwise improvement to the model.
It is exciting to see this “build the model block by block” instinct showing up again in modern generative models. Old ideas, new objectives, much bigger compute — and suddenly the blocks are back on tour.
With @jordan_t_ash, @JohnCLangford, and Robert Schapire: https://arxiv.org/abs/1706.04964
Sakana AI@SakanaAILabsIntroducing DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation http://pub.sakana.ai/diffusionblocks What if we didn’t have to hold an entire neural network in memory to train it? Standard neural net training optimizes all parameters jointly. As a result, the memory required during training grows linearly with the depth of the network. In our #ICLR2026 paper, we propose DiffusionBlocks, a principled framework to train networks one block at a time, drastically reducing memory requirements while matching end-to-end performance. With DiffusionBlocks, we split the network into blocks and train them one at a time, so you only need memory for a single block. How? We explicitly assign each block a role: to move the representation a little closer to the target than the block before it did. That role turns out to be precisely what a diffusion model does, step by step. Each block only needs to optimize its own objective and can be trained independently. We validated this across five different architectures: • ViT • DiT • Masked diffusion • Autoregressive transformers • Recurrent-depth transformers In each case, performance is competitive with end-to-end training while using a fraction of the memory. This perspective also extends naturally to recurrent-depth (Looped) transformers, which apply the same network iteratively and normally require expensive backpropagation through time (BPTT). Viewed through DiffusionBlocks, we can replace those multiple iterations with a single forward pass during training. Read our paper and code, to learn more. Paper: https://arxiv.org/abs/2506.14202 GitHub: https://github.com/SakanaAI/DiffusionBlocks 🐟