The Noise-Transmission Burden of Diffusion Network: Why Image Prediction Can Outperform Noise Prediction

Gyeonghun Kang; Pengxi Liu; Bingshuo Qian; Xiang Cheng

stat.MLOct 26, 2026

The Noise-Transmission Burden of Diffusion Network: Why Image Prediction Can Outperform Noise Prediction

Authors: Gyeonghun Kang, Pengxi Liu, Bingshuo Qian, Xiang Cheng

Organizations: Duke University, Durham, NC, USA

Abstract

Clean-image prediction ( $x$ -prediction) and noise prediction yield equivalent population-optimal denoisers, yet can perform very differently in finite diffusion networks. For noisy input $Y_t$ with noise standard deviation $σ_t$ , a noise predictor $f_{\varepsilon,t}$ induces the clean-image estimate $\widehat X_1=Y_t-σ_t f_{\varepsilon,t}(Y_t)$ . This reveals a noise-transmission burden: the predictor must preserve information needed to cancel the input corruption. Under Gaussian corruption of $d$ -dimensional data, a linear input bottleneck of rank $r<d$ imposes a denoising-error floor $σ_t^2(d-r)$ even with unrestricted downstream computation, while the corresponding optimal risk for image prediction remains bounded. EDM attenuates this bottleneck contribution at high noise; exact preconditioning cancels it at every noise level. Without an input bottleneck, assuming exact affine recovery of the noisy patch from the model representation, we characterize how the covariance spectrum of corrections beyond linear denoising determines the width required above the patch dimension. ImageNet experiments with JiT show that measured denoising losses reflect the predicted bottleneck floor and its attenuation or cancellation through preconditioning. Further widening yields larger relative denoising gains at low noise, where corrections beyond linear denoising spread across more directions, consistent with our spectral characterization of the required model width. Together, these results provide a denoising-based account of practical generation-quality gaps across prediction targets, preconditioning choices, and model widths.

Figures & tables

Prediction target	Population optimal	Network output	Induced clean image
Image ( $X_{1}$ )	$f_{X_{1},t}^{\ast}(y)$	$f_{X_{1},t}(y)$	$f_{X_{1},t}(y)$
Noise ( $\varepsilon$ )	$f_{\varepsilon,t}^{\ast}(y)$	$f_{\varepsilon,t}(y)$	$\widehat{X}_{1}^{(\varepsilon)}(y)$ , ( 2 )
General ( $T_{t}=\alpha_{t}X_{1}+\beta_{t}\varepsilon$ )	$f_{T,t}^{\ast}(y)$	$f_{T,t}(y)$	$\widehat{X}_{1}^{(T)}(y)$ , ( 3 )

Table 1: Notation for direct target prediction at a fixed interior time.

Figure 1: Architectural formulations of standard DiT. (a) Linear input map

A

and output preconditioning, ( 4 ) and ( 7 ). (b) Representation

h

, including final AdaLN, and the affine output layer in ( 11 ).

Figure 2: Generation and denoising under the same input bottleneck, at epoch

200

. (a) Guided FID-10K (

\downarrow

). (b) Per-coordinate clean-image MSE and the discarded-noise lower bounds from Corollaries 2 and 4 . Both EDM input scalings are plotted and nearly coincide. Noise prediction follows its growing floor; the EDM floor becomes negligible at high noise, where its loss approaches those for image prediction and exact preconditioning. FID and MSE use logarithmic scales; panel (b) uses time

t

, with

\sigma_{t}=(1-t)/t

Figure 3: Block-

12

noise recovery approaches but stays below the

1/6

ceiling at high noise, consistent with Theorem 1 .

Figure 4: Model width and denoiser approximation on ImageNet256. (a) Generation quality versus model width; the dotted line marks

d_{p}=768

. (b) Relative clean-image MSE reduction from width

768

. (c) Residual structure of a proxy for the population-optimal denoiser: purple circles give the

95\%

-variance direction count (left axis), and dashed green squares give residual magnitude (right axis, log scale). Both curves show the central patch; shading spans patch positions.

Figure 5: At

t=0.5

; labels:

