DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation

Train each block of a network on its own, as a denoiser for one range of noise levels.

A residual layer adds a correction to its input, and so does one Euler step of a diffusion model's reverse process. DiffusionBlocks builds on that match: no gradient crosses from one block to another, and training memory falls by the number of blocks.

A 24-layer image model trained as three independent 8-layer blocks reached a test FID (Fréchet Inception Distance, lower is better) of 10.63 on ImageNet; the same model trained end to end scored 12.09.

The memory cost of backprop

Backpropagation computes the gradient for layer 1 from quantities produced by every layer above it. To do that, the forward pass has to keep each layer's activations in memory until the backward pass reaches that layer. Training memory therefore grows linearly with depth: a 24-layer Transformer holds 24 layers of activations at the peak of every step, along with 24 layers of weights, gradients and optimizer state.

Block-wise training cuts the network into pieces and trains each piece with its own loss, so only one piece is in memory at a time. The idea is old (greedy layer-wise pretraining from 2006, Hinton's Forward-Forward from 2022, and several local-loss methods in between). None of these methods had a derived target for what a middle block should output, so each method invented a local objective, for example a contrastive score for "is this a real input", and the result trailed end-to-end training. Most of these objectives also only make sense for classification. On the classification benchmark in this paper, a Vision Transformer trained with Forward-Forward reaches 7.85% accuracy on CIFAR-100, against 60.25% for ordinary backprop.

DiffusionBlocks (Makoto Shing, Masanori Koyama and Takuya Akiba; Sakana AI and the University of Tokyo; ICLR 2026) derives the local objective instead of inventing one. A diffusion model is trained with a loss that splits into independent pieces, one per noise level. The paper reads a stack of residual layers as the steps of a diffusion model's generation process, assigns each block of layers a range of noise levels, and trains each block with the diffusion loss restricted to its range. The same ViT trained this way in three blocks reaches 59.30%.

The argument needs four pieces of diffusion background, taken in order: how data is noised, why a denoiser computes the direction back toward the data, how generation becomes an ordinary differential equation solved with Euler steps, and why a residual layer has the form of one Euler step. After that come the method, its partition of the noise range, the memory accounting and the results.

Adding noise to data

Take a data point $\mathbf{y}$ (an image, a label embedding, a vector of any kind) and add Gaussian noise with standard deviation $\sigma$:

This is the variance-exploding (VE) convention, which the paper takes from Karras et al.'s EDM paper. The signal $\mathbf{y}$ is never scaled down; the noise variance $\sigma^2$ grows without bound as $\sigma$ rises. The paper's noise levels run from $\sigma_{\min} = 0.002$ to $\sigma_{\max} = 80$, for data scaled to a standard deviation of about 0.5. At $\sigma = 80$ the data is buried under noise 160 times its own spread, so every starting sample looks the same. (DDPM uses the other common convention, variance-preserving, which multiplies the signal by a factor below one as it adds noise so the total variance stays near 1. Formulas from one convention do not transfer to the other without rescaling.)

Real data occupies a thin region of its space. Almost no random assignment of pixel values is a photograph, so the photographs sit on a low-dimensional, curved surface (a manifold) inside the space of all pixel arrays. Noising pushes each point off that surface; generation has to bring a random point back onto it.

Score and denoiser

Write $p_\sigma$ for the distribution of $\mathbf{z}_\sigma$: the data distribution blurred by a Gaussian of width $\sigma$. Its score is the gradient of the log-density with respect to the point, $\nabla_{\mathbf{z}} \log p_\sigma(\mathbf{z})$. At every location the score is a vector pointing in the direction in which the noised density rises fastest. (This is the gradient with respect to the data point, not the model parameters. Statisticians use "score" for the parameter gradient too.)

The log removes the normalizing constant. A density written as $\tilde p/Z$ has $\nabla \log(\tilde p / Z) = \nabla \log \tilde p$, because $Z$ does not depend on the point, so the score can be learned without ever computing $Z$, an integral over the whole space.

At small $\sigma$ the density is concentrated around each piece of data and the score points sharply to the nearest one. At large $\sigma$ the pieces have blurred into a single broad hill and the score points gently toward the overall center.

Given a noisy point $\mathbf{z}$, what is the best single guess of the clean $\mathbf{y}$ that produced it? Under squared error, the best guess is the average of every clean point that could have been noised into $\mathbf{z}$, weighted by how likely each is: the conditional mean $\mathbb{E}[\mathbf{y}\mid\mathbf{z}]$. A function that returns it is a denoiser, written $D(\mathbf{z},\sigma)$. Tweedie's formula, from Robbins's 1956 work on empirical Bayes, says that for Gaussian noise the denoiser and the score carry the same information:

The right-hand side is the vector from the current point to the best guess of the clean data, divided by the noise variance $\sigma^2$ (the variance, not the standard deviation, in this VE convention). The derivation uses only the chain rule. The noised density is $p_\sigma(\mathbf{z}) = \int \mathcal{N}(\mathbf{z};\mathbf{y},\sigma^2\mathbf{I})\,p(\mathbf{y})\,d\mathbf{y}$, and the point $\mathbf{z}$ appears only inside the Gaussian, whose gradient is the Gaussian times $(\mathbf{y}-\mathbf{z})/\sigma^2$. Differentiate under the integral and divide by $p_\sigma(\mathbf{z})$:

The ratio of integrals is a posterior average by Bayes' rule, and the posterior mean $\mathbb{E}[\mathbf{y}\mid\mathbf{z}]$ is the denoiser $D$, which gives (2). The identity is exact for the ideal denoiser. The paper writes it with $\approx$ because a trained network $D_{\boldsymbol{\theta}}$ only approximates the ideal one.

Equation (2) lets you learn the score without ever having a true score to supervise against, because a denoiser can be trained by plain regression. Take clean data, add noise yourself so that you know the correct output, and penalize the squared distance between the network's guess and the clean original:

The noise level $\sigma$ is drawn from a training distribution, and $w(\sigma)$ weights each level so that the rarely drawn and the often drawn levels contribute gradients of similar size. Here $D_{\boldsymbol{\theta}}$ predicts the clean data $\mathbf{y}$; DDPM's network predicts the noise $\boldsymbol{\epsilon}$ instead, and since $\mathbf{z} = \mathbf{y} + \sigma\boldsymbol{\epsilon}$, either prediction gives the other.

Loss (3) is an average over noise levels, and the term for any one $\sigma$ depends only on the network's output at that $\sigma$. Nothing in the loss connects the denoising problem at $\sigma = 5$ to the one at $\sigma = 0.1$ except that both use the same weights $\boldsymbol{\theta}$.

Because $D$ is an average, it looks blurry at high noise: many clean points are plausible sources of the same $\mathbf{z}$, and the mean of many faces is not itself a face.

Generation as an ODE

Song et al. (2021) showed that the noising process can be run backward in two ways with the same distribution of points at every noise level: a random one (a stochastic differential equation) and a deterministic one, the probability-flow ODE. In the VE convention it reads:

This is equation (1) of the paper. It assigns a velocity to every point at every noise level. Start from a sample of pure noise at $\sigma_{\max}$, integrate while $\sigma$ decreases toward 0, and the point ends on the data. The minus sign and the single factor of $\sigma$ come from matching the heat equation that $p_\sigma$ obeys as noise is added. Written as a continuity equation (density changes only because probability mass moves), the velocity that reproduces that spreading is $-\sigma\,\nabla_{\mathbf{z}}\log p_\sigma$.

Substitute Tweedie's formula (2) to replace the score with the denoiser:

A computer solves this with small steps. Euler's method moves from $\sigma_{\text{prev}}$ to a lower $\sigma_{\text{next}}$ by multiplying the current velocity by the change in $\sigma$:

Generation runs $\sigma$ downward, so $\sigma_{\text{next}} - \sigma_{\text{prev}}$ is negative, and it multiplies $\mathbf{z} - D$, which points away from the guess. The two signs cancel. Define the positive step size $\Delta\sigma = \sigma_{\text{prev}} - \sigma_{\text{next}}$ and the update becomes:

Each step moves the point the fraction $\Delta\sigma/\sigma_{\text{prev}}$ of the way toward the denoiser's current guess. The denoiser sets the direction and the noise schedule sets the fraction. With numbers: at $\sigma_{\text{prev}} = 0.677$ and $\sigma_{\text{next}} = 0.301$, the fraction is $0.376/0.677 = 0.555$. A coordinate at 5 with a guess of 2 moves to $5 + 0.555 \times (2-5) = 3.33$.

The paper prints this step with the opposite sign. Its Eq. (4) reads $\mathbf{z} + (\Delta\sigma/\sigma)(\mathbf{z} - D)$ with the same positive $\Delta\sigma$, which moves the coordinate from 5 to 6.67, away from the guess. The official code uses the sign in (5), and the Provenance note at the end records the discrepancy.

If the guess $D$ stayed fixed over a step, the ODE could be solved by hand: $\mathbf{z}(\sigma) = D + (\mathbf{z}_{\text{prev}} - D)\,\sigma/\sigma_{\text{prev}}$, a straight line that reaches $D$ at $\sigma = 0$. The distance to the guess shrinks in proportion to the noise level. Euler step (5) lands on that line exactly: from 0.677 to 0.301 the offset of 3 becomes $3 \times 0.301/0.677 = 1.33$, and $2 + 1.33 = 3.33$. Euler's error comes only from $D$ changing during the step.

Euler's method is first-order: when $D$ does change, the error per step shrinks like $\Delta\sigma^2$ and the accumulated error like $\Delta\sigma$. More steps track the true path more closely, which is why higher-order samplers exist.

Over a full run the guess changes a lot. At high $\sigma$ the denoiser cannot tell which data point a sample came from, so $D$ is close to the mean of all the data and every sample drifts toward the middle. As $\sigma$ falls, $D$ switches to the nearest cluster, and the field splits into basins, one per cluster, with sharp borders between them. The $1/\sigma$ in the velocity also grows as $\sigma \to 0$. The score at level $\sigma$ only resolves structure coarser than $\sigma$, so diffusion settles the coarse layout first and the detail last, and the middle noise levels, where the field reorganizes from one center into many, carry most of the decisions. Figure 7 below runs this process end to end.

Residual layers as Euler steps

A residual layer, as in ResNet and every Transformer since the original, computes a correction and adds it to its input:

In (6), $\mathbf{z}_{\ell-1}$ is the activation entering layer $\ell$ (in a Transformer, a tensor of shape $n \times d$ for $n$ tokens of width $d$), and $f_{\theta_\ell}$ is that layer's attention and MLP with weights $\theta_\ell$. Placed next to the Euler step (5):

Both add a computed correction to the current state. That a residual network is Euler's method applied to some ODE, with step size 1 and depth playing the role of time, is the observation behind Neural ODEs and Haber and Ruthotto's work on stable architectures; each layer advances the state one tick along a trajectory.

The match is in form only. A plain residual layer has step size 1 and receives no noise level. The Euler step has a coefficient $\Delta\sigma/\sigma_{\text{prev}}$ that changes from step to step, and its velocity $(\mathbf{z} - D)/\sigma$ depends on $\sigma$ explicitly. DiffusionBlocks does not ask the layers to learn either. A block computes only the guess $D$, and the update around it is assembled from the known schedule, which supplies $\Delta\sigma/\sigma_{\text{prev}}$. The noise level enters the block as an extra input, so one set of weights can behave differently at $\sigma = 0.6$ and at $\sigma = 0.3$.

Training blocks independently

The conversion has three steps (§3.1 of the paper). Partition the layers. Group the $L$ layers into $B$ consecutive blocks. Block $b$ is the composition of its layers, written $\bar f_{\theta_b}$.

Assign noise ranges. Split $[\sigma_{\min}, \sigma_{\max}]$ into $B$ consecutive intervals. The paper numbers blocks from the noisy end: block 1 owns the highest noise levels, where generation starts, and block $B$ owns the lowest. How the cuts are placed is the subject of a later section.

Add noise conditioning and a second input. Each block now receives the pair $(\mathbf{x}, \mathbf{z})$: the task input $\mathbf{x}$ (the image to classify, the previous tokens, a class label) and the current noisy estimate $\mathbf{z}$ of the target. Both enter as tokens. The noise level enters through AdaLN, the mechanism DiT uses: a small network maps $\sigma$ to a scale and shift applied at each layer normalization. The block's output tokens at the $\mathbf{z}$ positions are its guess $D$, and generation applies update (5) with that guess.

Loss (3) splits by noise level, and now each range of noise levels has its own parameters. Block $b$ minimizes the paper's Eq. (6):

In (7), $p^{(b)}$ is the training noise distribution restricted to block $b$'s interval and renormalized. The paper writes the inner loss as a generic "Loss", typically squared error; for labels and tokens it is cross-entropy (explained below). The input to block $b$ is the real target $\mathbf{y}$ with fresh noise, not the output of block $b-1$, and the target is $\mathbf{y}$ itself. A training step for one block needs no other block, stores activations for that block's layers only, and sends gradients into that block's parameters only.

A worked example with the shape of the paper's autoregressive experiment (a 12-layer Transformer, tokens of width $d$, $B = 4$ blocks of 3 layers each) and the squared-error loss of (7) that image targets use. With EDM's noise distribution, the equal-probability cuts derived in a later section land at $\{0.002,\ 0.134,\ 0.301,\ 0.677,\ 80\}$, so block 1 owns $[0.677,\ 80]$, block 2 owns $[0.301,\ 0.677]$, block 3 owns $[0.134,\ 0.301]$, and block 4 owns $[0.002,\ 0.134]$. One step for block 2:

• Draw a batch: conditioning $\mathbf{x}$ and target $\mathbf{y}$ of shape $n \times d$.
• Draw $\sigma$ from the noise distribution inside block 2's band, say $\sigma = 0.45$.
• Build $\mathbf{z} = \mathbf{y} + 0.45\,\boldsymbol{\epsilon}$ with fresh Gaussian $\boldsymbol{\epsilon}$ of the same shape.
• Run block 2's three layers on $(\mathbf{x}, \mathbf{z})$ with $\sigma = 0.45$ fed to AdaLN, and read off $\hat{\mathbf{y}}$, shape $n \times d$.
• Compute the scalar $w(0.45)\,\lVert \hat{\mathbf{y}} - \mathbf{y}\rVert^2$. EDM's weight is $w(\sigma) = (\sigma^2 + 0.5^2)/(0.5\,\sigma)^2$, which is 8.94 at $\sigma = 0.45$.
• Backpropagate. Gradients reach $\theta_2$ and nothing else; blocks 1, 3 and 4 are not in the computation graph.

The listing adds one detail the list skipped: EDM's preconditioning, which scales the block's input and output by $\sigma$-dependent factors so that the network's own target has roughly unit size at every noise level. The official code uses these constants, with $\sigma_{\text{data}} = 0.5$.

# one training step for block b of B = 4 (12 layers, 3 per block)
x, y  = sample_batch()                  # y: target, shape [n, d]
b     = 2                               # the paper picks b at random per step
sigma = sample_lognormal(band[b])       # band[2] = [0.30, 0.68], say 0.45
z     = y + sigma * randn_like(y)       # noised target, shape [n, d]

# EDM preconditioning around the 3 layers (sigma_data = 0.5)
c_in   = 1 / sqrt(sigma**2 + 0.25)      # 1.487 at sigma = 0.45
c_skip = 0.25 / (sigma**2 + 0.25)       # 0.553
c_out  = 0.5 * sigma * c_in             # 0.335
F      = block[b](x, c_in * z, log(sigma) / 4)   # AdaLN reads the sigma input
y_hat  = c_skip * z + c_out * F         # the denoised guess D(z, sigma)

w    = (sigma**2 + 0.25) / (0.5 * sigma)**2      # 8.94 at sigma = 0.45
loss = w * ((y_hat - y)**2).mean()      # labels and tokens: cross-entropy
loss.backward()                         # gradients reach block[b] only
opt[b].step()

In the paper's training loop each iteration picks one block uniformly at random and takes this step for it. The blocks can also be trained on separate GPUs at the same time, since they exchange nothing.

Why can the blocks be trained apart and still work together? Generation by (5) needs, at each noise level, the correct denoiser for that level and nothing else. That function is the same whether one network learns every level or four blocks each learn a quarter of them, and every block is trained against the same fixed target, real data with fresh noise. Earlier block-wise methods trained each block toward a made-up local goal and hoped the pieces would compose. Here each block's loss is exactly the denoising loss (3) restricted to its band, so a block that minimizes its loss is the right denoiser for its band.

Generation differs from training in one respect. At generation, block 2 receives whatever block 1 produced, errors included, and a small error can move the state slightly off the distribution $p_\sigma$ that block 2 was trained on. This mismatch between training on truth and running on your own output is called exposure bias. A single diffusion model has it too, since each of its sampling steps consumes the previous step's output, and splitting the model does not change the target each weight is trained toward. To smooth the handoff, the paper lets neighbouring bands overlap slightly during training (covered with the partition below).

Generating with the blocks

Generation starts from $\mathbf{z} = 80\,\boldsymbol{\epsilon}$ and applies Euler step (5) along a decreasing list of noise levels. At each step the sampler calls the one block whose band contains the current $\sigma$ (the paper's Figure 3 algorithm: "select block $b$ where $\sigma_i \in [\sigma_b, \sigma_{b-1}]$"). Between steps only $\mathbf{z}$ is carried forward; the conditioning $\mathbf{x}$ is fed to every call.

# generation: N Euler steps, each served by one block
sig = equal_mass_levels(N, 80, 0.002) + [0]   # N quantiles, descending
z   = 80 * randn(n, d)                  # pure noise at sigma_max
for i in range(N):
    b     = owner(sig[i])               # the block whose band holds sig[i]
    y_hat = D(block[b], x, z, sig[i])   # preconditioned, as in training
    z     = z + (sig[i] - sig[i+1]) / sig[i] * (y_hat - z)   # Euler step (5)
return z                                # N = 50 for images; N = B for the LM

The number of Euler steps $N$ is a separate setting from the number of blocks $B$. For image generation the paper uses 50 Euler steps with classifier-free guidance at scale 2.0, and each of the 50 steps runs one 4-layer block of a 12-layer DiT (with $B = 3$). The language model with $B = 4$ uses 4 steps, one per block, so it runs the same 12 layers per token as the original network. The official code places the $N$ noise levels at equally spaced quantiles of the training noise distribution, the same quantiles that define the block edges, so each block serves about $N/B$ steps (Appendix H says the 50 image-generation steps are assigned the same way).

The error toggle in Figure 7 adds a fixed offset of 0.63 to one block's guess. With the error in block 1, which serves only the high-noise steps, the samples still land on the clusters (one of 240 misses, as with no error), but the upper-right cluster (the direction of the offset) collects 107 of the 240 samples instead of 68. The offset in block 1 changes which cluster a sample lands in, and the later blocks do not move it to a different cluster. With the error in the last block, which serves the final step to $\sigma = 0$, all 240 samples end off the data, because the last step moves the point all the way to the faulty guess.

What gets noised

For image generation the target $\mathbf{y}$ is an image (for ImageNet, the 32×32×4 latent of a pretrained VAE, as in DiT) and the blocks denoise it directly. A classifier or a language model has no image to denoise, so the paper noises a continuous embedding of the answer instead (Appendix B).

Vision Transformer. The target is a learned embedding of the correct class, a vector of the model's width, L2-normalized so that the learned embeddings cannot all collapse to the same vector. The noisy label embedding $\mathbf{z}_\sigma = \mathbf{y}_{\text{emb}} + \sigma\boldsymbol{\epsilon}$ is added to the sequence as one extra token, and the image patches are the conditioning $\mathbf{x}$. The shared input embedding and the classification head sit outside the blocks. Training uses $w(\sigma)$-weighted cross-entropy on class logits read from the denoised embedding, not squared error; at inference the code turns the logits into a denoised embedding as the softmax-weighted average of the class embeddings.

Autoregressive language model. Noise is added to the embedding of the next token, and the block predicts that token's clean embedding while conditioning on clean embeddings of the earlier tokens. During training each sequence is fed twice, once clean and once noisy, with an attention mask that lets each noisy token see only the clean tokens before it (a technique from Block Diffusion). This doubles the sequence memory. The loss is cross-entropy over the vocabulary.

Masked diffusion language model. Here the noise is masking, not Gaussian: at time $t$ each token is replaced by [MASK] with probability $1 - \alpha(t)$. The paper shows that the training loss accumulates uniformly in $\alpha$, so it cuts the schedule into equal decrements, $\alpha_b = 1 - b/B$. With $B = 3$, one block handles mask rates from 0 to 1/3, one from 1/3 to 2/3, and one from 2/3 to every token masked. Appendix D gives the derivation.

The Euler form (5) needs $\mathbf{z}$ and $\mathbf{z}_{\text{next}}$ to have the same shape. In a Transformer every block maps $n$ tokens of width $d$ to $n$ tokens of width $d$, so the condition holds. A U-Net changes resolution between stages, so its stages do not fit the recipe yet.

Splitting the noise range

Where should the $B - 1$ cuts go? During training $\sigma$ is not drawn uniformly. Following EDM, $\log\sigma$ is drawn from a normal distribution, $\log\sigma \sim \mathcal{N}(P_{\text{mean}}, P_{\text{std}}^2)$ with $P_{\text{mean}} = -1.2$ and $P_{\text{std}} = 1.2$, which puts most of the draws at middle noise levels, around $\sigma \approx e^{-1.2} = 0.30$. At very high $\sigma$ the best a denoiser can do is predict roughly the data mean, and at very low $\sigma$ the input is nearly clean; at the middle levels coarse structure turns into specific content.

A block's share of the probability mass is its share of the training samples. The obvious split, equal widths in $\sigma$ (the formula $\sigma_{\min} + b(\sigma_{\max}-\sigma_{\min})/B$ in §3.3), gives the lowest block $[0.002,\ 20]$ for $B = 4$, which holds 99.98% of the draws; the other three blocks would see almost no training data. DiffusionBlocks instead picks cuts that give every block exactly $1/B$ of the mass. Because $\log\sigma$ is Gaussian, the cut points come from the inverse normal CDF $\Phi^{-1}$, evaluated at evenly spaced probabilities:

where $q_{\min}$ and $q_{\max}$ are the normal CDF values at the ends of the range, $\Phi\big((\log\sigma_{\min/\max} - P_{\text{mean}})/P_{\text{std}}\big)$. For $B = 4$ this gives $\{0.002,\ 0.134,\ 0.301,\ 0.677,\ 80\}$. The two middle blocks each span a factor of 2.25 in $\sigma$; the tails span factors of 67 (block 4) and 118 (block 1). Each holds a quarter of the draws, so the busy middle gets narrow blocks and the ends get wide ones.

The paper's ablation compares against the stronger of the two uniform baselines, equal widths in $\log\sigma$ (Appendix E.6.2). On CIFAR-10 with a 12-layer DiT in three 4-layer blocks, that baseline gets FID 43.53 and equal-probability cuts get 38.03. FID (Fréchet Inception Distance) compares statistics of generated and real images, and lower is better. Equal-probability wins at every layer split tested (for example 3-6-3 layers: 43.59 against 41.64), and the plain 4-4-4 split is the best of them, so the layers can be divided evenly without tuning. With $B = 3$, the log-uniform cuts at 0.068 and 2.34 give the three blocks about 11%, 85% and 4% of the training draws.

Overlap. Each block trains on a range stretched past its own edges. For a block owning $[\sigma_b, \sigma_{b-1}]$, the paper sets $\alpha_b = (\sigma_{b-1}/\sigma_b)^\gamma$ and trains on $[\sigma_b/\alpha_b,\ \alpha_b\,\sigma_{b-1}]$ (Appendix C). The stretch is a fixed fraction $\gamma$ of the block's width in log σ, so wide tail blocks stretch further in absolute terms. The default is $\gamma = 0.05$, and text generation uses $\gamma = 0.1$. For the $B = 4$ language model at $\gamma = 0.1$, the middle blocks get $\alpha = 2.25^{0.1} = 1.084$, about 8% past each edge, so block 3 trains on $[0.124,\ 0.327]$ instead of $[0.134,\ 0.301]$. A training draw at $\sigma = 0.31$ belongs to block 2's band, and it also falls inside block 3's stretched range, so both blocks practise on inputs at that level. At generation, when block 2 hands its output to block 3 near $\sigma = 0.30$, block 3 has trained on that noise level. The tail blocks get $\alpha \approx 1.52$ (block 4) and 1.61 (block 1). On ImageNet the three-block DiT scores 11.11 without overlap and 10.63 with $\gamma = 0.05$.

The authors also connect the equal split to curriculum learning. Each block gets a task of similar difficulty and an equal share of the data, and they list that as one of two hypotheses for why block-wise training sometimes beats end-to-end training (discussed with the results).

Memory and compute

Training block $b$ needs its own weights, their gradients, the optimizer state for those weights, and the activations of its $L/B$ layers. Appendix G of the paper counts this with per-layer parameter memory $P$ and activation memory $A$. Adam keeps two numbers per weight, so weights, gradients and optimizer state cost $4P$ per layer:

• end-to-end: $(4P + A)\,L$
• activation checkpointing: about $4PL + A$ (activations are recomputed instead of stored)
• DiffusionBlocks: $(4P + A)\,L/B$
• both together: $4PL/B + A$

Gradient checkpointing shrinks only the activation term and costs about a third more compute (a forward pass costs $F$, backward about $2F$, and recomputation adds one more $F$, so $4F$ against $3F$). DiffusionBlocks divides every term by $B$ at no extra compute, provided the inactive blocks' weights and optimizer state are kept off the accelerator, and it can be combined with checkpointing.

At the default setting ($B = 3$, activations twice the weights), checkpointing saves 32% and DiffusionBlocks saves 67%. At $A/P = 16$, where activations dominate, checkpointing alone saves 77%, DiffusionBlocks alone still saves 67%, and the combination saves 90%.

Total training compute does not change. Training all $B$ blocks for $K$ steps each costs $B \times (L/B) \times K = L \times K$ layer evaluations, the same as $K$ end-to-end steps. Measured on one H100 with a 12-layer ViT, an end-to-end step took 0.0507 s and a 4-layer block step took 0.0181 s, so a round of three block steps takes 0.0543 s; the paper attributes the 7% difference to the added noise conditioning.

Image generation gets cheaper. A standard 12-layer DiT runs all 12 layers at each of 50 sampling steps, 600 layer evaluations. With $B = 3$, each step runs one 4-layer block, 200 layer evaluations in total.

Results and limits

The paper compares each network trained block-wise against the same network trained end to end, across five kinds of Transformer.

Image classification. A 12-layer ViT on CIFAR-100 in $B = 3$ blocks reaches 59.30% accuracy against 60.25% end to end, about one point lower, while computing gradients for 4 layers at a time. Forward-Forward on the same ViT reaches 7.85%. On Tiny ImageNet, with $B = 2$, the block-wise ViT scores 36.16% against 35.32%.

Image generation. The paper reports FID against both the training set and the test set. A 12-layer DiT-S/2 on CIFAR-10 scores 30.59 / 37.20 (train / test) block-wise against 32.84 / 39.83 end to end. A 24-layer DiT-L/2 on ImageNet at 256×256 scores 9.00 / 10.63 against 9.01 / 12.09: level on the training split, 1.46 better on the test split. Both use $B = 3$, so training memory falls by a factor of 3. Each FID is computed from 50,000 samples, and the paper reports the best of three evaluations, following EDM.

Masked diffusion language model. A 12-layer model on text8 in $B = 3$ blocks reaches 1.45 bits per character against 1.56 for the baseline, lower being better, with 3× less memory.

Autoregressive language model. A 12-layer Llama-2-style model trains in $B = 4$ blocks of 3 layers. Perplexity is not defined for this diffusion version, so the paper scores generated text with MAUVE (similarity between generated and real text distributions, higher is better) and with the perplexity that Llama-2-7B and GPT2-XL assign to the generated text (lower is better). On LM1B, block-wise training is better on all three: MAUVE 0.71 against 0.50, Llama-2 perplexity 12.32 against 14.58, GPT2-XL perplexity 30.99 against 38.87. On OpenWebText it is comparable: MAUVE 0.82 against 0.85, perplexities 14.99 against 15.05 and 26.33 against 25.24.

Recurrent-depth model. Huginn applies one block of layers repeatedly, 32 iterations on average, starting from noise, and trains with backpropagation through time truncated to 8 iterations. A loop $\mathbf{z}_k = \mathbf{z}_{k-1} + f_\theta(\mathbf{z}_{k-1})$ already has the Euler form, so DiffusionBlocks trains the looped block as a single denoiser: sample a $\sigma$, run one forward pass, and match the target. At inference the model still iterates 32 times. The paper calls this up to a $K$-fold (here 32-fold) cut in compute per training step. It trained for 15 epochs against Huginn's 5 and reports about 10× less total computation, consistent with $32 \times 5 / 15 \approx 10.7$. On LM1B it scores MAUVE 0.70 against 0.49, Llama-2 perplexity 16.08 against 17.04, and GPT2-XL perplexity 42.43 against 46.73.

NoProp. Li, Teh and Pascanu's concurrent NoProp also trains blocks as denoisers, but it comes with its own CNN and was tested only on classification. Applied to NoProp's own architecture on CIFAR-100, DiffusionBlocks reaches 46.88%, against 46.06% for NoProp-DT, 21.31% for NoProp-CT, 37.57% for NoProp-FM, and 47.80% for backprop. Table 6 marks NoProp-DT as block-wise but discrete-time, and NoProp-CT and NoProp-FM as continuous-time but not block-wise, so DiffusionBlocks is the only row with both marks; the paper words this as the only method that successfully combines continuous time with block-wise training.

Figure 10 plots the paper's block-count ablations, which vary $B$ with the total depth fixed and overlap switched off. Look at where each curve sits relative to the dashed end-to-end line.

On ImageNet, $B = 2$ gives FID 9.90 and $B = 3$ gives 11.11, both better than 12.09 at $B = 1$; $B = 4$ (11.90) is about level and $B = 6$ (14.43) is worse, with each block down to 4 of the 24 layers. CIFAR-10 has the same shape, best at $B = 2$ (35.47 against 39.83). For the language model MAUVE peaks at $B = 4$ (0.67), which is why the main experiment uses four blocks. Sampling speed grows linearly with $B$ for the image models, since each step runs one block.

Why a moderate split can beat end-to-end training is an open question, and the paper labels its answer a hypothesis. Each block is tied directly to the target by its own denoising loss, instead of receiving gradient through every layer above it. And giving each block its own noise band, cut so the bands carry equal mass, amounts to a curriculum of balanced difficulty. The paper proves neither. It attributes the decline at large $B$ to reduced capacity per block: at $B = 6$ an ImageNet block has 4 of the 24 layers.

The limits are stated in the paper's conclusion. Every block must map its input to an output of the same shape, which excludes U-Net-style architectures for now. All experiments train from scratch; turning an existing pretrained model into DiffusionBlocks by fine-tuning is future work. The experiments use Euler steps throughout because Euler matches the residual form, and other samplers would need different connections between blocks. The largest model tested is the 24-layer DiT-L/2.