The Core of DiffusionBlocks: Why Block-Independent Training Works

Taught to me by AI (Claude)

Background: The Problem with Conventional Block-wise Training

Attempts to partition neural networks into blocks and train them independently have existed for some time, but they all shared a common problem. The choice of objective function for each block was ad hoc, and there was no theoretical justification for why independent training should work at all.

The Core Idea: Decomposability of the Score Matching Loss

The training loss of diffusion models (score matching loss) takes the following form:

Thinking of this discretely across noise levels σ1,σ2,σ3​:

This is simply a sum of terms. Each term only measures “how well the network can recover the original signal when noise of that level is added,” and the term for σ1\sigma_1 σ1​ never appears in the computation for σ2​ or σ3. This is the mathematical foundation of the independence.

The Key Implementation Idea: Using the Clean Target y as Input

There is one essential design choice required to realize this independence. During training, each block receives as input the clean target y\mathbf{y} y with noise added directly — that is, y+σϵ.

This is explicitly stated in the pseudocode of the paper (Figure 3):

With this design, each block no longer depends on the output of the previous block. It only needs to minimize the loss within its assigned noise range.

This is not a constraint but a free design choice. It shares the same nature as the asymmetry already present in diffusion models, where noisy inputs are used during training while a different procedure is followed at inference time.

Correspondence with Residual Connections: Why This Applies to Transformers

The residual connection

formally coincides with the update rule obtained by applying Euler discretization to the probability flow ODE of the diffusion process. By exploiting this correspondence, existing Transformer-based networks can be reinterpreted as a sequence of discretized Euler steps of a diffusion process, allowing the framework of score matching-based independent training to be applied directly.

It should be noted, however, that this formal correspondence holds for ODEs in general and is not unique to diffusion processes. The essential contribution lies in the idea of using the decomposability of the score matching loss as the theoretical justification for block-independent training.

Summary: What Is New

The decisive difference from prior methods is that block-independent training is no longer an approximation — it is a theoretically justified operation. The core contribution of this paper is the novel use of a well-known property of the score matching loss — its decomposability across noise levels — as the principled foundation for block-wise training.

Coming back to myself — I don’t fully understand it, and at the same time, I can’t quite wrap my head around how AI (Claude) is able to explain something this complex.