aiqbench analysis
DiffusionBlocks introduces a novel framework that transforms transformer networks into independently trainable blocks by interpreting residual connections through the lens of denoising processes. This approach effectively mitigates memory bottlenecks associated with end-to-end backpropagation, enabling scalable training on diverse generative tasks without sacrificing performance. By leveraging score matching objectives for localized gradient updates, the method offers a theoretically grounded solution to current scalability limits in large model development.
DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation
arXiv:2506.14202v4 Announce Type: replace-cross 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 https://github.com/SakanaAI/DiffusionBlocks .
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 .
To appear at the 14th International Conference on Learning Representations (ICLR 2026). v4: Fixed typos in experimental details (Appendix E.4)
Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Machine Learning (stat.ML)
arXiv:2506.14202 [cs.LG]
(or arXiv:2506.14202v4 [cs.LG] for this version)
Full text under the publisher's open licence · 229 words
Related coverage
[AINews] Agents disprove Polysemanticity; greatly improve Interpretability & Alignment
Latent Space · Oct 16, 2026
Evaluating task vectors, unlearning and inoculation
LessWrong · Sep 20, 2026
J++ Lens: Jacobian Filtering Enables More Faithful Workspace Lenses
LessWrong · Oct 8, 2026
How Does Changing the Elo of a Chess Transformer Affect its Computations?
LessWrong · Sep 29, 2026