[2506.14202] DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation
Abstract:End-to-end backpropagation requires storing activations throughout all layers, creating memory bottlenecks that limit model scalability. Existing block-wise training methods offer means to alleviate this problem, but they rely on ad-hoc local objectives and remain largely unexplored beyond classification tasks. We propose $\textit{DiffusionBlocks}$, a principled framework for transforming transformer-based networks into genuinely independent trainable blocks that maintain competitive performance with end-to-end training. Our key insight leverages the fact that residual connections naturally correspond to updates in a dynamical system. With minimal modifications to this system, we can convert the updates to those of a denoising process, where each block can be learned independently by leveraging the score matching objective. This independence enables training with gradients for only one block at a time, thereby reducing memory requirements in proportion to the number of blocks. Our experiments on a range of transformer architectures (vision, diffusion, autoregressive, recurrent-depth, and masked diffusion) demonstrate that DiffusionBlocks training matches the performance of end-to-end training while enabling scalable block-wise training on practical tasks beyond small-scale classification. DiffusionBlocks provides a theoretically grounded approach that successfully scales to modern generative tasks across diverse architectures. Code is available at this https URL .
[2506.14202] DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation --> Computer Science > Machine Learning arXiv:2506.14202 (cs) [Submitted on 17 Jun 2025 ( v1 ), last revised 12 Jun 2026 (this version, v4)] Title: DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation Authors: Makoto Shing , Masanori Koyama , Takuya Akiba View a PDF of the paper titled DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation, by Makoto Shing and Masanori Koyama and Takuya Akiba View PDF HTML (experimental) Abstract: End-to-end backpr
Explore this link on the map →saved by
related reading
- Aman's AI Journal • Primers • Ilya Sutskever's Top 30aman.ai
- What are Diffusion Models? | Lil'Loglilianweng.github.io
- The Annotated Transformernlp.seas.harvard.edu
- Annotated Research Paper Implementations: Transformers, StyleGAN, Stable Diffusion, DDPM/DDIM, LayerNorm, Nucleus Sampling and morenn.labml.ai
- Transformers from scratch | peterbloem.nlpeterbloem.nl
- The Decade of Deep Learning | Leo Gaobmk.sh
- Transformer Math 101 | EleutherAI Blogblog.eleuther.ai
- The Annotated Transformernlp.seas.harvard.edu
- Nemotron_Diffusion_Tech_Report_v1.pdfd1qx31qr3h6wln.cloudfront.net
- Large Language Diffusion Modelsarxiv.org
- The Little Book of Deep Learningfleuret.org
- The Annotated Diffusion Modelhuggingface.co