Mitigating the Noise Shift for Denoising Generative Models via Noise Awareness Guidance

TL;DR

The paper identifies a pervasive noise shift in denoising generative models, where intermediate states during sampling encode more noise than the pre-defined schedule, leading to sub-optimal denoising and degraded generalization. It introduces Noise Awareness Guidance (NAG) and a classifier-free variant to explicitly keep sampling aligned with the intended noise level, improving generation quality across diffusion and flow-based models. Extensive experiments on ImageNet generation and multiple supervised fine-tuning tasks demonstrate that NAG consistently mitigates noise shift and yields substantive performance gains, both standalone and in combination with classifier-free guidance. The work provides a practical, low-cost adjustment to existing denoising pipelines with broad compatibility and potential for further theoretical and empirical exploration of inference-time misalignment in generative modeling.

Abstract

Existing denoising generative models rely on solving discretized reverse-time SDEs or ODEs. In this paper, we identify a long-overlooked yet pervasive issue in this family of models: a misalignment between the pre-defined noise level and the actual noise level encoded in intermediate states during sampling. We refer to this misalignment as noise shift. Through empirical analysis, we demonstrate that noise shift is widespread in modern diffusion models and exhibits a systematic bias, leading to sub-optimal generation due to both out-of-distribution generalization and inaccurate denoising updates. To address this problem, we propose Noise Awareness Guidance (NAG), a simple yet effective correction method that explicitly steers sampling trajectories to remain consistent with the pre-defined noise schedule. We further introduce a classifier-free variant of NAG, which jointly trains a noise-conditional and a noise-unconditional model via noise-condition dropout, thereby eliminating the need for external classifiers. Extensive experiments, including ImageNet generation and various supervised fine-tuning tasks, show that NAG consistently mitigates noise shift and substantially improves the generation quality of mainstream diffusion models.

Figures (8)

• Figure 1: Empirical observation of noise shift. Denoising generative models suffer from a training–inference misalignment, where the posterior estimation during sampling tends to lean toward larger noise levels. The yellow curves indicate the estimated probability density of the posterior $p_{\phi,t}(t \mid \hat{\mathbf{x}})$ for sampled intermediate states $\hat{\mathbf{x}}$, while the orange curves indicate the posterior $p_{\phi,t}(t \mid \mathbf{x})$ for intermediate states $\mathbf{x}$ stochastically interpolated from training data $\mathbf{x}_0 \sim p_\text{data}(\mathbf{x}_0)$ on ImageNet. The (a), (b), and (c) show comparisons between posterior estimates obtained at inference and training, for prior noise levels $t=0.7$, $0.5$, and $0.3$, respectively. All density functions are estimated via kernel density estimation with 5,000 samples.
• Figure 2: Conceptual comparison of guidance behaviors based on class information and noise awareness. (a) A conceptual example of noise shift, where $\hat{\mathbf{x}}_t$ is drifted to a larger noise level by $\delta$. (b) Class-conditional guidance pushes the trajectory toward regions aligned with the class condition $c$. (c) Noise-aware guidance instead pushes $\hat{\mathbf{x}}_t$ toward the position better aligned with the intended noise level $t$ from the pre-defined prior. NAG explicitly targets the noise shift issue.
• Figure 3: FID comparison of vanilla DiTs and SiTs on ImageNet $256 \times 256$ after 80 epochs of training. Classifier-free guidance (CFG) is not used. All metrics are computed with 10K samples.
• Figure 4: Comparisons of the estimated posterior $p_\phi(t \mid \mathbf{x})$ on ImageNet $256 \times 256$ with a converged SiT/XL-2 model. (a) Noise shift across the entire sampling process, computed as the difference between the estimated posterior $p_\phi(t \mid \hat{\mathbf{x}})$ and the pre-defined prior $t$. The visualization shows that noise shift $\delta$ becomes increasingly severe as sampling progresses. (b) Noise shift measured between the estimated $p_\phi(t \mid \hat{\mathbf{x}})$ and $p_\phi(t \mid \mathbf{x})$, where $\mathbf{x}$ is generated from real data. This comparison reflects the training–inference misalignment while accounting for the inherent inaccuracy of $g_\phi$.
• Figure 5: Empirical observations of NAG mitigating the noise shift $\delta$. (a–b) Effects of NAG without interference from CFG. (c–d) Compatibility of NAG under CFG, showing that NAG addresses the noise shift directly, rather than relying on the indirect effects of CFG.