DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation
Authors: Makoto Shing
Masanori Koyama
Takuya Akiba
Venue: ICLR 2026 Poster (2026)
Keywords: block-wise training
backpropagation-free training
memory-efficient training
Primary area: generative models
TL;DR: We introduce DiffusionBlocks, a framework that partitions transformers into independently trainable blocks, reducing memory requirements proportionally while maintaining competitive performance across diverse architectures and tasks.

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.

Published as a conference paper at ICLR 2026
DIFFUSIONBLOCKS: BLOCK-WISE NEURAL NET- WORK TRAINING VIA DIFFUSION INTERPRETATION
Makoto Shing', Masanori Koyama", Takuya Akiba'
Sakana AI, 'The University of Tokyo (mkshing, takiba) ) sakana ai, masanori koy.mauwebaaaaaaaaaaaa www.aap....ut.oooo... ac.
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 DifsusiinnBlokss principled framework for transforming transformer- based networks into genuinely independent trainable blocks that maintain com- petitive 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 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 a time, thereby reducing memory requirements in proportion to the number of blocks. Our experiments on range of transformer architectures (vision, diffusion, autoregressive, recurrent-depth, and masked diffusion) demon- strate that DiffusionBlocks training matches the performance of end-to-end train- ing while enabling scalable block-wise training on practical tasks beyond small- scale classification. DiffusionBlocks provides theoretically grounded approach that successfully scales to modern generative tasks across diverse architectures.
Code is available https://githhb..co
www../oow.aaaaaaaaaaa http://tthhh...coc/aaaaaaaaaaa// ffus on ocks.
INTRODUCTION
The memory bottleneck in neural network training. Modern AI led by generative mod- els (Brown al. 2020 Rombach al., 2022; Touvron al. 2023 Peebles & Xie 2023) has
become integral to everyday life. These models rely on end-to-end backpropagation, which requires storing intermediate activations across network layers during training. This fundamental require- ment causes memory consumption to grow linearly with network depth, creating computational bottlenecks that limit both research flexibility and practical deployment.
Block-wise training: promises and limitations. Block-wise training methods partition partioiipartitoon networks into smaller components that can be trained independently. promising dramatic memory savings. Despite this potential, existing approaches (Hinton 2022 Bengio et al. 2006 Nøkland & Eidnes)
2019 Belilovsky et al.) 2019 Siddiqui et a 20044) consistently underperform end end-to-end training.
The core challenge twofold: (1) lack of theoretical grounding: existing methods rely on ad-hoc lo- cal objectives without principled coordination between blocks, (2) limited applicability, where they require seradigm-speciffe designs, task-specific objectives that do not naturally extend beyond clas- sification. Their results are typically demonstrated only on custom architectures without providing systematic procedures to be applied to modern architectures such as Transformers (Vaswani et al., 2017) (Section 4), leaving their applicability modern generative AI largely unexplored. Without a systematic framework grounded in theory, block- wise training remains an unfulfilled promise.
We use block-wise training to encompass all approaches that partition networks into independently train- able components. This includes layer-wise training as the special case where each block contains one layer.

---

