Noise prediction limits diffusion model performance compared to image prediction
The Noise-Transmission Burden of Diffusion Network: Why Image Prediction Can Outperform Noise Prediction
Machine Learning
Summary
Predicting clean images rather than noise in diffusion models can work better because noise prediction requires the model to carefully undo input corruption, which is hard when the network has limited capacity. The authors show mathematically how this limitation creates an unavoidable error floor with noise prediction but not with image prediction. They also explain how design choices like preconditioning and model size can reduce or remove this problem. Their experiments on ImageNet support these findings, helping clarify why some methods yield better image generation quality.
What this means in practice
- •For machine learning engineers: Design diffusion models that avoid noise-transmission bottlenecks by choosing image prediction targets and appropriate preconditioning to improve denoising quality.
- •For computer vision researchers: Develop better diffusion-based image generation methods by understanding how model width and prediction targets affect generated image quality.
Tested on one dataset.
Authors
Gyeonghun Kang, Pengxi Liu, Bingshuo Qian, Xiang Cheng
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.