d_{h}

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Variant	$c_{\mathrm{skip}}$	$c_{\mathrm{out}}$	$\tilde{c}_{\mathrm{skip}}$
Noise	$0$	$1$	$1$
Additive	$t$	$1$	$t$
Exact	$1/\sigma_{t}$	$-1/\sigma_{t}$	$0$

Appendix

Table 2: Fixed noise-output coefficients and the resulting clean-image skip. The learned branch receives

AX_{t}

, so

c_{\mathrm{in}}=t

Figure 6: Learned coefficients relative to exact preconditioning at epoch

200

(EMA1). The dotted line marks exact coefficients. (a) The skip ratio approaches

1

at moderate and high noise. (b) Output ratios can differ in sign and scale because the learned branch can absorb these changes.

Figure 7: Full denoising-loss comparison under the input bottleneck at epoch $200$ (EMA1), using $2000$ common validation images without guidance. Cell values give per-coordinate clean-image MSE to three significant figures; color gives the ratio to the image-prediction MSE at the same $t$ , on a logarithmic scale. Small $t$ means high noise.

Layer	$t=0.1$	$t=0.5$	$t=0.7$	$t=0.8$	$t=0.9$	$t=0.95$
0	0.16662	0.16611	0.16384	0.15873	0.13403	0.08302
1	0.16662	0.16611	0.16383	0.15872	0.13381	0.08163
2	0.16662	0.16610	0.16383	0.15871	0.13387	0.08206
3	0.16662	0.16610	0.16383	0.15871	0.13389	0.08226
4	0.16662	0.16610	0.16383	0.15872	0.13392	0.08225
5	0.16662	0.16610	0.16383	0.15871	0.13390	0.08224

Appendix

Table 3: Layerwise affine noise-recovery scores for the noise predictor at epoch

200

(EMA1;

\uparrow

), evaluated together on one H200. Scores use

2000

validation test images and finite-test noise-variance normalization. Layer

0

is the embedding; all values are below

1/6

Figure 8: ImageNet128 counterpart to Figure 4 , with patch dimension

d_{p}=192

. (a) FID-10K versus model width. (b) Relative clean-image MSE reduction from width

192

. (c) Residual structure of a proxy for the population-optimal denoiser. Colors, axes, and shading follow Figure 4 .

Figure 9: Paired denoising and affine input recovery versus model width at epoch

200

(EMA1) and

t=0.5

. (a,b) Clean-image MSE from the trained denoising head on ImageNet256 and ImageNet128. (c,d) Affine noisy-input recovery from the final-head input at the same resolutions. Blue squares and orange circles mark image and noise prediction. Vertical dotted lines mark

d_{h}=d_{p}

; horizontal dotted lines mark

100\%

input recovery.

Prediction	$d_{h}$	.05	.1	.2	.5	.8	.9	.95
Noise	768	0.54859	0.16399	0.062464	0.019595	0.0057381	0.0025318	0.0011617
Noise	896	0.1222	0.074817	0.043629	0.015972	0.0049459	0.0022827	0.0010471
Noise	1024	0.12351	0.07479	0.043377	0.015776	0.0048654	0.0022625	0.0010483
Noise	1152	0.12228	0.074653	0.043191	0.015557	0.0047105	0.0021868	0.0010165
Noise	1536	0.11825	0.073	0.042467	0.015237	0.004552	0.0020918	0.00099022
Image	128	0.11259	0.073589	0.044556	0.018345	0.0094614	0.0080807	0.0077457

Appendix

Table 4: ImageNet256: per-coordinate native clean-image MSE at epoch

200

(EMA1,

2000

validation images, no guidance;

\downarrow

). Patch dimension is

d_{p}=768

. Columns give

t

; the same images and corruptions are used for all widths.

Prediction	$d_{h}$	.05	.1	.2	.5	.8	.9	.95
Noise	192	1.957	0.48142	0.12913	0.027037	0.0069793	0.0029279	0.0012212
Noise	256	0.15822	0.10029	0.058144	0.01944	0.0055573	0.0025714	0.0011851
Noise	320	0.15593	0.099435	0.057605	0.018981	0.0053131	0.0024401	0.0011328
Noise	384	0.15521	0.098938	0.05719	0.01875	0.0051801	0.0023641	0.001104
Image	64	0.14759	0.098758	0.058917	0.021162	0.0086825	0.0068084	0.0063907
Image	128	0.14647	0.097371	0.057447	0.019596	0.0065516	0.0042498	0.0035854