Published as a conference paper at ICLR 2026
Standard Network
Diffusion Models
output
output
DiffusionBBlocks
DiffusionBlocks
Only one block per step
Recurrent- depth Models
Rerative
max
Input
Irpul
-
DiffusioniBlocks
Requires backgrop through all layers
Train each block independently for specific noise rangees
Single pass during training
Figure 1: Overview DiffusionBlocks. Left: Standard networks require backpropagation through all layers. Center: DiffusionBlocks partitions networks into blocks, each trained independently to denoise within assigned noise ranges. Right: Applications. For diffusion models (top), inference requires only the relevant block per denoising step. For recurrent-depth models (bottom), our frame- work replaces iterative training with single-pass training, eliminating the computational overhead of backpropagation through time.
Diffusion models: a mathematical foundation for decomposition. Score-based diffusion mod- els (Song & Ermon, 2019 Song al. (0221b model the data distribution through a continuous- time process that gradually adds noise, then learns to reverse this process by estimating the score function each noise level. Crucially, the denoising step each noise level can be optimized independently from other noise levels. This independence property provides the theoretical foundation that has been missing from block-wise training approaches: it allows partition networks into blocks, each responsible for a specific noise level range, without compromising global coherence.
Our approach: interpreting networks as diffusion processes. We propose DiffusionBlookss framework that enables principled block-wise training by interpreting sequential layer updates in transformer-basee networks as discretized steps of continuous-time diffusion process. Building on the established connection between residual networks and differential equations (Haber & Ruthotto) 2017 Chen et al., (2188,, we leverage the fact that residual connections naturally correspond Euler discretization of the probability flow ODE in diffusion models. This correspondence allows to partition networks with residual connections, particularly transformer- that each handle specific noise-level ranges. These blocks can be trained completely independently, illustrates the core concept of Diffusion- Blocks. Unlike previous block-wise methods with ad-hoe objectives, block's 's training objective from score matching theory. As result, consistent local optimization at each noise level collectively yields faithful approximation of the global reverse process, while also allowing practitioners to seamlessly adopt techniques such as those of Karras el al. 2022 to further enhance training.
networks, into blocks
requiring gradients for only one block at a time. Figure [
framework derives each
Our main contributions are:
Block-wise training via continuous-time diffusion interpretation: We show that transformer-based networks can be interpreted as implementing discretized steps of continuous-time diffusion processes (Section 2.2 enabling genuinely independent block training. Each block learns to denoise within its assigned noise level range, requiring gra- dients for only one block at a time during training (Section 3.T.
biubbiprobabiltty partitioning for balanced learning: We propose a principled, diffusion theoretic strategy that partitions noise levels based on equal cumulative probability mass, ensuring balanced parameter utilization across blocks (Section 3.3
2

---

Published as a conference paper at ICLR 2026
Broad applicability with maintained performance: We conduct extensive experiments (Section 5) demonstrating that DiffusionBlocks successfully applies to diverse architec- tures (vision, diffusion, autoregressive, recurrentddepth, and masked diffusion), achieving competitive performance to end-to-end hackpropagation while requiring gradients for only one block a time. Additionally, framework naturally extends to recurrent-depth mod- els, transforming their multiple-iteration training into single-pass training (Section(5.5)
Significant efficiency gains: During training, only one block requires gradient computa- tion, reducing memory requirements proportionally to the number of blocks. For diffusion models, inference requires only one relevant block per denoising step (Section [.2). For recurrent-depth models, framework eliminates K iterations during training, demon- strating up to K'ofodd reduction in training computation (Section 5.5
PRELIMINARIES
2.1 SCORE-BASED DIFFUSION MODELS
We adopt the Variance Exploding (VE) formulation (Song et al.] 2021b Karras al. (022) where a clean data Pdata is perturbed with Gaussian noise at noise level a: Zo y + ae where (0, I). For generations, noising process:
the deterministic probability flow ODE that reverses the
dz a Pog)) do
where x log Per (Zo) is the score function. Using Tweedie's formula, the score is approximated via (Z( 0) that predicts clean data from noisy input: V, lo (2.) 2 De(z..0)-z. (Rob-
denoiser
bins 1992 Hyvärinen 2005 Vincent 2011). The denoiser is trained by minimizing:
(( : Ezy pr
[a( + oo, y11]]
where w(() weights different noise levels and Phoise is the noise level distribution used during train- ing. The choice of Pnoise determines which noise levels are emphasized during training. Karras et al. (2022) uses log-normal distribution to concentrate training on perceptually important intermediate noise levels where image structure emerges. The weighting w(ar ) is designed to counteract the sam- pling bias from Proises ensuring balanced gradient magnitudes across all noise levels (Karras al., 2022
2.2
RESIDUAL CONNECTIONS AS EULER STEPS OF THE REVERSE DIFFUSION PROCESS
The connection between residual networks and differential equations has been established in prior works (Haber& & Ruthotto 2017 Chen al.] 2018). We extend this perspective to show that residual networks naturally implement discretized steps of the reverse diffusion process. Applying Euler
discretization to Eq. (I) with noise levels ag >> 01
ar, we define Age = O--1 oe
and obtain:
Zo Zo--1
Aot Ji-1 Vx log Pa (Zo.-1)
(3)
Ag:
-az-1
(Zar-1
Do Do 0t-1)
(4)
Te-1
As has historically been utilized in the development of the networks with sequential updates, this update rule has an affinity with skip connections. In fact, modern architectures such as Transform- ers (Vaswani et al. (017)) employ residual connections where each block updates its input through an additive transformation: Z, Z, + for (z,-1) where E Rd denotes the intermediate output of the block L, and fo, is the block transformation parameterized by 0z. This structure appears in ResNets (He et al., 2016). Transformers, and other modern architectures (Peebles & Xie 2023 Tou- vron et al. 2023 DeepSeek-AL et al.) 2025). This scheme is also used in the recent development of recurrent-depth models (Dehghani et al. 2019 Fan et al. 2025| Geiping et al. 2025), which apply the same network parameters recursively K times: 2₁ + (z for E [K]. How- ever, these methods suffer from the expensive backpropagation through time (BPTT), and various measures have been taken to reduce its computational burden, for example, by gradient trunca-
tion (Williams & Zipser 1995 Mikolov et al. 2010 Geiping al. 2025). That being said, the
3

---

Published as a conference paper at ICLR 2026
Partition L layers Into B blocks Define noise partition into B intervals
output
Augment each block with noise conditioning
output
min
output
DiffusionBlocks
input
input
input
noise
Input
output
Figure 2: 3-step conversion of a standard neural network to DiffusionBlocks at training phase. Step Partition L layers into B blocks. Step Define noise distribution Po (e.g., log-normal) and partition the range [mmin mmax] into B intervals {[ 16-1)) assigning each block specific noise range (Section 3.3). Step 3: Augment blocks with noise conditioning: extend input (x, Za) where Z, block trained independently from other blocks to predict target y within its assigned noise range.
and incorporate noise-level conditioning (e.g., via AdaLN), Then, each
Standard Network Training
DiffusionBlocks Training
1: Given: Network with parameters 0
Given: A single block E [B] with parameters O,
2: Sample data (x,y)
2: Sample data (x,y)
3: Zo <- x
3: Sample o ~ v from (os, 06-1]
4: for f to L do
4: t- ox. + oe), where ((, Apply
5:
2 + (2(-1)
block to denoise
6: end for
5: C +- w(o)- Loss(ý.y)
Weighted loss
7: y
+ ZL
6: Update only O, via backprop
8:
L + Loss(ý.y)
9: Update all 9 via backprop
Standard Network Inference
DiffusionBlocks Inference
1: Input:
I: Input: x, noise levels {at)
& Typically, T B
2: Za fx x
2: Za (0 o
3: f= to do
do 4: Select block where as (ab, 06-1)
3: for 0
4:
Z: +- 2:-1 + So, (21-1)
5: end for
5: ý + felln; (x, (x))
6: Output: zz
6:
Euler step 0,011.
Eq. 5
7: end for
8: Output: Zr
Figure 3: Training and inference algorithms for standard residual networks (left) versus Dif- fusionBlocks (right). Given: A L-layer network partitioned into blocks with noise ranges {lo,, 00-1l) noise distribution Po. and training data N.)) The function w(ar de- notes the loss weighting, and fe.j. represents the noised-conditioned block with parameters 0,.
critical observation is that, in the setting of the diffusion introduced in the previous section, Dø itself in Eq. (4) can be trained with Eq. 2 without BPTT, thereby providing a theoretically sound optimization method of dynamical system through ensemble of local optimization. In the next section, we provide recipe for converting networks with skip connections into diffusion, thereby replacing the backpropagation through layers with the optimization scheme analogous to Eq. (

---

Published as a conference paper at ICLR 2026
3 METHOD
3.1 CONVERTING A NEURAL NETWORK TO DIFFUSIONBLOCKS
Our goal in this section is transform given feedforward system into a discretized version of the recursive denoising steps in the diffusion model. Throughout this paper, we denote by (x, y) the input-output pairs where x represents the network input (e.g., images for classification) and y is the target output (e.g., class label for classification). Figure [ provides an overview: instead of backpropagating through all layers, we partition networks into blocks that independently learn to denoise within assigned noise level ranges. Consider neural network in form of a stack of set- to-set maps (c.g. transformer- based networks) F = {fo. & [L]} with the same output and input dimensions, so that fo, maps a variable set of tokens IR to the same number of tokens in Ted The original network therefore processes the input with for o .... fa,, followed possibly by a readout module. Or, in more conventional formulation with the presence of residual, the original network may update the f-th layer input Zy the next layer via the rule - Z, (Zz). We transform this network into a stack of Diffusion Blocks through the following three steps (Figure
Step 1: Block partitioning. We partition into B blocks F Wb where Fs contains layers indexed by (6-1 + 1, ...................... (.). Let fo. == for." .... for be the composition of layers in Fo.
Step 2: Noise range assignment. We define noise distribution Pooise and define noise range [min maax] We partition the range into B intervals {[o,, 06-1]/=- b=1 We recommend the choice of log-normal for Puoise. following Karras et al. (2022). along with the partitioning strategy in Sec- tion 3.
Step 3: Augmenting blocks with noise conditioning. Finally, we suit {fe, Ja to the update rule in Eq. (4) by letting fe, play the role of Do,- Leveraging the assumption that fe, is map from set of tokens to a set of tokens, we alter the input fe. from x to x = (x, z). Additionally, we ex- tend each block fo. to incorporate noise-level conditioning through, for example, via normalization (AdaL.)) (Peebles & Xie 2023). We denote this conse-conditioned version as Se Altogether, the update of the diffusion block constructed from F is given by:
Aob
Zb = 7b-1 +
(zs-1 [Fosioo (x, 25-
(5)
0b-1
where [f()]2 is the set of tokens corresponding to z (i.e. f(.) = (lf(-)xx (f)l))) put, our modified update rule Eq. (5) can be rewritten as 26- aZ ffo, (x, 2b-1) where and B3 are constants dependent on a ratio. We note that our modification of the network into the stack of diffusion blocks maintains most of the structure of the original, particularly in the presence of skip connection, so that z; z + fa, (1)-1) is the original update rule. At the time of inference, 2o serves as the intermediate estimator of the target variable, with Zg maxe being the pure noise.
(..()...((()))))). More abstractly
Please see Figure 6l in Appendix for the conversion of this inference process.
3.2 BLOCK-INDEPENE
TRAINING OF THE DIFFUSION BLOCKS
By the network modification recipe in the previous section, we transform the original feedforward map to the recursive denoising map in a diffusion process. The advantage of this modification is the fact that the objective in Eq. (2) can be optimized at any noise level a independently without knowledge of other noise levels. This allows us define a training objective for each block b:
(O,) := E (x.))
(w) Loss(fo. (x, y + oe).
(6)
where (b) is the noise distribution Pnoise with the support of [o%, 06-1] and renormalized, and Loss(-; :) is the inner loss function, typically L2 loss as in Eq. (2). Each block independently learns to denoise within its assigned range, with training samples drawn according to the original distribution
Pnoise. Collectively, the B blocks cover the entire noise distribution: [ σ6-1] = [amin- mmax], ensuring that the complete network can denoise at any noise level while each block specializes in its designated range. This independence enables training with memory requirements for only T./ B layers, storing activations only for the active block, compared to all I. layers required by
5

---

Published as a conference paper at ICLR 2026
standard training. More succinctly in comparison to the original network, gain this block-wise independence from the fact that f0,1=(X, y σε) is now modified to predict y for each b. This way, training for each block can be carried out without waiting to receive the output of the previous layer. Please see Figure 5 in Appendix [ for the training process in the specific adaptations for different architectures. Figure 3 provides an algorithmic procedure of training and inference. This approach achieves a B memory reduction during training, as gradients are computed for only one block at time.
3.3 BLOCK PARTITIONING STRATEGY
A critical design choice in DiffusionBlocks is how to partition the noise level range [Amin: m into intervals. A naive approach would divide the range uniformly: σb miin + (amax min) However, this fails to account for the varying difficulty of denoising at dif- ferent noise levels. Following Karras et al. (2022), we adopt a log-normal distribution for sam- pling noise levels during training: log a ([Pmeam Paa)). This distribution concentrates proba- bility mass at intermediate noise levels, which empirically contribute most to generation quality. To preserve this distribution across the entire net- work while ensuring each block handles equal de- noising difficully, we partition based on cumu- lative probability mass. Specifically, we choose boundaries {a such that each block han- dles exactly 1/B of the total probability mass: I Pnoise ( = 1/B. The block boundaries are computed as 0b = exp(Pmean + Pad -- (gh)), where Φ- is the inverse standard normal CDF and
EOM POF (iggtomal) al] Uniform Boundaries Equi-hob Bounderie
0.30
0:25
Density
0:20
0:15
0.10
0.05
& Qmin + B (Amax qmin), with gmin max
0.00
q6
T0
10"
10a 12'
10'
30
Nolse level a
log minn max Pawnn
This equi-probability par-
Pas
ritioning ensures that each block handles an equal amount of the training distribution's 's probability mass, leading to balanced parameter utilization. As shown in Figure 4 blocks assigned to intermedi- ate noise levels, where denoising is most challeng- ing, receive narrower intervals, while blocks han- dling very high or low noise levels receive wider intervals. This strategy optimizes learning effi- ciency across all blocks. In Section (.6 demonstrate that this strategy contributes significantly to the training of DiffusionBlocks. Also, see Appendix C for implementation details.
Figure 4: Equi-probability partitioning (B 3). Blocks partition the log-normal ps by equal probability mass (orange boundaries), not uni- form spacing (gray), concentrating capacity where denoising is most challenging.
4 RELATED WORKS
Block-wise training methods. Various block-wise training approaches (Hinton 2022 Bengio et al. 2006; Nokland & Eidnes 2019: Belilovsky al. 2019; (iddiqui et al. 2024) partition net- works into independently trainable components but lack theoretical grounding, relying on heuris- tic objectives that fail to guarantee global performance when optimized locally. Approaches like Nowardd-Forarr algorithm (Hinton, 2022) rely on contrastive objectives, which fundamentally limit them classification tasks and make adaptation to generation non-trivial. In contrast, Dif- fusionBlocks leverages denoising score matching theory, which naturally decomposes into indepen- dent local objectives without task-specific constructs, enabling application both classification and generative tasks.
Comparison with NoProp. Concurrently with our submission, Li et al. (2025) has also released backpropagation-free strategy in close relation to our philosophy. However, they present their technique together with the custom CNN-based architecture one package and evaluate only on classification tasks, making it unclear how to apply their approach to modern architectures or tasks other than the classification they showcase in their work. In contrast, DiffusionBlocks provides a systematic procedure for converting any residual networks, particularly modern transformers, into block-wise trainable models with minimal modifications. We partition the continuous noise range using equi-probability partitioning and demonstrate success on both generative tasks and classifica-

---

Published as a conference paper at ICLR 2026
DIT results for image gener- CIFAR-100. DiffusionBlocks ation. FID is computed both training achieves comparable accuracy while training only 4 layers
Table 1: VIT results on Table 2:
Table Masked diffu-
sion model (MDM) results
and test splits (train test), DiffusionBlocks on text8. DiffusionBlocks achieves comparable scores while reducing training memory and inference cost by 3x,
improves BPC while training
a time, outperforming Forward-
with 3x less memory through
Forward algorithm.
masking schedule partitioning.
Dataset
Method
FID (1)
Method
Accuracy (f)
DIT
32.84/ 39.83
Method
BPC (4)
CIFAR-10
VIT
60.25
DiffusionBlocks
30.59/37.22
MDM DiflusionBlocks
1.56
DiflusionBlocks
Forward-Forward
59,30
Ietge
DiT
DiffusionBlocks
9.00 10.63
9.01 / 122.09
1.45
tion tasks. In Section 5.6.1 we apply DiffusionBlocks to their architecture, and demonstrate that our continuous-time block-wise training with equi-probability partitioning is more effective.
Stage-specific diffusion models. Several works train specialized models for different noise levels in diffusion (Balaji et al., 2023 Fang et al., 2024 Park et al.) 2024, Reuss et al. 2025). However, these approaches train models jointly fine-tune from shared parameters, DiflusionBBocks trains blocks independently, with no shared parameters or joint fine-tuning, achieving complete isolation.
5 EXPERIMENTAL RESULTS
We evaluate DiffusionBlocks across diverse architectures and tasks to demonstrate its generality and effectiveness. Detailed experimental configurations are provided in Appendix For each architec- ture. we report task performance alongside the memory reduction factor B, where only L/ B/B layers require gradients during training.
Baselines. Because DiffusionBlocks is framework for transforming networks into block-wiss wise trainable models, we evaluate its efficacy by comparing the modified network (trained block-wise) against the original network (trained with end-to-end backpropagation). Other block-wise training
methods in practice today also include norward-Forwad (FF) (Hinton) 2022) and the concurrent NoProp (Li et al., 2025). Fair comparison against these methods warrants careful experimental design. Firstly, we compare against FF only on classification tasks (Section 5.T) since its contrastive objective does not naturally extend to generation. Also, because NoProp is proposed together with custom architectural design rather than with a principled transformation procedure to be applied to a vanilla network, the adaptation of NoProp to other architectures involves nontrivial design choices and freedom. To enable fair comparison with NoProp, we therefore use their specific architecture as the base diffusion model on which to apply our DiffusionBlocks (Section 5.6.
5.1 VISION TRANSFOR MERS FOR IMAGE CLASSIFICATION
We first validate DiffusionBlocks on classification tasks using Vision Transformer (ViT) (Dosovit- skiy et al. 222T)) on CIFAR- 100 (Krizhevsky (200).. A 12-layer ViT is partitioned into B=3 blocks, with noise added to class label embeddings during training. We compare against the Forward- Forward algorithm, representative block-wise training method that uses contrastive objectives. Table I maintains baseline accuracy while requiring gradients for only 4 layers. Notably, Forward- Forward achieves only 7.85% accuracy, highlighting the importance of principled denoising objec- tives over ad-hoc contrastive approaches.
5.2 DIFFUSION MODELS FOR IMAGE GENERATION
Having established its effectiveness on classification tasks, we now turn to generative models. We begin with image generation, where DiffusionBlocks provides both training and inference efficiency benefits. We apply DiffusionBlocks to DiT (Peebles & Xie 2023) within the EDM Karras et al. 2022) framework. We evaluate 12 12-layer DiT (DiT-S/2) on CIFAR-10 (Krizhevsky 2009) and 24. layer DiT (DiT-L/2) on ImageNet at 256 256 resolution (Deng et al. 2009), both with B=3 blocks. During inference, we use Fuler sampling with 50 steps and classifier-free guidance (scale 2.0) (Ho
7

---

Published as a conference paper at ICLR 2026
Table 4: Autoregressive (AR) transformer results generation. DiffusionBlocks maintains genera- tion quality with 4x memory reduction on both LMIB and Openwebtexi (OWT) datasets.
Dataset
Method
MAUVE (1)
PPL (Llama-22 (L)
PPL (GPT2 -XL) (L)
AR
0.50
14.58
38.87
LMIB
DiffusionBlocks
0.71
12.32
30.99
AR
0.85
15.05
25.24
OWT
DiffusionBlocks
0.82
14.99
26.33
Table 5: Recurrent-depth model results for text generation. DiffusionBlocks eliminates 32 training itera-
tions, achieving better performance with single-pass training.
Method
MAUVE (t)
PPL (Llama (L)
PPI. (GPT2-XL) (1)
Huginn (Geiping et al. 2025
46.73
DiflusionBlocks
0.70
16.08
42.43
& Salimans| 2021). Table 2 shows that DiffusiioBBlocs achieves comparable FID scores with 3
memory reduction. Additionally, inference requires only one block per denoising step, providing computational savings proportional to the number of steps.
5.3 MASKED DIFFUSION MODELS FOR TEXT GENERATION
We extend DiffusionBlocks to masked diffusion language models using MD4 (Shi et al. 2024) on the text8 dataset (Mahoney] 2011). While continuous diffusion models naturally map to our framework through noise levels a, extending DiffusionBlocks to discrete masked diffusion requires careful adaptation. Specifically, we partition the masking schedule rather than continuous noise levels, ensuring each block handles equal share of the demasking work (details in Appendix D We use a 12-layer DiT-based transformer (Lou et al.] 2024; Sahoo et al.) (022)) partitioned into B=3 blocks. Table shows that BiffusionBlocks achieves 1.45 bits-per-character (BPC) compared to MD4's 1.56, while using 3x less memory. This improvement confirms that our principled noise- level partitioning effectively extends discrete diffusion processes.
5.4 AUTOREGRESSIVE MODELS FOR TEXT GENERATION
We demonstrate that DiffusionBlocks successfully transforms standard autoregressive (AR) models, which are architectures originally designed for next- text-token prediction, not denoising. Using 12-layer Llama-2-style transformers (Touvron et al., 2023) with B=4 blocks, we evaluate Billion Words
Dataset (LM1B) B) (Chelba et al. al.] 2014) and WebTett (OWT) (Gokaslan & Cohen 2019). While AR models are typically evaluated using perplexity, computing traditional perplexity non-trivial for our diffusion framework as it is not derived from ELBO. Instead, we evaluate using MAUVE (Pil- lutta et al. 2021 scores following SEDD (Lou et al. 2024 to measure similarity between generated and real text. We also report generative perplexity from two teacher models, L. amaam2222 - and GPT2-XX XL (Radford et al. 2019)). following Lou et al. (2024); Sahoo et al. (2024). Table shows that DiffusionBlocks achieves comparable performance despite training only layers at time, demonstrating the framework's broad applicability beyond diffusion-natio architectures.
5.5 RECURRENT-DEPTE MODELS FOR TEXT GENERATION
We now showcase different application DiffusionBlocks beyond block-wise training. noted in Section 2.2 the updates in recurrent-depth models naturally correspond to diffusion steps. Fol- lowing Section 3.1 we apply DiffusionBlocks to Huginn (Geiping et al. 2025), a recurrent-dppth model that applies the same network multiple times, starting from noise. While Hugi nn uses 8-step truncated BPTT to avoid the full BPTT over 32 iterations, DiffusionBlocks makes this optimization more efficient. because it only requires single forward pass per training step. Table 5 shows better performance on LMIB for text generation while eliminating 32 iterations. This demonstrates that our framework enables fundamental training transformations beyond block-wise training.
8

---

Published as a conference paper at ICLR 2026
Table 6: Comparison with NoProp on CIEAR-100. DiffusionBlocks achieves both continuous time formu-
lation and layer-wise training. All scores except DiffusionBlocks are taken from NoProp (Li et al. 2025). Note that Backprop in this table the result applying BPTT the sampled paths a specific form of SDE that equationally resembles the process used in NoProp-Dre See Li et al. (2025) for the details.
Method
Continuous
Block-wise
Accuracy (f)
Backprop
NoProp-DT
46.06
rop-rT
21.31
NoProp-PM
37.57
(Ours) DiflusionBlocks
46.88
Table 7: Effect of block partitioning strategy on CIFAR-10. Layer distribution indicates the number of layers in each of the blocks (total-
Table 8: Effect of block count ImageNet. Fewer blocks achieve better FID but require more layers per dif- fusion step, creating a trade-off between quality and effi- ciency. The scores that surpass the end-to-end backprop- agation (B=1) are highlighted bold.
ing 12 layers).
Partitioning Strategy
Layer Distribution
FID (1)
Uniform
[4.4.4]
43.53
Number Blocks
(L)
L/B (U)
Relative Speed
Uniform
[3.6,3]
43,59
Uniform
[6,4.2]
47.49
B 1
12.09
24
1.0x
Uniform
[2.4.6]
42.37
B 2
9,90
12
2.0x
Equ-Probabilityy
[4,44]]
38.03
B 3
11.11
3,0x
Equi-Probabliity
Pruu-Probbblitt
[6.4.2]
[3.6,3]
41.64
45.42
B
11.90
6
4,0x
Equi- Probability
[2.4.6]
40.40
B
14.43
6.0x
5.6
ANALYSIS
5.6.1
COMPARISON WITH NOPROP
We compare DiffusionBlocks with NoProp ablation study, applying their custom CNN- based architecture to isolate the effect of our continuous- block-wise training using equi- probability partitioning. Table shows results CIFAR-100 classification (details in Ap- pendix T.. BiffusionBlocks outperforms all NoProp variants. Notably, while maintaining com- parable performance the backpropagation, DiffusionBlocks is the only method that successfully combines continuous- time formulation with block-wise training. This demonstrates that our equi- probability partitioning with independent denoisers per block is crucial for continuous-time block- wise training.
5.6.2 ABLATION STUDIES ON DESIGN CHOICES
We conduct ablation studies to analyze key design choices in DiffusionBlocks. All experiments follow the configurations described in Section 5.2
Block partitioning strategy. Table compares equi-probability partitioning with uniform partitioning on CIFAR-10. Equi-probabiliic partitioning achieves significantly better FID across all layer distributions. The improvement stems from allocating computational resources based on denoising difficulty: equi-probabilii assigns more blocks to challenging intermediate noise levels where most learning occurs, while uniform partitioning wastes capacity trivial very high/low noise regions. Notably, within equi-probability partitioning, uniform layer distribution (4-4-4) achieves the best FID, demonstrating that practitioners can simply divide layers equally without tuning since the noise-based based partitioning automatically balances learning difficulty across blocks.
Number of blocks B. [able|Bsummarrzer the effect of varying the number of blocks on ImageNet (see Appendix F22 the results on CIFAR-10). It reveals the trade-off between generation quality and efficiency on ImageNet. Notably, moderate block counts (B=2 B=3) achieve better FID than
These ablations disable block overlap in Appendix isolate the effectiveness of each component, sulting in the FID difference from Table

---

Published as a conference paper at ICLR 2026
end-to-end training (B=1), suggesting that moderate block partitioning can actually improve per- formance through specialization. As B increases further, quality gradually declines due to reduced capacity per block, though inference speed improves linearly. The optimal B varies across tasks (see Appendix F.3 for language modeling results).
6 CONCLUSION
We introduced DiffusionBlocks, a theoretically grounded framework that transforms residual net- works into independently trainable blocks through continuous-ting diffusion interpretation. By rec- ognizing that residual connections naturally implement discretized diffusion steps, we provide systematic recipe requiring minimal modifications that maintains competitive performance across diverse architectures while achieving B x memory reduction during training.
Future works. Our work opens several important directions for future research. First, while we consistently used Euler discretization match residual connections, other diffusion samplers (Song) ctt al. al.
20 Lu et al. al.] [023 Zhao et al. 2023) could be employed within blocks with modified inter- block connections. Second, Diftusion locks currently requires matching input-output dimensions, which limits its application to architectures like U-Net (Ronneberger et al., 2015). Third, while we demonstrate DiffusionBlocks' effectiveness models trained from seratch, scaling to even larger models would further demonstrate its practical impact. Particularly, a promising direction to convert pre-trained large models to DiffusionBlocks through fine-tuning rather than training from scratch. Fourth, determining the optimal granularity of block partitioning presents an interesting theoretical and practical challenge. While our experiments demonstrate that treating entire architec- tural blocks (e.g., complete ViT blocks) as single denoising units works well, a principled method for selecting the ideal partitioning granularity based on architecture and task characteristics could further enhance the framework's 's applicability. Finally, understanding why moderate block partition- ing sometimes outperforms end-to-end training warrants theoretical investigation. We hypothesize two contributing factors: (1) DiffusionBlocks employs different optimization structure in which each block is directly linked to the target through a denoising objective in Eq.( creating a learn- ing signal that differs from standard end-to-end training; and (2) assigning different noise ranges to different blocks may induce beneficial specialization effects. Combined with equi-probability partitioning, this introduces natural form of curriculum learning (Bengio et al., (009) by allo- cating balanced difficulty across blocks. Developing a formal theory and analysis for these effects could reveal new principles for scalable and structured neural network optimization beyond memory efficiency.
DiffusionBlocck represents step toward democratizing large-scale model training by reducing com- putational requirements without sacrificing performance, making advanced AI capabilities more ac- cessible.
AUTHOR CONTRIBUTIONS
Makoto Shing conceptualized the DiffusionBlocks framework, developed its diffusion--heorreiic for- mulation connecting residual networks and continuous-time diffusion processes, implemented the method, conducted all experiments, and wrote the manuscript. Masanori Koyama provided theoret- ical insights into the diffusion-based interpretation and contributed to refining both the manuscript and the conceptual positioning of the work. Takuya Akiba supervised the research and provided technical guidance and feedback throughout the project. All authors contributed to the interpreta- tion of results and manuscript revision.
ACKNOWLEDGEMENT
The authors would like to thank Stefano Peluchetti for helpful feedback on an earlier version of the draft.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Hossein Aghagolzadeh and Mehdi Ezoji. Contrastive forward-forward: A training algorithm of vision transformer. Neural Networks, 192:107867, 2025.
Marianne Arriola, Subham Sekhar Sahoo, Aaron Gokaslan, Zhihan Yang. Zhixuan Qi, Jiaqi Han, Justin I Chiu, and Volodymyr Kuleshov. Block diffusion: Interpolating between autoregressive and diffusion language models. International Conference on Learning Representations, 2025.
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 17981- Curran Associates, Inc., 2021.
Yogesh Balaji, Seungjun Nah, Xun Huang, Arash Vahdat, Jiaming Song, Qinsheng Zhang, Karsten Kreis, Miika Aittala, Timo Aila, Samuli Laine, al. ediff-i: Text-to-image diffusion models with an ensemble of expert denoisers. arXiv preprint arXiv:2211.01324 2023.
Eugene Belilovsky, Michael Eickenberg, and Edouard Oyallon. Greedy layerwise learning can scale to ImageNet. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 583-593. PMLR, 09-15 Jun 2019.
Yoshua Bengio, Pascal Lamblin, Dan Popovici, and Hugo Larochelle. Greedy layer-wise training of deep networks. In B. Schölkopf, J. Platt, and T. Hoffman (eds.), Advances in Neural Information Processing Systems, volume 19. MIT Press, 2006.
Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In In- ternational Conference on Machine Learning, pp. 41-48. Association for Computing Machinery, 2009.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, Usvsn Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In An- dreas Krause, Emma Brunskill, Kyunghyun Cho. Barbara Engelhardt. Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, pp. 1877-1901. Curran Associates, Inc., 2020.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Kochn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv: 13 12.30005 2014.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Desa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
Ekin Dogus Cubuk, Barret Zoph. Jon Shlens, and Quoc Le. Randaugment: Practical automated data augmentation with reduced search space. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp- 18613-18624. Curran Associates, Inc., 2020.
DeepSeek-AI, Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, et al. DeepSeek-R1: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.1294 2025.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser, Universal transformers. In International Conference Learning Representations, 2019.
11

---

Published as a conference paper at ICLR 2026
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hier- archical image database. 2009 IEEE Conference on Computer Vision and Pattern Recognition, pp. 248-255, 2009.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. In M. Ran-
zato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 8780-8794. Curran Associates, Inc., 2021.
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H. Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Continuous diffu- sion for categorical data. arXiv preprint arXiv:2211. 5089, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Ying Fan, Yilun Du, Kannan Ramchandran, and Kangwook Lee.
Looped transformers for length generalization. In The Thirteenth International Conference on Learning Representations, 2025.
Gongfan Fang, Xinyin Ma, and Xinchao Wang. Remix-dit: Mixing diffusion transformers for multi- expert denoising. A. Globerson, L.. Mackey, D. Belgrave, A. Fan, Paquet, Tomezak, and C. Zhang (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 107494- 107512. Curran Associates, Inc., 2024.
Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R. Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up test-time compute with latent reasoning: A recurrent depth approach. arXiv preprint 20
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus.
io/ io/
OpenWebbextCCrppu 2019.
Ishaan Gulrajani and Tatsunori B Hashimoto. Likelihood diffusion language models. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 16693-16715. Curran Associates, Inc., 2023.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse Problems, 34(1) 34(1):014004, dec 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In IEEE Conference on Computer Vision and Pattern Recognition, pp. 770-778, 2016.
Geoffrey Hinton. The Forward-Forwart algorithm: Some preliminary investigations. arXiv preprint arXiv:2212. 133455 2022.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurlPS 202/ Workshop Deep Generative Models and Downstream Applications, 2021.
Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(24):695-709) 2005.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances Neural Information Processing Systems, volume 35, pp. Curran Associates, Inc., 2022.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge, ceag 2015. URL http://vision.
proppppooeeo..... pdf
Qinyu Li, Yee Whye Teh, and Razvan Pascanu. NoProp: Training neural networks without back- propagation nowaard-rropgattin In 4th Conference Lifelong Learning Agents, 2025.
12

---

Published as a conference paper at ICLR 2026
Xiang Li. John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion- Im improves controllable text generation. In Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp- 4328-4343. Curran Associates, Inc., 2022.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310. 1683, 2024.
Justin Lovelace, Varsha Kishore, Chao Wan, Eliot Shekhtman, and Kilian Q Weinberger. Latent diffusion for language generation. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 56998- 57025. Curran Associates, Inc., 2023.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint urXiv:2211.01095. 2023.
Matt Mahoney. About the test data, 2011. URL Bttp:://mattnatoney.nee// textdata html
Tomáš Mikolov, Martin Karafiát, Lukáš Burget, Jan Černocky, and Sanjeev Khudanpur. Recurrent neural network based language model. In Interspeech 2010, pp. (045-1048, 2010.
Arild Nøkland and Lars Hiller Eidnes. Training neural networks with local error signals. In Kama- lika Chaudhuri and Ruslan Salakhutdinov (eds.), International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 4839-4850. PMLR, 09-15 Jun 2019.
Byeongjun Park, Hyojun Go, Jin- Young Kim, Sangmin Woo, Seokil Ham, and Changick Kim. Switch diffusion transformer: Synergizing denoising tasks with sparse mixture-of-expers. arXiv preprint arXiv:2403.09176, 2024.
Gaurav Parmar, Richard Zhang, and Yan Zhu. On aliased resizing and surprising subtleties in gan evaluation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11410-11420, June 2022.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 4195-4205, 2023.
Krishna Pillutla, Swabha Swayamdipta, Rowan Zellers, John Thickstun, Sean Welleck, Yejin Choi, and Zaid Harchaoui. Mauve: Measuring the gap between neural text and human text using diver- gence frontiers. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 4816-4828. Curran Associates, Inc., 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Moritz Reuss, Jyothish Pari, Pulkit Agrawal, and Rudolf Lioutikov. Efficient diffusion transformer policies with mixture of expert denoisers for multitask learning. In The Thirteenth International Conference on Learning Representations, 2025.
Herbert E. Robbins. An Empirical Bayes Approach to Statistics, pp. 388-394. Springer New York, 1992.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition (CVPR), pp. 10684-10695, 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. In Nassir Navab. Joachim Homegger, William M. Wells, and Alejan dro F. Frangi (eds.). Medical Image Computing and Computer-Assisted Intervention MICCAI 2015, pp. 241. Springer International Publishing, 2015.
13

---

Published as a conference paper at ICLR 2026
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang (eds.), Advances in Neural Information Processing Systems, volume 37, pp- 130136-13016 Curran Associates, Inc., 2024.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. Simplified and gener- alized masked diffusion for discrete data. In A. Globerson, L. Mackey, D.. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 103131-103161 Curran Associates, Inc., 2024.
Shoaib Siddiqui. David Krueger, Yann LeCun, and Stephane Deny. Blockwise self-supervised learn- ing at scale. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- tional Conference on Learning Representations, 2021a.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In Interna- fional Conference on Learning Representations, 2021b.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307. 092888, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Compu- tation, 23(7):1661-1674
2011. 2011.
Ronald Williams and David Zipser. Gradient-based learning algorithms for recurrent networks and their computational complexity, pp. 433 486. L. Erlbaum Associates Inc., 1995.
Wenliang Zhao, Lujia Bai, Yongming Rao, Jie Zhou, and Jiwen Lu. Unipe: A unified predictor- corrector framework for fast sampling of diffusion models. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp.
Curran Associates, Inc., 2023.
A NOTATIONS
In this section, we provide the notations that we will be using in the ensuing mathematical formula- tions and statements.
B EXTENSION TO DIVERSE ARCHITECTURES
While we have described DiffusionBlocks for standard residual networks where inputs and outputs naturally live in the same d-dimensional space, the framework extends to specialized architectures. Figures F iilluusttee how different model types can be converted to DiffusionBlocks for training and inference, respectively.
For Vision Transformers (ViT) (Dosovitskiy et al. 2021 in classification tasks (top left), adapt DiffusionBlocks by adding noise to the class label embeddings while maintaining the standard ViT architecture. Specifically, create the input sequence by concatenating the [CISS] token, patch
14

---

Published as a conference paper at ICLR 2026
Notation
Description
T ~ X
Conditioning/Inuut to the network (task-dependent: see below)
y
Clean target data (task-dependent: see below)
a E IR
Noise level in continuous diffusion.
2a E IRd
Noisy data at noise level σ: 2 + de, where c ~ cNN(0)1
2f Rdd
Intermediate activation layer/block
Do Rd y
Denoiser network with parameters 0
fe. Rd + R2d
Layer/block transformation with parameters Oz Number blocks
L
Total number of layers
Examples of (z,y) y) on a task:
Image classification 2: imput image, y/ class label
Image generation
2: noisy image (optionally, and class label), yr clean image
Text Generation (AR)
previous tokens, next token
Text Generation (AR)
x: sequence with mask tokens, y: unmasked sequence
Diffusion
00cl
DiffusionBlecks
DiffusionBlocks
nobe
Wbel
rase
Autoregressive Models
Recurrent- depth Models
Wods
output
World
DiffusionBlocka
Enroeddinn
DiflusionBlocka
Input
Emboddcing
Hello
Warld
Input
<bos>
Converting different architectures to DiffusionBlocks: Training. During training, noise is added target outputs (labels, embeddings, or images) and each block learns denoise within its assigned noise range. Blocks are sampled randomly and trained independently, requiring gradients for only one block time.
Figure 5:
embeddings x, and the noisy label embedding Zo, where Z, Yemb + ae and Yemb E R is the learnalle continuous embeddings for the class label y. Each block b learns to denoise this label representation conditioned on the patch embeddings x. The training loss is the standard cross- entropy between the classification head's output logits (applied to the [CLS] token) and the true class labels, following the conventional ViT training procedure.
For diffusion models (top right), DiffusionBlocks provides a natural fit: these models already op- erate by denoising, so partitioning simply assigns different noise ranges to different blocks without architectural modifications. The standard denoiser Dz o) becomes De, (z. σ)) for block 6.
For discrete output spaces like language modeling (bottom left), we operate in the embedding space following prior works (Dieleman al. 2022 Li al. 2022 Guirajani & Hashimoto 2023 Lovelace et al. 2023). Noise added after the embedding layer: given input tokens x, we compute de. For autoregressive models, the denoiser De, (Zi Z recovers the clean embedding of token from its noisy version, conditioned on previous clean token embeddings Zeci We minimize cross-entropy loss instead of 1.2 loss.
(x), then add noise z,
15

---

Published as a conference paper at ICLR 2026
Vision Transformer
Diffusion Models
Label Embeodding
DiffusionBBockk
DiffusionBlooks
TCLSS
Autoregressive Models
Recurrent-depth Models
outppt
oupput
Enbeediing
Inpult
Hefo
loffusionBBocks
DiffusionBlokks
Embeddingg
Emboddiing
Irput Kyr
Hello
Figure 6: Converting different architectures DiffusionBlocks: Inference. During inference, blocks are applied sequentially from max to min- The figure shows the first denoising step where block b = transforms pure noise Zg into the next state Zi. Only relevant block is active each noise level, providing memory efficiency. @ denotes the Euler step in Eq. 55)
For recurrent-depth architectures that apply the same network K times (bottom right), we interpret the entire recurrence diffusion process. Instead of training with K forward passes through recur- rent iterations, we train the network as a denoiser Do(z,, a) by sampling o Pa and performing a single forward pass to map noisy input to clean output, reducing computational cost by factor K while maintaining the original K Kiiteration inference procedure.
Beyond these adaptations, DiffusionBBlokss also applies to diffusion language models (Austin al,, 2021 Lou al. 2024; Sahoo et al. 2024 Shi al. 2024). where the framework provides additional benefits for text generation. We provide a detailed treatment of this application in Appendix() These diverse applications demonstrate that DiffusionBlocks provides general recipe for transforming various architectures into memory-efficient, independently trainable components.
C IMPLEMENTATION DETAILS IN DIFFUSIONBI.OCKS
We introduce several practical considerations for effective training and inference.
Overlap between blocks. To smooth transitions across block boundaries, we slightly extend each block's noise interval in log-or space. For block responsible for [ab, 0b-1] with 06-1 Th. define Ogg = 1/o,),, where 7 2 0, and train over the expanded range [aa/ an, 01607- 1]. Here controls the degree of overlap: 7=0 recovers non-overlapping intervals, while
yields smoother
transitions between blocks. In practice, we found e [0.0,0. 0.1] effective, and we use 0.05 by default and 0.1 for text generation.
Weighting and preconditioning. Following the EDM framework (Karras al. 2022).
the weighting function: w(r) = (o2 x2 (a )2ma)² where ddata = 0.5 for all experiments. The weighting crucial for equi-probability partitioning to work effectively, sampling bias introduced by the log-normal distribution Po. We also adopt EDM's preconditioning scheme, which involves input scaling ensure stable training dynamics across all noise levels. See
it counteracts the
Karras et al. (2022) for more details.
Normalizing embeddings. For tasks where the target variables are discrete (e.g. class labels in image classification or token ids in text generation), DiffusionBlocks operates the diffusion process
16

---

Published as a conference paper at ICLR 2026
in the continuous embedding space (see Appendix B). A known issue in continuous relaxation of discrete variables embedding collapse, where all learned embeddings correspond to the same vector (Dieleman al. 2022). To prevent this, follow the regularization strategy introduced Dieleman et al. (2022) and apply 1.2 normalization to the embeddings.
Training and inference details. For training efficiency, blocks randomly sampled per itera- tion, requiring memory for only L/ layers. Blocks can alternatively be trained in parallel across multiple GPUs when available. During inference, blocks from max to min- While we use Euler steps in our experiments due to the natural corre- spondence between residual connections and Euler discretization (Section(.22) our framework not limited to this choice. By modifying the inter-block connections to match the discretization scheme of other solvers, any diffusion sampling methods (Song al. 2021a Lu et al. 2023 Zhao et al, 2023) can be employed. We leave this exploration for future work.
generate samples by sequentially applying
D MASKED DIFFUSION LANGUAGE MODELS AS DIFFUSIONBLOCKS
D.1 CONTINUOUS- TIME FORMULATION
We first recall the continuous-time formulation of masked diffusion language models (Sahoo et al.) 2024 Shi et al. 2024). Let X = (xo1..... Fon) denote a sequence of tokens and let
(tt) [J + [] denote the masking schedule at continuous time é E [0, 1], where a(t) rep- resents the probability of remaining unmasked. The forward process progressively masks tokens as:
(Xx
xo)
( 1=1
II
zox) where
Foi
with prob. ar({), with prob. 1 a(t).
[MASK],
The training objective in continuous form is:
((0) Exo
=Exo [/' aa((t) E Ex (x)
Σ
CB(fo(X,, t) sun) dt,
(8)
1 a(t)
Liru [NASK]
where a'(t) = da//t < 0 and CE denotes cross-entropy loss. This form is equivalent to the continuous-timg NELBO (Shi et al., 2024; [ahoo et al., 2024), but expressed with a nonnegative
weight multiplying CE, which avoids sign ambiguity.
D.2 PARTITIONING INTO DIFFUSIONBLOCKS
To enable block-wise training, we partition the objective in Eq. 8 into B disjoint intervals t. The expected number of masked positions at time is n11=((()), so the effective density of contributions
(t) (1-aih a()) -a(()). a-
Hence, the contribution of interval [lo, (n) is
-'(1) a(la) ()) altx)-
(10)
This shows that the training mass is distributed uniformly in a, not in Therefore, the natural partition boundaries are defined by equal decrements of o:
B = O....... O...,, B,
(11)
with corresponding time boundaries obtained by inversion:
t4 = a (1 h)
(12)
For a linear schedule att) = 1 - t, this simply yields t, = b/B. Each block b is then trained independently on its assigned interval:
(.().) Exa 6' -a'(1) h--1 a(t)
Ex,, (xx)xo)))
Σ
CE CE (x,, m0o) dt.
(13)
Ri [MASK]
17

---

Published as a conference paper at ICLR 2026
where De, denotes the denoiser assigned to block b. The global loss decomposes as L =>== Co-
This derivation shows that DiffusionBlocks in masked diffusion models amounts to partitioning the masking schedule a(t) rather than time. Each block is responsible for an equal decrement in a(t), ic. an equal share of the total "demasking work". which ensures balanced parameter utilization and true independence across blocks. This construction is directly analogous to the equi-probability partitioning continuous diffusion models described in Section
E EXPERIMENTAL DETAILS
Unless otherwise specified, all experiments use the following settings. For DiffusionBlocks, adopt the EDM framework (Karras et al. 2022) with default parameters: log-normal noise distribu- tion with Prean -1. and Pau 1.2, noise range minn mmax] [0.002, 80), and preconditioning following the recommended contiguration. Inference uses Euler sampling with 50 steps unless stated otherwise. During training, blocks are sampled uniformly random for each iteration.
E. VISION TRANSFORMERS FOR IMAGE CLASSIFICATION
For image classification experiments Section 5.1 we use a 12-layer VIT with patch size 4, 128 hidden dimensions, 4 attention heads, and 0.1 dropout, partitioned into B=3 blocks (4 layers each). We train for 500 epochs with batch size 128 and AdamW optimizer with learning rate 5 x 10-4, We employ cosine learning rate scheduler with 10-epoch linear warmup. As data augmentation, we apply random horizontal flipping (p = 0.5) and RandAugment (Cubuk et al. 2020) as data augmentation.
Figure 5] (top left) illustrates the crrfussioBllokk adaptation for VIT. We add noise to the class label embeddings and concatenate them with the patch embeddings. Each block learns to denoise the label embedding conditioned on the patch embeddings. We use an overlap ratio 7 = 0.05 and perform denoising steps during inference (matching L/ B= 12/3). The classification head is applied after the final denoising step to produce class predictions. We minimize cross-entropy loss between predicted and true class labels during training. For the Forward-Forward baseline, we adapt the Contrastive Neeward-Forwand (FF) (Aghagolzadeh & Ezoji, 2025) implementation to ViT
E.2 DIFFUSION MODELS FOR IMAGE GENERATION
For image generation experiments in Section 5.2 we use DiT-S/2 (12 layers) for CIFAR-10 and DřT- 1.2 (24 layers) for ImageNet-2566 Both models are partitioned into B 3 blocks. Training follows the EDM framework with classifier-fres guidance (Ho & Salimans (02T) (10% label dropout). For CIFAR-10, we train for 100 epochs with batch size 512 and AdamW optimizer with learning rate
10 10-4 For ImageNet, we resize to 256x256 and encode images by pre-trained VAE (Peebles &
Xie 2023 We also train 100 epochs with batch size 512 and Adam optimizer with learning rate 5 x 10 Overlap ratio is set to y 0.05.
In evaluation, we apply Euler sampling with 50 steps and classifier-free guidance (scale 2.0) on both CIFAR-10 and ImageNet experiments. FID is computed using 50,000 generated samples against the training and test sets, with the minimum of three evaluations reported following Karras et al.) (002)). For the training set, we use the official ADM (Dhariwal & Nichol, (021) evaluation suite, which computes FID against the entire training set as the reference distribution. For the test split, compute FID using clean- fid (Parmar et al. 2022).
E.3 MASKED DIFFUSION MODELS FOR TEXT GENERATION
In Section 5.3 we follow MD4's training protocol with 256 sequence length, AdamW optimizer with learning rate 3 x 10-1, weight decay 0.03, and 2,000 linear warmup steps. Training runs for 100 epochs with batch size 256. The 12-layer DiT-based transformer (Lou et al. 2024 Sahoo et al. 2024) uses 768 hidden dimensions and 12 attention heads, partitioned into B=3 blocks with
com stabi stabi ft- ema
18

---

Published as a conference paper at ICLR 2026
overlap ratio - 0.05. Masking schedule follows MD4's linear schedule. For block partitioning in discrete diffusion, we apply equi-probability partitioning to the masking ratio distribution rather than continuous noise levels in Appendix D set following [Shi et al. 2024).
Bits-per- Sits-per-charaeter (BPC) is evaluated the text8 test
E.4 AUTOREGRESSIVE MODELS FOR TEXT GENERATION
In Section 5.4. we use a 12-layer Llama style transformer (Touvron et al. 2023) augmented with time conditioning in DiT [Peebles& & Xie 2023) with 768 hidden dimensions, 12 attention heads, and the Llama-2 tokenizer with 32K vocabulary size. The model is partitioned into B=4 blocks with an overlap ratio 1. Training uses sequence length 256 for LM1B and 3072 for OWT, batch size 256, Adam W with learning rate 3 x 10-4, and 2500 warmup steps for 10 epochs.
Since DiffusionBlocks is not derived from ELBO-based objectives, computing traditional perplexity is non-trivial. Instead, we evaluate using MAUVE scores following SEDD (Lou et al. 2024), which measures the similarity between generated and real text distributions. For each test sample, generate 5 continuations of 50 tokens from 1K prompts and compute MAUVE against 1K reference samples with the scaling factor 0.2. Additionally, we report generative perplexity, commonly used in diffusion language models (Lou et al. 2024; Sahoo et al. 2024), by computing the perplexity of generated text using teacher models (Llama -2- 7 and GPT2-XL. (Radford et al. 2019 T. For generations, we use top-p sampling (0.95) for the baseline and 4 diffusion steps with greedy sampling for DiffusionBlocks. The OWT test set is created by splitting 10% of the data since no official test set exists.
Applying DiffusionBlocks to autoregressive models requires maintaining causal consistency during training. When denoising future tokens, the model must condition on clean past tokens rather than noisy ones to preserve the autoregressive property. Following Block Diffusion (Ariiol et al. (025). implement this using sequence concatenation: noisy and clean sequences concatenated with a modified causal attention mask that allows noisy tokens to attend to their corresponding clean past tokens while preventing information leakage. This approach doubles sequence memory but main- tains single forward pass efficiency. An alternative implementation computes key-value pairs sepa- rately for clean and noisy sequences, combining them during attention computation. This requires two forward passes but uses standard sequence memory. We adopt the concatenation approach for computational efficiency.
E.5 RECURRENT-DEPTH MODELS
we use the default configuration: prelude layers, 4-layer recurrent block, and 2 coda layers following Pythia-700 700 (Biderman et al,, 2023 architecture with 512 hidden dimensions and attention heads. Unlike other architectures, recurrent-depth models do not require block partitioning since the entire network is applied recur- rently. Instead, we train the full network as denoiser by sampling different noise levels each training step. While baseline Huginn uses stochastic recurrence depth (average 32 iterations) with truncated BPTT (8 steps), DiffusionBlocks trains with single-pass diffusion. We train on LMIB for 15 epochs compared to Huginn's 's 5 epochs. Despite this, our approach uses approximately 10x less total computation since we avoid the 32x recurrent iterations during training.
For Hugim nn (Geiping al. 2025 described in Section 5.5
E.6 ABLATION STUDIES
E. COMPARISON WITH NOPROP
We follow the experimental protocol of NoProp (Li et al., 2025). In the absence of publicly available code, we implemented their NoProp-DT architecture augmented with time condition- ing from NoP rop- following their specifications (Figure 5 in their paper). Training follows NoProp-CT's hyperparameters with AdamW optimizer, learning rate 10 10-4 batch size 128, and 1000 epochs on CIFAR-100. For DiffusionBlocks, we use B=3 blocks with overlap ratio 7 = 0.1.
19

---

Published as a conference paper at ICLR 2026
Table 9: VIT results on Imay-lmageNet. DiffusionBlocks shows consistent performance on inter classification dataset.
Method
Accuracy (1)
ViT
35.32
DiffusionBlocks
36.16
Following NoProp-CT's evaluation protocol, we use 1000 Euler sampling steps instead of our default 50.
We attempted to adapt BewaaddFFowwart (FF) algorithm (Hinton 2022) as an additional baseline to NoProp's architecture for Table G However, without publicly available code and with no spec- ified adaptation procedure, the implementation requires numerous design decisions. Our attempts achieved only 1 accuracy, highlighting the fundamental incompatibility: NoProp's architecture is specifically designed for their method (type (e) their Figure 2), while requires contrastive pos- itive/negative samples (type (d)). Successfully bridging these paradigms may require innovations beyond straightforward adaptation. This highlights a key distinction between approaches. NoProp does not provide guidance for adapting other methods architectures. DiffusionBlocks instead offers systematic procedure for converting existing Transformer- networks into block-wise trainable models. This recipe enabled successful application to modern architectures with minimal modifications, demonstrating the generality of our framework.
E.6.2 DESIGN CHOICE ANALYSIS
All ablation studies follow the configurations described in Appendix E.2 We report FID scores on the test splits. For partitioning experiments, we test both uniform partitioning (equal intervals in log-space) and our equi-probability method. Layer distribution indicates the number of layers in each block. For block count experiments, we vary B from to while keeping total layers fixed at 12. We disabled the block overlap (y 0.0) in Section C to isolate the effectiveness of each component.
F ADDITIONAL EXPERIMENTS
F.I IMAGE CLASSIFICATION EXPERIMENT ON TINY IMAGENET
To further evaluate the effectiveness of DiffusionBlocks on classification tasks beyond CIFAR- 100,
conducted additional experiment the Tiny ImageNet dataset (Le & Yang 2015). This dataset consists of 200 classes, 100,000 training images with each image resized to 64 resolu- tion. Tmageet provides more challenging and higher-resolution benchmark than CIFAR- 100.
We trained 12-layer Vision Transformer (ViT) with patch size 4, hidden size 768, and 12 attention heads. Both the baseline ViT and DiffusionBlocks models were trained for 100 epochs using a batch size of 256 and the AdamW optimizer with learning rate of 10-4 14,4, For DiffusionBlocks, we used B = 2 blocks (each containing 6 layers).
Table demonstrates that DiffusionBlocks maintains competitive performance relative to the base- line VIT, consistent with our findings CIFAR-100 as well our large-scale classification ex- periments in language modeling (LMIB and WebTeeeText in Table 3 5). These results further indicate that DiflusionBlocks remains effective classifier across different data modalities, reso- lutions, and dataset scales.
F.2 EFFECT OF BLOCK COUNT ON CIFAR-10 10
To examine whether the design trends observed on ImageNet in Table g generalize to different datasets, we additionally evaluate the effect of the number of blocks CIFAR-10. This experiment allows us to assess whether the behavior of DiffusionBlocks remains consistent across datasets of different scales and complexities. We use the same DiT-S/2 architecture described in Section 5.2
20

---

Published as a conference paper at ICLR 2026
Table 10: Effect block count on CIE.R-I0.
Table 11: Effect of block count on text gen- eration (LM1B). Best performance is achieved with B=4.
Moderate block counts (2-3) achieve the best FID,
showing consistent trends with ImageNet in Table
Number of Blocks
FID (1)
L/B (
Relative Speed
Number Blocks
MAUVE
per Bloek Speed
39.83
12
1.0x
B =
2
35.47
6
2.0x
6-
0.42
0.67
6.0%
38.03
3.0x
45.43
4.0x
B
53.32
2
6.0x
training under the EDM framework while varying the number of blocks B {1,2,3, 6) and disabling block overlap (7 0.0) to isolate the effectivenss of the number of blocks B.
As shown in Table TO smaller block counts tend to achieve better FID scores, and B 2 B provides strong performance. This trend matches the observations in Table These results indicate that the effectiveness of using moderate number of blocks consistent across datasets of varying scale, supporting the validity of the design choices analyzed in Section 5.6
F.3 EFFECT OF BLOCK COUNT TEXT GENERATION
Table I1 shows the effect of varying the number of blocks for autoregressive language modeling on LMIB with overlap ratio 0.0.
The optimal number of blocks differs between tasks: image generation achieves best FID with B=2 or B=3 (Table (), while language modeling achieves best MAUVE with B=4. This motivated our choice of B=4 for language modeling experiments the main paper.
G COMPARISON WITH ACTIVATION CHECKPOINTING
DiffusionBlocks and activation checkpointing (also known as activation recomputation, gradient checkpointing, rematerialization) offer fundamentally different trade-offss and can be powerfully combined.
The key distinction lies in what each method reduces. Activation checkpointing reduces only acti- vation memory, leaving parameters, gradients, and optimizer states unchanged. In contrast, Diffu- sionBlocks reduces all memory components by factor of B. This distinction becomes increasingly critical as modern models grow larger.
To illustrate this difference, consider L-layer network where each layer has parameter size P and activation size A. With Adam optimizer (requiring 2P for momentum and variance), each layer needs 4P memory for parameters, gradients, and optimizer states. Standard training thus requires (4P + A)L total memory. Activation checkpointing reduces this to 4PL + A by rematerializing activations only when needed (though this is an optimistic estimate that ignores the memory cost of the checkpoints). DiffusionBlocks, by training / independent blocks, requires (4P + (L/ B). Since L. > B, combining DiffusionBlocks and activation checkpointing uses the least memory among these four patterns.
Regarding computational costs, it is empirically known that activation checkpointing increases the training time by factor of approximately 4/3, and this holds true when combined with the proposed method. This is justified as follows. With a forward pass computation cost of F, a backward pass requires approximately 2F (computing Jacobians and weight gradients). Standard training 3F cost per iteration, while activation checkpointing increases this to 4/" due to recomputation. DiffusionBlocks maintains this ratio when combined with checkpointing.
Beyond memory reduction, DiffusionBlocck offers unique advantages regarding training time: each block can be trained in an embarrassingly parallel manner. This means each block can be trained in parallel with absolutely no communication overhead. This provides an additional advantage over activation checkpointing, especially when computational resources abundant.
21

---

Published as a conference paper at ICLR 2026
Table 12: Wall-time comparison on VIT. The aggregated DiffusionBlocks time is computed by multiplying the measured per-block iteration time by B 3.
Method
Wall time (sec/iter)
ViT
0.0507
DiffusionBlocks: per-block time (4 layers)
0.0181
DiffusionBlocks: aggregated time (0.0181
0.0543
H TRAINING AND INFERENCE EFFICIENCY
This section provides detailed analysis of the computational efficiency and wall-time characteris- ties of DiffusionBlocks.
Training efficiency. Consider an L-layer network trained for K iterations. Standard end-to-end backpropagation performs K L layer evaluations. DiffusiinBlocks trains only L/B layers at time; training all blocks for K iterations each performs x LxK layer evaluations. Thus, DiffusionBlocks requires the same total amount of computation as standard training, while reducing memory usage by a factor of B.
To validate this theoretical equivalence, measured the per-iteration wall time using 12-layer VIT on a single H100 80GB GPU, averaging over 100 iterations. As summarized Table [2 standard training requires 0.0507 seconds per iteration for all 12 layers. Under DiffusionBlocks with B 3, each block (4 layers) takes 0.0181 seconds per iteration (measured). The total per-iteration wall time for DiffussinBlocks is therefore obtained by summing the independently trained blocks. computed as 0.0181 x 3 0.0543 seconds. The resulting end-to-end wall time is thus comparable to standard training, with the small difference attributable to the noise-level conditioning introduced during the DiffusionBlocks conversion (Section 3.1
Inference efficiency. For inference, we ensure that the total amount of computation matches that of the baseline model. For 12-layer network, the baseline performs single forward pass through all 12 layers. Under DifsusionBlocks with B = 3, we perform three denoising steps, each invoking the corresponding 4-layer block once. The total compute therefore corresponds to the same 12 layer evaluations as standard inference.
For diffusion models used in image generation, the computational benefit is even more pronounced. Standard diffusion models must apply the full network for every denoising step. With 50 denoising steps, 1 12-layer DiT requires 12 50 layer evaluations. In DiffusionBlocks, each denoising step applies only the block responsible for that noise level, which contains 4 layers when B = 3. This reduces the total compute to 4 x 50, achieving a B-fold reduction inference cost. The 50 denoising steps are assigned to blocks according to the equi-probability partitioning in Section so that each block is used approximately the same number of times during inference. Euler sampling used for simplicity, and, shown in Section 2.2 is computationally equivalent to a residual update, requiring no additional overhead.
22


Reviews:
[{"review_id": "MTZf64PnJq", "reviewer": "Reviewer_WpqB", "summary": "This paper proposes to divide the network into K smaller subnetworks and train them independently.\nIn terms of architecture the paper assumes transformer-like residual networks, which are composed of a sequence of blocks.\nA diffusion objective is used and noise level intervals are assigned to to each subnetwork.\nDuring training, each subnetwork is trained alone and independently to denoise data in its assigned noise range.\nThe blocks are adapted so that their residual nature can be reframed as a denoising update of the form $z_{i+1}=z_i + \\mu_{i+1} f(z_i)$, and in addition their input and ouput shapes should match the input data shape.\nDuring inference, following the noise schedule, subnetworks for denoising are selected on the criteria that the current noise levels falls into their assigned noise range.\nThe proposed method makes both training and inference faster.\nIn addition this makes training scalable to larger architectures which as whole couldn't fit in current GPU memory but for which subnetworks can fit in memory.\n\nThe noise level assignment is done such that each network receives on average the same number of training noisified samples.\n\nIn addition, the method also targets recurrent-depth models.", "questions": "1. In Fig.2 I am confused by your drawing, for step 3, you extract a block in the middle and refer to its output which is fed as an input with noise. What output are we talking about here, the actual output of the subnetwork or the final expected output of the whole network? From your algorithm in Fig.3 (the orange box), it appears it's the expected output $y$.\n2. In 5.1 (ViT for image classification), you mentioned that noise is added to class label embeddings:\n2.1 Where do class label embeddings come from, or are they really one-hot representations?\n2.2 What type of noise is added to these embeddings? From Figure 2 it appears you add Gaussian noise to one-hot labels, is that right?\n3. More generally, by slicing a network into K subnetworks, it stands to reason that each subnetwork has less capacity (given that the total number of weights remains about the same). How do you explain that your reported results in tables (1,2,3,4,5) are better than the full models they compare too while at the same much faster to train and with less capacity per subnetwork?\n4. Why not also reporting performance with CFG on ImageNet256? Does your method works with CFG?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The method looks simple and easy to implement.\n2. The method targets an important problem: scalability and training/inference speed of large models.\n3. For the most part, the paper is well written and easy to read.", "weaknesses": "1. However the notation sometimes lacks clarity, for example it's unclear what are $(x,y)$. From section 2.1, it appears that $y$ is data (e.g. an image for example). So that would make $x$ a label supposedly, however in Fig.3 $z_0\\leftarrow x$ which suggest data rather than label. In Figure 2, in the case of the classifier, it would like $y$ is actually a label. A simple fix would be to clearly specify what $x$ and $y$ are and keep the notation consistent thorough the entire paper.\n2. The results seem to compare to weak baselines, for example CIFAR100 classification baseline has an accuracy of 45% while SotA easily exceeds 80% if memory serves me well. Similarly DiT FID on CIFAR10 is reported at 39.83 while a google search reveals people obtaining FIDs within 10-15 on the DiT GitHub repo.\n3. Nit, provide important information in the paper when possible, for example I had to dig into appendix D2 to find the resolution used for ImageNet (aka 256).\n4. CIFAR10 is a surprising choice for benchmarking the method given that it's well known that CIFAR10 is subject to overfitting and uses heavy regularization whether in classification or diffusion. Therefore, given your method spreads the capacity over each subnetwork, it is unclear how much can be really explained by your $B$ ablation in table 8. I believe this ablation should be done on either ImageNet64 or 256.\n5. On ImageNet256, the FID numbers are pretty high (10-12), it appears you DiT-L/2 while the DiT paper has most comparison for XL/2 (which gets an FID around 10), digging into the DiT paper appendix, the reported FID without CFG for DiT-L/2 is 23.33. I am confused, did you run DiT-L/2 or DiT-XL/2 as your results don't seem to match the ones reported by the original paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "S6tqyZ3av7", "reviewer": "Reviewer_wyrV", "summary": "The paper proposes a block-wise training strategy that splits standard transformer-based neural networks into groups of layers (blocks) which can then have forward/backward computed separately. This approach saved memory equal to the number of blocks, e.g. Bx memory savings. The paper achieves this by reconceptualizing the role of a block as a subnetwork that needs to denoise the output within a certain range. So, for a network with 12 layers spilt into 3 blocks, each group of 4 layers is responsible for removing a designated amount of noise, and the denoised output of the final set of 4 layers constitutes the actual prediction. The paper shows this method can compare to end-to-end training for image classification, image generation, and text generation.", "questions": "Since this method does not seem to be immediately practically useful, I am mainly viewing the strength in terms of the novelty (which I think is quite high). However, my mind could be changed (either for better or worse) depending on the following:\n\n1. Can this be used effectively in tandem with recompute/checkpointing? What is the resulting impact on training time?\n\n2. How long does this method take to train for ViT, DiT, etc.? Wall time preferred.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "S1. The method seems very novel. Converting these disparate tasks (classification, image generation, text generation) into denoising tasks and separating the forward/backprop for the network into different groups of layers (blocks) seems to solve a problem that prior block-wise training methods could not.\n\nS2. Under the specified setups, the block-wise training matches or beats the baseline.", "weaknesses": "W1. On practical terms, this method would seem to have no advantage over recompute/checkpointing. In fact, the memory savings from recompute are substantially higher.\n\nW2. It is not clear how this affects training efficiency (wall time). \n\nW3. Baselines are odd. CIFAR-100 is a little toy-ish, and the default ViT performs quite poorly on it. The DiT is also under-optimized, and I assume this applies to other things as well. So it's not actually clear that this method can match any end-to-end training in practical setups, where these models are fully optimized.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "T0HKoWPizh", "reviewer": "Reviewer_ZJo4", "summary": "This paper presents a novel method to transform modern networks with residual connections (e.g., ViT) to independent trainable blocks that learns a diffusion denoising process for a predefined noise range. During training, all blocks rely only on the input x and not output of previous/other blocks. During inference, input x is passed as a condition to blocks starting with the largest noise range to the smallest noise range. The method does not require backpropagation across all blocks, reducing memory usage. Experiments on numerous tasks demonstrate promising results.", "questions": "- Each block learns difference noise levels independently. In this case, why use mulltiple blocks and not just one large block to learn all noise ranges?\n- The last block of the network reconstruct the original input. Why can this be used for classification? Conditioned on the input, the denoising blocks reconstruct the input. This does not yield a standard representation like other models that can be used for deterministic tasks. The authors have demonstrated classifcation performance on CIFAR but I think this is not sufficient. I would really like to see whether the approach works for ImageNet-1k as well as downstream dense prediction tasks such as semantic segmentation. Due to the time limit of the rebuttal, partial results are ok, for example, training only 100 epochs instead of 300 (standard ImageNet-1k training setting).\n- If I want to finetune a pretrained model on a dataset, can the weights and knowledge learned by the pretrained model ViT be used with this approach?\n\nOverall, I think the idea is novel, but there are some questions that needs clarifying and I'm have concerns on the practical application and performance of the proposed approach.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The main idea of this paper is novel (concurrent to NoProp). It tries to convert the foward propagation proccess of residual-based architectures to a diffusion denoising process.\n- The method is evaluated on numerous tasks (image + text classifcation and generation)\n- There are some theory analysis of the partitioning approach", "weaknesses": "- In Table 2, the number of training epoch/iterations are not reported. For DiT-L/2 on ImageNet256x256, what are the FID for with and without classifier free guidance (CFG). The current result does not state whether CFG is used.\n- To my understanding Equation 4 should only denote one residual connection. In ViT, one residual connection is for the self-attention operation and one residual connection is for the MLP. To be rigourous with the thoery, each denoising block should denote a single residual connection so each ViT block should be decomposed into two diffusion blocks. In the current implementation ViT/DiT with 12/24 layers are partitioned into 3 blocks. From this perspective, I don't think the thoery fully back up the proposed partition approach.\n- Despite that the proposed method uses less memory than standard models, diffusion processes are more complex. I am worried that the acutal speed of the method is much slower than standard ViT. The authors should report real time comparison of training and inference speeds on GPUs for example on the image classifcation and generation tasks.\n- The authors should report the number of denoising steps for each block during inference.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)