DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation

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 .

💡 Research Summary

DiffusionBlocks tackles the memory bottleneck that plagues modern deep learning models trained with conventional end‑to‑end backpropagation. Standard training requires storing activations for every layer, causing memory consumption to grow linearly with depth and limiting both batch size and model size. Existing block‑wise training approaches attempt to partition a network into smaller components that can be trained independently, but they rely on ad‑hoc local objectives, lack theoretical justification, and have been demonstrated mainly on classification tasks and simple architectures. Consequently, they often underperform full backpropagation and do not generalize to transformer‑based generative models.

The core insight of DiffusionBlocks is a rigorous reinterpretation of residual connections as discretized steps of a continuous‑time diffusion process. Building on the established link between residual networks and ordinary differential equations, the authors show that the Euler discretization of the reverse probability‑flow ODE used in score‑based diffusion models yields exactly the same additive update form as modern residual blocks (including transformers). By treating each residual block as a denoising step at a particular noise level, the whole network can be viewed as a sequence of denoisers that collectively approximate the reverse diffusion trajectory.

To exploit this view, the authors partition the overall noise interval (