How Do DiTs Denoise? An Aggregator-Predictor View of Attention and MLPs
Authors: Gyeonghun Kang, William Xingxu Chen, Xiang Cheng
Organizations: Department of Statistical Science, Duke University · Department of Computer Science, Duke University · Department of Electrical and Computer Engineering, Duke University
How do diffusion transformers (DiTs) represent image content and learn rules for reusing it to generalize? We develop an aggregator--predictor view of DiTs as input-adaptive local denoisers, in which context aggregation and patch prediction provide functional roles for interpreting attention and MLPs. We first establish that, for Gaussian Markov random field (GMRF) priors, the global Bayes denoiser can be approximated using only a subset of context patches, with error decaying exponentially in context radius. We further show that patchwise linear denoisers, often invoked to explain diffusion generalization, recover the empirical local denoiser with sufficiently informative aggregation. Within this view, AdaLN-modulated attention naturally parameterizes an input- and noise-dependent aggregator, while the MLP admits a learned key--value representation for tokenwise patch prediction. Our experiments support this view. DiTs trained on real images exhibit noise-dependent denoising contexts resembling those in the analytic GMRF setting, while attention mass exhibits a qualitatively corresponding broadening pattern. In a controlled CelebA64 study, shared-memory patch predictors used in place of DiT MLPs achieve competitive denoising and generative performance. Finally, cross-dataset component transplantation reveals compatibility patterns consistent with the proposed roles of QK in spatial weighting, VO in feature transport, and MLPs in patch prediction. Together, these results provide a component-level account of DiT denoising and motivate architecture-specific fine-tuning and controlled intervention.
Figures & tables
Figure 1: In each row, a CelebA source checkpoint is fine-tuned on the AAHQ, FFHQ, AFHQ, or STL-10 dataset. The source and fine-tuned generated images are shown in columns (000) and (111). In the middle columns, attention QK/VO and MLP weights of the source checkpoint are replaced with the corresponding fine-tuned weights; e.g., (001) means only MLP is replaced, while QK/VO are retained. For face-to-face fine-tuning, replacing only QK generates source-like faces, while replacing MLP brings visual changes: cartoon faces, bald heads, or dogs. Among the tested CelebA → STL-10 hybrids, coherent target-like structure appears primarily in the joint endpoint (111); reverting VO (101) leaves fragmented texture. See Sec. 5.3 and Appendix E.4 .
Figure 2: Exact context set in a Gaussian MRF prior. The known graph and precision matrix allow exact analysis of context set in this synthetic setting. (a) Radius-3 contexts B3(m) for two selected patches. Defined by graph distance, these sets need not be spatially contiguous. (b) Exact and certified radii at ϵ=10−3 . Corollary 2 guarantees that the certified radius is nondecreasing up to σabs (dotted lines); the exact radii can be substantially smaller. (c) 90%/95% sensitivity radii (Sec. 5.1 ), showing that relative influence spreads over a broader context as noise increases. Details in Appendix E.1 .
Figure 3: Sensitivity-defined contexts broaden with noise, with attention mass showing a related trend. Gold/white contours enclose 50%/95% of normalized squared sensitivity or attention mass; ticks point inward and cyan dots mark target patches. Right: 95% radii averaged over target patches; Details in Appendix E.2 .
Figure 4: Shared-memory patch predictors at comparable parameter budgets. Models are trained on CelebA64 for 400 epochs; see Appendix E.3 for details.
Figure 5: QK, VO, and MLP component transplants. (a) Source-QK compatibility is higher across spatially similar datasets. (b) VO reversion costs more as joint non-additivity increases. (c) Natural-image source QK retains partial agreement on face targets. See Appendix E.4.2 .
Appendix figures & tables19 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Component map for the aggregator–predictor view. At the branch level, AdaLN-conditioned QK supplies spatial weights, VO transforms the weighted features, and the MLP admits a learned-memory interpretation for tokenwise prediction. The patchwise linear denoiser (PLD) instead uses a prescribed context and retrieves training patches.
Predictor
Memory bank
Bank parameters
Layerwise parameters
LPLD gℓ,m,tlpld
{z(r)}r=1R
RMdh′
3dh2+dhdh′+4dh
SCA gℓ,tsca
{k(r),v(r)}r=1R
2Rdh′
2dh(dh+dh′)+5dh+H
SMLP gℓ,tsmlp
{k(r),v(r)}r=1R
2Rdh′
2dh(dh+dh′)+Hdh′+2dh
MLP gℓ,tdit
{kℓ(r),vℓ(r)}ℓ,r
2LRdh
0
Appendix
Table 1: Predictor parameter comparison. The layerwise column includes Q/K/V/output projections, their biases, and SCA’s headwise gain parameters at each layer. For LPLD, a learned compatibility map adds dhdh′ parameters once; this cost is zero when A=I .
Dataset
Content
Training
Evaluation
CelebA ( Liu et al., 2015 )
Aligned celebrity photographs
162,770
19,962
FFHQ ( Karras et al., 2019 )
Photographic faces
60,000
5,000
AAHQ ( Liu et al., 2021 )
Artistic portraits
15,366
1,921
AFHQ-v2 ( Choi et al., 2020 ; Karras et al., 2021 )
Cat, dog, and wild-animal faces
14,335
1,467
iCartoonFace ( Zheng et al., 2020 )
Cartoon-face recognition crops
388,710
22,496
STL-10 ( Coates et al., 2011 )
Natural-object images
104,500
8,000
Appendix
Table 2: Datasets used throughout the paper. Counts refer to prepared images after filtering. Evaluation uses the held-out splits described below; the ImageNet64 evaluation set is its official validation split.
Figure 7: Sensitivity and attention radii and patch counts on four datasets. Solid/dashed curves show 90%/95% coverage, computed from population maps over 1,024 held-out images and averaged over all 256 target patches.
Figure 8: Additional context maps for iCartoonFace and STL-10 at σt=0.75,49 , complementing the CelebA and ImageNet examples in Fig. 3 . Colors are logarithmic; gold/white contours mark 50%/95% coverage, ticks point inward, and cyan dots mark target patches.
Predictor
R
dh′
Parameters, H=1
Parameters, H=8
LPLD
32
128
8,935,472
8,935,472
LPLD
32
256
10,213,424
10,213,424
LPLD
32
512
12,965,936
12,965,936
SCA
1024
256
8,642,616
8,642,672
SCA
1024
512
10,215,480
10,215,536
SCA
1024
1024
13,361,208
13,361,264
Appendix
Table 3: Complete patch-predictor grid. Each numeric entry in the last two columns specifies one trained configuration. All models use dh=256 , L=8 , and H~=8 ; predictor head width is 256 for H=1 and 32 for H=8 . R counts image-shaped memories for LPLD, shared prototype pairs for SCA/SMLP, and hidden units for MLP.
Figure 9: Training trajectories for the patch-predictor comparison. Held-out velocity loss, denoising log- σ AUC MSE, and FID-50k on CelebA64. Colors identify predictor families; circles and triangles indicate H=1 and H=8 . Solid, dashed, and dotted curves denote increasing dh′ within each shared-memory family, or increasing hidden width R for MLP.
Transfer
Epoch
FID
Transfer
Epoch
FID
CelebA → AAHQ
75
40.767
STL-10 → AAHQ
150
40.592
CelebA → FFHQ
100
25.252
STL-10 → FFHQ
100
28.231
CelebA → AFHQ-v2
200
44.821
STL-10 → AFHQ-v2
200
49.119
CelebA → STL-10
200
68.050
STL-10 → CelebA
200
17.262
ImageNet64 → AAHQ
150
34.493
ImageNet64 → FFHQ
100
22.640
ImageNet64 → AFHQ-v2
200
38.840
ImageNet64 → CelebA
100
14.620
Appendix
Table 4: Selected joint-adaptation endpoints. Epochs use the smaller source/target training set as reference; FID is the two-seed mean FID-50k used for checkpoint selection. STL-10 and ImageNet64 source checkpoints are both from epoch 800.
Figure 10: Spatial self-similarity distance between real datasets. Unbiased Gaussian-kernel MMD 2 of spatially normalized, fixed-grid DINOv2 patch descriptors. Cell and colorbar values are in units of 10−2 . Off-diagonal cells compare 1,000 unpaired images from each dataset; diagonal cells compare disjoint 500-image halves. Lower means more similar descriptor distributions.
Figure 11: Agreement of all hybrid states with the jointly adapted endpoint. Global-feature endpoint agreement EA (left) and relative spatial endpoint agreement SA (right), averaged over two adaptation seeds. Each row is a transfer direction, and columns give QK/VO/MLP bits (0: source, 1: jointly adapted). The same 1,024 noise latents are used for all eight states of each cube. Each score is 0 for the source and 1 for the joint endpoint; negative values are possible.
Interpretation
State difference
ρEA
ρSA
QK reversion penalty †
Agr111−Agr011
0.739
0.844
MLP adaptation cost (VO adapted) †
Agr010−Agr011
0.718
0.809
VO adaptation cost (QK adapted)
Agr100−Agr110
0.771
0.683
MLP-only adaptation cost
Agr000−Agr001
0.746
0.673
MLP adaptation cost (QK adapted)
Agr100−Agr101
0.644
0.673
VO reversion penalty
Agr111−Agr101
0.641
0.673
Appendix
Table 5: Spatial distance and nested state differences. Spearman correlations with dataset spatial MMD 2 use 12 two-seed direction means for each score. Each row changes one component with its partners fixed. Rows are sorted by decreasing min(∣ρEA∣,∣ρSA∣) .
Figure 12: Spatial distance and functional coordination. Spatial MMD 2 versus (a) joint non-additivity, (b) minimum reversion penalty, and (c) three-way interaction, defined in Appendix E.4.2 . Points average two adaptation seeds; colors, markers, and lines follow Fig. 5 .
Reverted component
State difference
ρEA
ρSA
VO
Agr111−Agr101
0.832
0.930
QK
Agr111−Agr011
0.601
0.573
MLP
Agr111−Agr110
0.664
0.063
Appendix
Table 6: Component dependence and joint non-additivity. Spearman correlations between JNAAgr and each displayed difference use 12 two-seed direction means, separately for Agr=EA,SA . Note that since all quantities share cube terms, the differences between correlations are descriptive.
Figure 13: CelebA source: three additional latents across four targets.
Figure 14: STL-10 source: three latents across four targets.
Figure 15: ImageNet64 source: three latents across four targets.
Figure 16: CelebA-source adaptation trajectories. Rows are AAHQ, FFHQ, AFHQ-v2, and STL-10 targets.
Figure 17: STL-10-source adaptation trajectories. The source checkpoint is from epoch 800.
Figure 18: ImageNet64-source adaptation trajectories. The source checkpoint is from epoch 800; rows are AAHQ, FFHQ, AFHQ-v2, and CelebA targets.
Diffusion Transformers use attention-based networks to generate and transform images, video, audio, and other structured data through iterative denoising.
Generating images beyond the training resolution of Diffusion Transformers (DiTs) remains challenging, as the number of tokens grows rapidly with resolution, full attention becomes expensive, and outputs often suffer from structural distortions and degraded detail. We identify uniform full attention as the common source of both problems. Attending to all tokens at every step causes quadratic cost and exposes the model to more tokens and relative distances beyond the training range, pushing the model out of distribution. We observe that attention's effective receptive field narrows during denoising, from global in early steps that form structure to local in later steps that refine details. We propose AReA, a training-free attention mechanism whose receptive field evolves from coarse to fine, following the denoising progression. In early steps, AReA applies cross-scale coarse attention over downsampled keys and values, capturing global context, while in later steps it applies local attention within progressively shrinking windows. Without additional parameters or training, experiments show that AReA consistently outperforms state-of-the-art training-free baselines across target resolutions and different models, while also being faster and more efficient.
Transformer architectures, particularly Diffusion Transformers (DiTs), have become widely used in diffusion and flow-matching models due to their strong performance compared to convolutional UNets. However, the isotropic design of DiTs processes the same number of patchified tokens in every block, leading to relatively heavy computation during training process. In this work, we introduce a multi-patch transformer design in which early blocks operate on larger patches to capture coarse global context, while later blocks use smaller patches to refine local details. This hierarchical design could reduces computational cost by up to 50% in GFLOPs while achieving good generative performance. In addition, we also propose improved designs for time and class embeddings that accelerate training convergence. Extensive experiments on the ImageNet dataset demonstrate the effectiveness of our architectural choices. Code is released at: https://github.com/quandao10/MPDiT
Denoising diffusion models are the dominant architecture for image generation, whereas most natural language generation and modeling are primarily handled by well-known transformer architectures employing attention mechanism. Here, we show that diffusion models also inherently use an attention mechanism very similar to that of transformers. Therefore, attention emerges as a universal machine learning principle, based on a general training objective. We also show similarities in basic functional principle of auto-encoders and attention-based models. These equivalences allows us to interchange these designs based on practical requirements. As an example, we can reformulate the diffusion framework to reduce the lengthy training process and computation-intensive image generation. Using this approach, a simplified algorithm is proposed for image generation which is based on attention mechanism. Results show that the attention-based implementation achieves comparable performance with significantly less effort and computational resources.
Farzan Haddadi, Leila Monfared, Ebrahim Rezaii +3
School of Electrical Engineering, Iran University of Science & Technology, Tehran, IRAN · Independent researcher