DiffusionBlocks: Blockwise Training for Generative Models via
  Score-Based Diffusion
Authors: name: Makoto Shing; name: Takuya Akiba
Published: 2025-06-17 05:44:18

Abstract: Training large neural networks with end-to-end backpropagation creates
significant memory bottlenecks, limiting accessibility to state-of-the-art AI
research. We propose DiffusionBlocks, a novel training framework
that interprets neural network blocks as performing denoising operations in a
continuous-time diffusion process. By partitioning the network into
independently trainable blocks and optimizing noise level assignments based on
equal cumulative probability mass, our approach achieves significant memory
efficiency while maintaining competitive performance compared to traditional
backpropagation in generative tasks. Experiments on image generation and
language modeling tasks demonstrate memory reduction proportional to the number
of blocks while achieving superior performance. DiffusionBlocks provides a
promising pathway for democratizing access to large-scale neural network
training with limited computational resources.