Appendix

Table 5: ImageNet128: per-coordinate native clean-image MSE at epoch

200

(EMA1,

2000

validation images, no guidance;

\downarrow

). Patch dimension is

d_{p}=192

. Columns give

t

; the same images and corruptions are used for all widths.

ImageNet256
Denoiser	.05	.1	.2	.5	.8	.9	.95
JiT noise ( $d_{h}=1536$ )	0.111033	0.071357	0.042082	0.015287	0.004591	0.002104	0.000993
JiT image ( $d_{h}=768$ )	0.110087	0.070997	0.041916	0.015472	0.005118	0.002918	0.002197
ADM U-Net	0.110465	0.070728	0.040943	0.014005	0.003711	0.001485	0.000594
ImageNet128
JiT noise ( $d_{h}=384$ )	0.145839	0.096349	0.056384	0.018685	0.005189	0.002365	0.001104

Appendix

Table 6: Denoising quality of the resolution-matched ADM references: full-image per-coordinate clean-image MSE on

4000

common validation images within each resolution, without guidance. Columns give

t

; bold marks the lowest measured MSE at each time. All rows use FP32 inference and conversion.

Diffusion-Based Denoising

See all 49 papers

Residual-Stream Burden Shapes Representation Learning in Diffusion Transformers
Sep 27, 2026Tongtong Liang, Siqi Kou, Ziqiao Xi +6Diffusion Transformer Residual Learning
UC San Diego · Shanghai Jiao Tong University · Aether AI
Oracle Supervision Transfers for Hyperparameter Prediction in Model-Based Image Denoising
May 19, 2026Jianmin Liao, Lixin Shen, Yuesheng XuHyperparameter Transfer Transfer Learning
Department of Mathematics Syracuse University · Department of Mathematics & Statistics Old Dominion University
Just Align $\bm{x}$ : Aligning Predictions, Not Representations
Sep 30, 2026Yuyao Zhang, Yuwei Hu, Ziyang Mai +1Diffusion Model Alignment Diffusion Transformer
Dartmouth College

Sep 27, 2026J/K move · Enter open · S save

Residual-Stream Burden Shapes Representation Learning in Diffusion Transformers

Tongtong Liang, Siqi Kou, Ziqiao Xi, Esha Singh+5

UC San Diego · Shanghai Jiao Tong University · Aether AI

Diffusion Transformer Residual Learning Hyper-Connections Diffusion-Based Image Generation Diffusion Models

Open paper AskPDF arXiv

In diffusion-based generation, a neural network can be trained to predict the clean data, the noise, or the velocity from a noisy input. These prediction targets are interconvertible and describe the same generative process, yet plain Diffusion Transformers operating on large pixel patches succeed with clean prediction and fail with noise or velocity prediction. We argue that this asymmetry arises because noisy targets require the residual stream to preserve noise-dependent input variation through depth for the final readout, forcing subsequent layers to compute on noisy representations. A spectrally concentrated clean target imposes a lighter demand, leaving greater freedom to organize hidden representations for subsequent computation. We call this preservation requirement residual-stream burden and show how it shapes representation learning in Diffusion Transformers. Controlled experiments indicate that the exploitable structure is spectral concentration in patch space and that the bandwidth of the persistent residual state is a key resource for noisy prediction. We further show that this account is consistent with recent decoupled pixel-space architectures, whose diverse designs all reduce the residual-stream burden on the main pathway. To examine this understanding from a complementary direction, we expand and reorganize the residual-stream bandwidth directly, introducing Spatially Indexed Hyper-Connections (SiHC) that reach FID 1.71 on ImageNet

256^2

. Together, these results identify residual-stream burden as a mechanism through which prediction targets and architecture jointly shape representation learning in Diffusion Transformers.

Abstract

Diffusion-Based Denoising

Explore similar work

Residual-Stream Burden Shapes Representation Learning in Diffusion Transformers

Oracle Supervision Transfers for Hyperparameter Prediction in Model-Based Image Denoising

Just Align x\bm{x}x: Aligning Predictions, Not Representations

Just Align $\bm{x}$ : Aligning Predictions, Not Representations