How Do DiTs Denoise? An Aggregator-Predictor View of Attention and MLPs

Gyeonghun Kang; William Xingxu Chen; Xiang Cheng

cs.CVOct 21, 2026

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

Abstract

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 $\to$ 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

B_{3}(m)

for two selected patches. Defined by graph distance, these sets need not be spatially contiguous. (b) Exact and certified radii at

\epsilon=10^{-3}

. Corollary 2 guarantees that the certified radius is nondecreasing up to

\sigma_{\mathrm{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_{\ell,m,t}^{\mathrm{lpld}}$	$\{z^{(r)}\}_{r=1}^{R}$	$RMd_{h}^{\prime}$	$3d_{h}^{2}+d_{h}d_{h}^{\prime}+4d_{h}$
SCA $g_{\ell,t}^{\mathrm{sca}}$	$\{k^{(r)},v^{(r)}\}_{r=1}^{R}$	$2Rd_{h}^{\prime}$	$2d_{h}(d_{h}+d_{h}^{\prime})+5d_{h}+H$
SMLP $g_{\ell,t}^{\mathrm{smlp}}$	$\{k^{(r)},v^{(r)}\}_{r=1}^{R}$	$2Rd_{h}^{\prime}$	$2d_{h}(d_{h}+d_{h}^{\prime})+Hd_{h}^{\prime}+2d_{h}$
MLP $g_{\ell,t}^{\mathrm{dit}}$	$\{k_{\ell}^{(r)},v_{\ell}^{(r)}\}_{\ell,r}$	$2LRd_{h}$	$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

d_{h}d_{h}^{\prime}

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 $\sigma_{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$	$d_{h}^{\prime}$	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

d_{h}=256

L=8

, and

\tilde{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-

\sigma

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

d_{h}^{\prime}

within each shared-memory family, or increasing hidden width

R

for MLP.

Transfer	Epoch	FID	Transfer	Epoch	FID
CelebA $\to$ AAHQ	75	40.767	STL-10 $\to$ AAHQ	150	40.592
CelebA $\to$ FFHQ	100	25.252	STL-10 $\to$ FFHQ	100	28.231
CelebA $\to$ AFHQ-v2	200	44.821	STL-10 $\to$ AFHQ-v2	200	49.119
CelebA $\to$ STL-10	200	68.050	STL-10 $\to$ CelebA	200	17.262
ImageNet64 $\to$ AAHQ	150	34.493	ImageNet64 $\to$ FFHQ	100	22.640
ImageNet64 $\to$ AFHQ-v2	200	38.840	ImageNet64 $\to$ 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	$\rho_{\mathrm{EA}}$	$\rho_{\mathrm{SA}}$
QK reversion penalty †	$\mathrm{Agr}_{111}-\mathrm{Agr}_{011}$	0.739	0.844
MLP adaptation cost (VO adapted) †	$\mathrm{Agr}_{010}-\mathrm{Agr}_{011}$	0.718	0.809
VO adaptation cost (QK adapted)	$\mathrm{Agr}_{100}-\mathrm{Agr}_{110}$	0.771	0.683
MLP-only adaptation cost	$\mathrm{Agr}_{000}-\mathrm{Agr}_{001}$	0.746	0.673
MLP adaptation cost (QK adapted)	$\mathrm{Agr}_{100}-\mathrm{Agr}_{101}$	0.644	0.673
VO reversion penalty	$\mathrm{Agr}_{111}-\mathrm{Agr}_{101}$	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(|\rho_{\mathrm{EA}}|,|\rho_{\mathrm{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	$\rho_{\mathrm{EA}}$	$\rho_{\mathrm{SA}}$
VO	$\mathrm{Agr}_{111}-\mathrm{Agr}_{101}$	0.832	0.930
QK	$\mathrm{Agr}_{111}-\mathrm{Agr}_{011}$	0.601	0.573
MLP	$\mathrm{Agr}_{111}-\mathrm{Agr}_{110}$	0.664	0.063

Appendix

Table 6: Component dependence and joint non-additivity. Spearman correlations between

\mathrm{JNA}_{\mathrm{Agr}}

and each displayed difference use 12 two-seed direction means, separately for

\mathrm{Agr}=\mathrm{EA},\mathrm{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 Transformer

Diffusion Transformers use attention-based networks to generate and transform images, video, audio, and other structured data through iterative denoising.

See all 473 papers

AReA: Adaptive Receptive-Field Attention for Efficient High-Resolution Diffusion Transformers
Oct 12, 2026Javad Rajabi, Koorosh Roohi, David B. Lindell +1Diffusion Transformer Diffusion-Based Image Generation
University of Toronto · Vector Institute
MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model
Mar 27, 2026Quan Dao, Dimitris MetaxasDiffusion Transformer Diffusion-Based Image Generation
Rutgers University
Denoising Diffusion Generative Models Secretly Calculate Attentions
Sep 1, 2026Farzan Haddadi, Leila Monfared, Ebrahim Rezaii +3Attention Control in Diffusion Models Diffusion-Based Image Generation
School of Electrical Engineering, Iran University of Science & Technology, Tehran, IRAN · Independent researcher

Oct 12, 2026J/K move · Enter open · S save

AReA: Adaptive Receptive-Field Attention for Efficient High-Resolution Diffusion Transformers

Javad Rajabi, Koorosh Roohi, David B. Lindell, Babak Taati

University of Toronto · Vector Institute

Diffusion Transformer Diffusion-Based Image Generation Attention Mechanisms Diffusion Model Inference Acceleration Diffusion-Based Denoising Image Super-Resolution

Open paper AskPDF arXiv

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.

Abstract

Diffusion Transformer

Explore similar work

AReA: Adaptive Receptive-Field Attention for Efficient High-Resolution Diffusion Transformers

MPDiT: Multi-Patch Global-to-Local Transformer Architecture For Efficient Flow Matching and Diffusion Model

Denoising Diffusion Generative Models Secretly Calculate Attentions