Mitigating Noise Shift in Denoising Generative Models with Noise Awareness Guidance
Authors: Jincheng Zhong
Boyuan Jiang
Xin Tao
Pengfei Wan
Kun Gai
Mingsheng Long
Venue: ICLR 2026 Poster (2026)
Keywords: diffusion models
generative models
training-inference misalignment
noise awareness
guidance
Primary area: generative models
TL;DR: We identify a noise shift problem in diffusion models, where intermediate states deviate from the pre-defined noise schedule, and propose Noise Awareness Guidance (NAG) to correct it, significantly improving generation quality.

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. Code is publicly available at https://github.com/KlingAIResearch/noise-awareness-guidance.

Published as a conference paper at ICLR 2026
MITIGATING THE NOISE SHIFT FOR DENOISING GEN- ERATIVE MODELS VIA NOISE AWARENESS GUIDANCE
Jincheng Zhong": Boyuan Jiang", Xin Tao?, Pengfei Wan", Kun Gai", Mingsheng Long" R School of Software, BNRist, Tsinghua University, China 2 *Kling Team, Kuaishou Technology. China
(zhengjinchengworkk, jiangsutx (egasuttee@mmiilll @gmaill com
{j1 giggbuuy wanpongfei suouushou....... com mi wenggghnngttinnghhinnn ent.cn
ABSTRACT
Existing denoising generative models rely on solving discretized reverse-time SDEs ODEs. In this paper, we identify 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 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- -irtr---iitribution generalization and inaccurate denoising updates. To address this problem, we propose Noise Awareness Guidance (NAG), 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 ciiindunconiiiional 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. Code publicly available now
INTRODUCTION
denoising generative models, such as diffusion models (Ho al. 2020 and flow-based models (Lip- man et al. 2023), have demonstrated remarkable scalability and achieved state-of- results across a wide range of visual generation tasks, including image synthesis (Ho et al., 2020), video generation (Ho et al., (022), and cross- modal generation (Saharia et al.) (022). The core principle of these models to progressively recover a target sample from pure noise. At each iteration, network processes an intermediate state, which consists of both signal and noise mixed in pre-defined proportions, and updates it to the next state according to the output and pre-defined coefficients.
During iterative sampling, the model is repeatedly applied and inevitably accumulates errors from multiple sources, including imperfect network approximation, discretization in numerical integration, and other stochastic factors. Recent studies have primarily focused on the discretization aspect, aiming to accelerate generation by reducing the number of denoising steps (Geng et al., 2025 Song et al. 2023 Lu et al. 2022),
designing more effective diffusion architectures to increase model capacity (Peebles & Xie 2023 Ma et al. 2024 Karras et al.) 2022). Nevertheless, accumulated errors in such complex system unavoidable. A key manifestation of these errors is that the noise level inherently encoded in intermediate states may deviate from the pre-defined schedule. This misalignment, long overlooked by the community, both widespread and rooted in the collective effect of diverse error sources. We refer to this phenomenon as noise shift, which often leads to fundamental mismatch between training and inference denoising networks.
In this work, we demonstrate that the noise shift manifests as a systematic drift toward larger noise levels t'. We conduct an empirical analysis on recent advanced diffusion models (Ma et al.) 2024)
Work done during internship at Kling Team, Kuaishou Technology
Code publicly available at: http://diibboooo//iiaaaeeeeeatcc

---

Published as a conference paper at ICLR 2026
Mean
Mean
0.320
Ps:0 1.)
Pau
pure MIBec
(8) Prior 0.70
Prior 0.50
(c) Prior 0.30
Figure 1: Empirical observation of noise shift. Denoising generative models suffer from train- ing-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 (t x) for sampled intermediate states &. while the orange curves indicate the posterior Ps. x) for intermediate states x stochastically interpolated from training data Xy ~ Pdata (Xo) on ImageNet. The (a), (b), and (c) show comparisons between posterior estimates obtained at inference and training, for prior noise levels 0.7, 0.5, and 0.3, respectively. All density functions are estimated via kernel density estimation with 5,000 samples.
for ImageNet generation. As illustrated in Figure T the noise shift issue is widespread and can be directly observed using an external posterior noise-level estimator gé. This observable noise shift ó indicates a clear mismatch: the actual noise encoded in intermediate states is not consistent with the pre-defined noise levels, exhibiting systematic tendency toward larger noise levels !' -t + 6. To quantify noise shift, we compare the posterior estimation (o x)) of intermediate states during sampling with the posterior estimation ( xx) of intermediate states from the forward process in training, along with the reference of the corresponding pre-defined prior I.
This misalignment can lead to sub-optimal results in two ways: 1) noise shift introduces out-of- distribution generalization issues, since the trained model is applied to a shifted intermediate state se (Xt+5, t) rather than the intended Sy (Xe, t). 2) noise shift causes sub-optimal denoising operations, as the next state is computed using inaccurate pre-defined coefficients.
To address this issue, we propose Noise Awareness Guidance (NAG), a novel guidance correction approach designed to mitigate the noise shift phenomenon. The key idea of NAG is to enable denoising models to recognize the inherent noise level of a given intermediate state during sampling and to generate a guidance signal that steers shifted samples back toward the accurate pre-defined noise level. However, as discussed in prior works (Ho & Salimans) 2021 Dhariwal & Nichol 202T), gradient-based guidance signals that rely external classifiers suffer from several drawbacks, including vulnerability to adversarial-like gradient manipulation, complex training pipelines, and the need for additional costly training on noisy inputs. Inspired by the success of classifier-free free guidance (CFG) (Ho & Salimans| 202I), we further propose a classifier- free variant of NAG. Instead of relying on the gradient of separately trained noise estimator, classifier-fre free NAG combines the score estimates of noise-conditional diffusion model with those of a jointly trained noise-unconditional model. This approach removes the dependency on external classifiers by applying noise-condition dropout during training.
Empirically, we show that NAG substantially alleviates the noise shift issue, consistently leading to significant improvements in the generation quality of mainstream denoising generative models. Our comprehensive evaluations are conducted across two widely used base models: DiT (Peebles Xie 2023) for diffusion models and SiT (Ma et al.) 2024] for flow-based models. To demonstrate both the effectiveness and generality of NAG, our evaluations cover two mainstream use cases of modern denoising generative models: 1) We show that NAG can be directly incorporated into DiT and SiT to improve ImageNet conditional generation, highlighting that foundation model development can benefit from our approach. 2) We conduct supervised fine-tuning experiments on small downstream datasets, verifying the effectiveness of NAG in supervised fine-tuning scenarios.
Overall, our contributions can be summarized follows:
We identify the noise shift issue, which is widespread in existing denoising generative models but has long been overlooked. Through empirical analysis with an external noise estimator ImageNet generation tasks, reveal the severity of this issue.
2

---

Published as a conference paper at ICLR 2026
We propose novel and concise approach, Noise Awareness Guidance (NAG), to mitigate the noise shift issue. We further introduce its classifier-fre variant, which can be more easily incorporated into mainstream denoising generative models.
We conduct comprehensive experiments validating the effectiveness and generality of NAG, providing strong evidence that it mitigates the noise shift issue and leads to significant improvements in both ImageNet generation and supervised fine-tuning tasks.
2 PRELIMIARY
We begin by reviewing denoising generative models under the unified framework of stochastic interpolants (Albergo al. 2023). Throughout this section, we adopt the notation of Mf et al. (2024).
Both diffusion and flow-based models can be understood as stochastic processes that gradually transform a noise sample from simple prior distributions, typically standard Gaussian e N I), into a data sample from the complex target distribution
Pdata (Xo).
Forward process. Let X ~ P be a sample from the data distribution. continuous-time stochastic interpolant over E [0, 7
We define a
X = 0g + o2C,
agp = or 1, oT = σ = 0.
where ay is monotonically decreasing and or is monotonically increasing (Lipman et al. 2023 Ma ct al. 2024). This formulation interpolates smoothly between the clean data (t = 0) and pure noise
Probability flow ODE. Given the forward process, the dynamics of x, can be equivalently described by a probability flow ordinary differential equation (PF ODE):
* v(xe; 1),
where the velocity field is given by
v(x, t) = E[ x; = x] = ái E(xu X2 = x + o Ele x, = x].
+
(3)
In practice, the velocity is parameterized by neural network vo (x, 1). trained with the objective
(0) = Ex | 1) d; - o,e
Since the ODE solution at time & matches the marginal distribution Pe(x) of x,, samples can be generated by integrating Equation 2
backward from XT (O, I) using standard ODE solvers.
Reverse-time SDE. Equivalently, the marginals Pr are consistent with the reverse-time stochas- tic differential equation (SDE):
dx; = t) wx t) dt + /une dw
(5)
where W, is reverse-time Wiener process, w: > O is diffusion coefficient, and s(x, Vx log Pi (x) is the score function. The score can be expressed either as a conditional expecta- tion equivalently in terms of the velocity field:
s(x, t) = -, Ele x; = =- -xxxxxx O,
(6)
Cqo; a.tz
Thus, data can also be generated by solving Equation 5 with the same velocity model vo(x, f).
Conditional generation Let (x y) is the density that x, is condtioned on some variable y. Pe(y x) is known, can sample from pe [x(x y) by solving conditional reverse-time SDE where the conditional score defined as:
s(x, t y) V log y) Vx log (x)) + V log (y x).
(7)
In practice, we can build a seperate neural network to model (y x) on noisy data, following
classifier guidance (Dhariwal & Nichol 2021 Song et al. 2020). Note that prly x) (x(x
3

---

Published as a conference paper at ICLR 2026
y)p; (xx), derive the classifier-free guidance sampling (Ho & Salimans, 2021). Empirically, classifier-free guidance achieves significant performance.
For simplicity, primarily consider the linear interpolant with T 1, ar 1- t, and t, following Ma ct al. (2024). Nevertheless, our analysis extends naturally to other formulations such DDPM (Ho al. 2020), which employ discretized dynamics, alternative schedules (04,0z) or
different model parameterizations.
3 NOISE SHIFT ISSUE IN THE DENOISING PROCESS
We identify misalignment between the training distribution (x(),, obtained from clean data samples Xo Pdata (Xo). and the intermediate distribution (x( encountered during the numerical solution of the SDE or ODE. Conceptually, this misalignment can be diagnosed by comparing the posterior Pt(t] x) inferred from perturbed states with the pre-defined prior p(t).
In practice, accumulated errors e from multiple cources-sure as imperfect network approximation. discretization error, and other modeling inaccuracies can be viewed perturbation applied to Xg. where X, = X + e, where e ~ NNO0 I²I). This perturbation increases the effective variance from of to of + og making the perturbed state behave if it were sampled at a shifted noise level = t 5, where
additional Gaussian
- of +
(8)
We refer to the discrepancy 8 = t as the noise shift.
Statement (Relation between noise shift and additive error). Given the forward process defined in Equation 7 consider an additive error ~ NO o²1). When the error variance o² is small, the shift admits a first-order approximation:
Voi + o2
8
ōt
where = -do/dd (See Appendix for full derivations.)
Intuitively, Statement 1 shows that accumulated errors push the effective variance in & toward a later noise level t' = + 8, where
O, causing systematic bias. For example, in the linear interpolation case og t, the shift reduces 8 v oo + o2 og, illustrating that perturbed states tend to be interpreted as noisier than intended. Although based on simplified assumptions, this analysis qualitatively captures the nature of noise shift in practical denoising processes.
Empirical analysis. To better illustrate the noise shift issue, we conduct empirical simulations on ImageNet at 256 256 resolution using the pre-trained SiT-XL/2 model, which was trained for 1,400 epochs Previous studies (Sun et al.) 2025) Stahl et al. 2000) suggest that for high-dimensional data such as images, the posterior (t x) concentrates sharply (similar to a Dirac delta), making the noise level encoded in reliably estimable. Motivated by this, train noise estimator go x) on the ImageNet 256 256 datase
Empirical comparisons between the estimated posterior distributions Pt..(1 x) are shown in Figure Consistent with Statement [] we observe that the estimated posterior distribution Pt x) (yellow curve) shifts toward larger values of the pre-defined prior t. demonstrating that the noise shift phenomenon is widespread in the denoising stage. Additionally, the orange curve shows the posterior estimation on samples generated from ImageNet through the forward process in Equation I serving evidence of the accuracy of 9& ground-truth intermediate states.
In particular, intermediate states with mid-level noise exhibit substantial systematic overestimation by gg. highlighting a clear misalignment between the training and inference distributions. Further results at more noise levels can be found in the Appendix D
The effect of noise shift While our empirical analysis is constrained by the accuracy of the noise estimator, the observed noise shift ó can still be regarded as a sufficient but not necessary condition for indicating sub-optimal behavior in the denoising stage.
Implementatiinn details of the noise estimator provided in the Appendix B.2
4

---

Published as a conference paper at ICLR 2026
P(Ko
E. logppic (x)
Fx ( (x)
(0) Noise
Guidance class conditions
(c) Guidance with noise conditions (NAG)
Figure 2: Conceptual comparison of guidance behaviors based on class information and noise awareness. (a) conceptual example noise shift, where & is drifted larger noise level by (b) Class-conditional guidance pushes the trajectory toward regions aligned with the class condition (c) Noise-aware guidance instead pushes x, toward the position better aligned with the intended noise level from the pre-defined prior. NAG explicitly targets the noise shift issue.
This pervasive noise shift affects the entire sampling trajectory in two primary ways: 1) The learned velocity field (x(x, I) suffers from out-of-distribution errors, since the model operates on perturbed intermediate states with shifted noise levels 8. If the noise-conditioned network vø(x, t) satisfies a Lipschitz condition in x, the resulting model error can be bounded by Lx|le||, where Lx is the Lipschitz constant. 2) The misalignment introduces errors in the SDE coefficients a and de during reverse-time integration. Consequently, the denoising process becomes sub-optimal under the influence of noise shift.
As discussed above, 5 can be interpreted as collection of errors originating from various sources, making unrealistic eliminate completely. Notably, reducing d to zero is not sufficient condition for generating better images. For instance, if an intermediate sample corresponds to an image that entirely out of distribution, generation will still fail due to the limited capability of the model. Nevertheless, since the existence of noise shift always induces some degree of misalignment, our qualitative findings provide valuable insights into the design of corrective methods.
NOISE AWARENESS GUIDANCE
In this section, we introduce the core concept of Noise Awareness Guidance (NAG), which directly addresses the noise shift issue. We interpret the shift 5 as the misalignment between the sampled state &, and its intended noise condition t. Inspired by conditional guidance methods (Dhariwal & Nichol 202T; (Sogg al. 2020). we propose a mechanism that explicitly steers the sampling trajectory to remain consistent with the pre-defined noise schedule. Our key insight is that by reinforcing the conditioning on I,, the posterior pe(4 x) along the reverse-time SDE (or ODE) trajectory remains closer to the pre-defined thereby mitigating the noise shift 6.
Noise awareness guidanee.
The noise-conditional score can be written as
s(x () lo () Vx (x) Vx (( x).
(10)
Analogous to Equation 7 if (t x) were available, we could sample from (x(x by solving the conditional reverse-time SDE in Equation 10 As discussed in Section 3 the posterior px x) can be reliably estimated from noisy data point x,. Intuitively, we can guide the sampling trajectory with V log (t x) the guidance signal, where 96 is the posterior estimator model in Section Since it relies on being aware of the accurate noise level encoded an intermediate state. We refer to this approach as Noise Awareness Guidance (NAG). As the gradient V log 9s x) is provided by an external posterior estimator Jg. we call this formulation classifier-based NAG.
Classifier-free noise awareness guidance. Despite its effectiveness, classifier-based NAG inherits the drawbacks of classifier guidance (Dhariwal Nichol, 2021 Song et al. 2020), including the high computational cost of training external posterior estimator for t. increased pipeline complexity, and the risk of adversarial- behavior in explicit classifiers. To address these issues, we extend the idea of classifier-free guidance (CFG) (Ho & Salimans 2021 to NAG.
5

---

Published as a conference paper at ICLR 2026
Noting that Pt(t x) (x ()), we can utilize score mixture to approximate the gradient of an implicit noise predictor as
S 6) (Wing + s) [t) Wnag s(x),
(11)
where Wnag is the guidance parameter for NAG. Importantly, modern denoising models already accept the noise level t along with the intermediate state x, inherently defining the conditional score s(x tt.. Thus, we only need access to the unconditional score s(x), without explicitly training a separate noise-level predictor.
implement NAG, we follow the training strategy CFG: during training, the noise condition / is randomly dropped with fixed probability, allowing the model to share weights between conditional and unconditional objectives.
Discussion and relation to CFG. The mechanism of NAG can be intuitively understood in analogy to CFG. From the perspective of conditional generation, sampling without NAG corresponds to relying solely the conditional score model. By strengthening the conditioning on t, NAG guides the trajectory toward lower- remperrtm regions where the model produces higher-confident samples, ensuring that each intermediate state remains aligned with its intended noise level.
As illustrated in Figure 2 the noise-level conditioning axis introduced by NAG is orthogonal to the conditional axis of CFG, providing complementary control over the sampling process. It is worth noting that because the noise shift arises from various sources, CFG empirically mitigates it to some extent, it biases sampling toward lower-temperature regions where models are more confident. However, compared to this indirect effect of CFG, NAG directly targets the reduction of ó and thereby constructs improved sampling trajectories. Figure visualizes the mitigating effect on noise shift by different methods.
5 EXPERIMENTS
In this section, we present a comprehensive empirical analysis to demonstrate the effectiveness and generality of NAG. Our study considers two settings: (1) standard ImageNet generation benchmarks (Section 5. and (2) supervised fine-tuning off-the-shelff models on small, fine-grained datasets (Section 5.2 These experiments provide evidence of NAG's compatibility with two widely used scenarios: large-scale foundation model training and supervised fine-tuning. Section 5.3 presents more discussion empirical analysis of noise shift 8.
5.1 NAG FOR IMAGENET GENERATION
Implementation details. Our experiments are conducted on two representative variants of denoising
generative models: DiT's (Peebles & Xie 2023) for diffusion-basel models and SiT's (Ma al. 2024) for flow-hased models. We faithfully follow the experimental setups described in the DiT and SIT papers, unless otherwise specified. All experiments are performed resolution of 256 256 (denoted as ImageNet 256 x 256), where 32 x 32 x 4 latent vectors are obtained using the pre-trained Stable Diffusion VAE tokenizer (Rombach et al. 2022). For model configurations, we adopt the S/2, B/2, L/2, and XL/2 variants introduced in the DiT and SiT papers (Peebles & Xie 2023 Ma al. 2024), all of which process inputs with a patch size of For experiments trained from random initialization, we train for 80 epochs and apply a 10% dropout probability on the noise conditions. Due to computational limitations, evaluations on fully converged XI.2 models are instead conducted by fine-tuning for additional 10 epochs off-the-shelf checkpoints pre-trained for 1,400 epochs with 20% noise dropout. Additional experimental details are provided in Appendix B
Evaluation. For experiments with DIT, we follow the default setup using 250 DDPM sampling steps (Peblles & Xie 2023). For SiT, consistent with its original setup, we always adopt the SDE-Euler-Maruyee
sampler with 250 sampling steps (Ma et al. 2024). For experiments across different architectures of DiTs and SiTs, we report the Fréchet Inception Distance (FID) (Heusel al. 2017) computed with 10,000 samples. For converged results, to enable direct comparison with the original papers, we report FID, precision (Prec.), and recall (Rec.) (Kykkäänniemi et al.) 2019) computed with 50,000 samples by default.
Comparison. Figure presents the results of training DřTs and SiTs from scratch across various architectures. The results show that NAG consistently brings substantial improvements over the

---

Published as a conference paper at ICLR 2026
Table 1: Converged comparsions on ImageNet 256 256 with DIT-XL/2 and SiT-XL/2. We
fine-tune off- afflth DiT-XL/2 and SiT-XL/2 checkpoints for an additional 10 epochs to support NAG sampling, with and without classifier free guidance (CFG), following the setup the original papers (Peebles & Xie. 2023
Ma et al. 2024 All metrics are reported on 50k generated images.
Training Epoches
Generation w/o CFG
Generation w/ CFG
Model
FID
Prec.
Ree.
FID
Prec.
Rec.
DiT-XI.2 (Pechles & Xie 2023
1400
9.62
0.67
0.67
2.27
0.83
0.57
+AG (ours)
10+((1000))
2.59
0.79
0.60
2.14
0.80
0.61
SiT-XL/2 (Ma et al. 2024
1400
8,61
0.68
0.67
2.06
0.82
0.59
+ NAG (ours)
10++(1400")
2.26
0.75
0.66
1.72
0.77
0.66
700-
with NAG
NAG
615
42.43
204
26.0
15:27
20:24
05
20.00
16.00
oF
sh
(0)
(b) Model R/2
(c) Model 1./2
(d) Model
Figure 3: FID comparison of vanilla DiTs and SiTs on ImageNet 256 x 256 after 80 epochs of training. Classifier-free free guidance (CFG) is not used. All metrics computed with 10K samples.
baselines. An interesting observation is that DiTs benefit more from NAG than SiTs when trained for 80 epochs. This may arise from the different training schedules: the DDPM-style setup used in DiTs could lead to better training of the noise- uniiis-uncniiiionall branch, thereby providing more accurate guidance direction for NAG. Notably, for extensively pre-trained models, is sufficient to fine-tune only the noise-unconditional branch at small fraction of the original cost (e.g., 10% additional epochs, approximately 0.7% of the full 1,400-epoch pre-training cost) to enable the model to apply NAG. Remarkably, using NAG alone allows the model to achieve generation quality close to that of a CFG-guided model. Morcover, when combined with CFG, NAG continues to provide additional improvements, demonstrating that its mechanism is complementary and orthogonal to CFG.
5.2 NAG FOR SUPERVISED FINE-TUNING
Implementation Details. Supervised fine-tuning of an off-the-shelf pre-trained checkpoint to new domain is a fundamental task in generative modeling. To further demonstrate the general effectiveness of NAG, conduct supervised fine-tuning evaluations following the setups in Zhong et al. (2025) 2024). Specifically, we evaluate NAG on fine-tuning DIT-XL/2 across seven well- established fine-grained downstream datasets: Food101 (Bossard et al. 2014), SUN397 (Xiao et al. 2010), DF20-Mini (Picek et al., 2022), Caltech101 (Griffin et al. 2007)) ). 20--200-201 (Wah et al. 2011), ArtBench- 10 (Liao et al. 2022). and Stanford Cars (Krause et al., 2013). We fine-tune for
24,000 steps with a batch size of 32 at 256 x 256 resolution for each task. The compared basclines include vanilla generation, generation with classifier-free guidance (CFG). and Domain Guidance (DoG) (Zhong et al. 2025). Notably, DoG is guidance method specifically designed for fine-tuning scenarios. To demonstrate both the fundamental effect and generality of NAG, we directly apply on top of these baselines without any modifications, except for introducing noise-dropout training to support the noise-unconditional branch. Detailed implementation information is provided in Appendix
B
Followling the setup in (Zhong et al., 2025), all results are generated with 50 DDPM sampling steps. and the FIDs are computed with 10,000 samples.
Evaluations.
The FIID comparisons across various fine-tuning tasks are summarized in Table The results indicate that NAG is highly general and exhibits strong compatibility across different basclines,
Results.
7

---

Published as a conference paper at ICLR 2026
Table FID Comparisons on fine-tuning tasks with pre-trained DIT-XL-2-2563:
Dataset
CUB
Stanford Car
Average
Food
SUN
Caltech
DF-20M ArtBench
Method
Bird
FID
Fine-tuning (w/o CFG)
16.04
21.41
31.34
9.81
11.29
17.92
22.76
18.65
NAG (ours)
11.18
14.95
24.32
5.68
5.92
14.79
19.22
13.72
Fine-tuning (with CFG)
10.93
14.13
23.84
5.37
6,32
15.29
19.94
13.69
NAG (ours)
5.78
8,81
21.87
3.52
3.91
12.55
15.69
10,31
-uninn (with DoG)
9.25
11.69
23.05
3.52
4.38
12.22
16.76
11.55
NAG (ours)
6.45
8.24
21.88
3.41
4.21
11.38
14.80
10.05
wis NAG wio CFG
NAG EPG
1.33
1.30
0.30
1.0
020
(o) ((1
Estimateddp.
Figure 4: Comparisons of the estimated posterior Ps(t x) on ImageNet 256 x 256 with converged SiT/XL-22 model.
(a) Noise shift across the entire sampling process, computed as the difference between the estimated posterior Po(t X) and the pre-defined prior t. The visualization shows that noise shift 8 becomes increasingly severe as sampling progresses. (b) Noise shift measured between the estimated (t x) and pt(t x), where x is generated from real data. This comparison reflects the training-inference inference misalignment while accounting for the inherent inaccuracy of gø.
benchmarks, and guidance approaches. Consistent with the ImageNet results, NAG alone achieves performance comparable to sampling with CFG. Furthermore, Table shsww that both CFG- guided sampling and DoG-guided sampling can be substantially improved by NAG. This broad compatibility highlights that the noise shift issue is indeed widespread in denoising saaaseinng-baee generation, and that NAG, by directly addressing this issue, can consistently improve generation quality across various baselines. Notably, Domain Guidance (DoG) (Zhong et al. 2025). CFG variant specifically designed for supervised fine-tuning, also benefits from NAG, with significant improvements observed in Table
5.3 EMPIRICAL OBSERVATIONS OF NOISE SHIFT WITH NAG
We present a detailed empirical analysis based on the estimator Jø. as an expansion beyond Section
As the sampling process progresses, the noise shift can be divided into stages. In the first stage, the shift increases steadily until it reaches a threshold (e.g., when the signal-to-noise ratio is around 1). In the second stage, the shift plateaus, remaining relatively stable the actual noise level decreases from 0.5 to O. When intermediate states approach the data distribution at very low noise levels, the estimated noise shift relative the pre-defined prior tends to be overestimated. This occurs because Is applied to intermediate states x generated from real data suffers from larger estimation errors due its limited capability in this regime. This overestimate can be viewed in Figure 4 a) and released by mean normalization in Figure b).
As shown in Figure NAG primarily influences the sampling process when the signal-to- -oniie ratio is larger than I (roughly R 0.5), effectively reducing the noise shift in this range. In contrast, its effect less pronounced in the early denoising stage, where the signal-to-noise ratio is low. Figure 5 further illustrates that NAG shifts the density of intermediate states toward the posterior Pt. estimated from real data, and hence closer to the pre-defined prior t.
Classifier-free guidance (CFG) (Ho & Salimans 1021 is known steer the sampling trajectory toward low- -owterpperature regions associated with the target class, thereby producing hugletyquality samples within high-com confideeonceeeeeeeee regions. This can be interpreted Since noise shift d reflects the accumulation errors from multiple sources, CFG also reduces noise shift to some extent (as observed in Figure [ b)). However, its effect remains indirect and limited,
reduction of model fitting errors.
8

---

Published as a conference paper at ICLR 2026
CIN
210
Pricct
Phort 2015
Phort=
(d) Priort 0.15
Figure 5: Empirical observations of NAG mitigating the noise shift 8. (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.
CFG primarily mitigates errors along the class dlass-conditiona dimension. In contrast, Figures and c--d demonstrate that NAG can be directly applied on top of CFG-guided models, substantially reducing the remaining noise shift.
is important to clarify that eliminating the estimated noise shift 6 is not a sufficient condition for achieving optimal generation, since potential pitfalls may lie in the imperfect accuracy of the noise estimator or in other complex factors. Nevertheless, the presence of a distinguishable noise shift during sampling is sufficient condition for sub-optimal generation. This observation motivates us to address the noise shift issue directly.
5.4 ABLATION ANALYSIS
The main ablation analysis results are shown in Figure
SIT-L-2
XIXXL-2
25 20
X2
22
29
of weight Whin on
weight Wrog
of NAG
wi WINAG
200
3000
X020
14
w1i0
140 sampplingstpss
(0) Pa (+ -Ra=(tlk) across
Figure 6: Sensitivity analysis of Noise Awareness Guidance (NAG).
NAG weight Wnag- We conduct analysis how Whig influences the sampling results with NAG, shown in Figure 6 bb.... As guidance-based technique, NAG shows a similar effect to CFG; namely, with increasing Wnag. the diversity will be sacrificed, which can be seen as a form temperature controlling in denoising generation, even though NAG does not incorporate any class information.
Can better fitness of the denoising model reduce noise shift? An intuitive question that, since model prediction errors contribute part of the noise shift, stronger networks, such as increasing model size or training for many more iterations? Unfortunately, Figure ( shows that the empirical noise shift always falls within relatively stable range and cannot be significantly reduced by increasing the fitness the model. Note that SiT-XL/2" is widely believed to be a convergent model with 1400 epochs of training on ImageNet256, yet it still suffers from the noise shift issue.
mitigate the noise shift by building

---

Published as a conference paper at ICLR 2026
Sampling steps. Although NAG does not directly optimize the sampling steps, Figure still shows that we can effectively save sampling steps after applying NAG. Notably, although NAG introduces a double forward pass to compute the guidance signal. the sampling steps (50 steps compared to 250 steps) to achieve comparable sampling results.
only around one-fifth of
6 RELATED WORK
Denoising generative models. Denoising generative models, including diffusion models and flow- based models (Ho et al. 2020; Song & Ermon 2019 Song et al.) 2020 Lipmann et al. 2023), generate
high quality samples from pure noise through an iterative denoising process. Recent progress in this field has primarily focused on noise schedules (Nichol & Dhariwal] | Karras et al. (2022), training
objectives (Salimans & Ho 2021), and model architectures (Peebles & Xie [023] Ma et al. 2024).
which aim to reduce approximation errors caused by limited model capacity. Another important direction is the development of faster denoising methods with fewer iterative steps, such high order solvers (Bao et al. 2022 Lu et al. 2022) and improved interval modeling (Frans et al. 2025 Geng al. 2023 Song al. (023). These works primarily address numerical errors introduced by discretized integration. In contrast, most prior studies have focused on eliminating specific sources of error. In this paper, we instead highlight pervasive issue, namely noise shift, and demonstrate how addressing it alleviates the persistent sub optimality in the generation process.
Training -nfferenee misalignment. Training-inferene misalignment is fundamental challenge that has accompanied the development of generative modeling. Modern generative models suffer from this issue severely. particularly due to their multistep sampling nature. In autoregressive models, each token is generated conditioned on previous model predictions, and accumulated errors propagate throughout the sampling chain, a 2015, Ranzato et al.) 2015; Schmidi 2019) Zhang et al., 2025). Recent works have investigated analogous misalignment phenomena in diffusion models (Ning et al. 2023a Specifically, Ning et al. (2033)) aims to reduce generalization error under misalignment by improving the model's Lipschitz continuity, while Li et al. 2023) manipulates the sampling schedule mitigate the mismatch. Ning eet al. 2023a) further proposes a training-frr frre epsilon scaling method that rescales the removed noise intermediate states. Compared these works, our formulation of noise shift provides an empirical and objective perspective for quantifying the misalignment, and it naturally motivates the design of NAG. Abuudweiili al. (2025) also adopts noise-level perspective identify misalignment, however, their analysis remains primarily theoretical, lacks clear empirical quantification.
problem referred to as exposure bias (Bengio et al.
Guidance techniques for condition generations. Guidance has been shown to play a central role in conditional generation (Dhariwal & Nichol 2021 Ho & Salimans (02T), significantly improving alignment between generated samples and conditioning information. More recently, Kynkäänniemi al. (2024); Karras al. (2024) proposed techniques to further improve the practical effectiveness of classifier free guidance. Our proposed Noise Awareness Guidance also falls into this category. To the best of our knowledge, the first method to explicitly use the noise level itself as a guidance signal, directly enhancing alignment with the intended noise condition.
7 CONCLUSION
This paper presents a novel perspective by observing the behavior of the posterior noise level pr(1 x). and finds the noise shift issue that the empirically estimated posterior noise level Po (L &) has tendency toward a larger noise level. We analyze that the noise shift issue is a manifestation caused by a collection of errors from various sources and is widespread in the current denoising sampling process, and performing iterative denoising sampling under noise shifts leads to sub-optimal generations. We further provide a noise awareness guidance approach and its classifier-free variants to directly relieve the noise shift issue and achieve significant improvement by reducing the noise shift gap. We hope that our work will attract researchers to pay attention to the widespread training and inference misalignment in denoising generation and facilitate many possible future research directions, including theoretical empirical analysis on the noise shift issue, building generative models that are robust to inference shift in sampling stages, exploring the boundary high-quality generation, or faster sampling.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of our results, we provide comprehensive details, including model config- urations, key hyperparameters, fine-tuning strategies, and the checkpoints used in Section and Ap- pendix [ We believe that, with these details, the main results can be reproduced with only few lines of modification to the official DiT and SiT codebases. For the posterior estimator Ig. which requires additional modifications. Code is publicly available at http://github.com//iingg. IieeannnnnsAssssccoooiii-esss wwareness-guidanee
ACKNOWLEDGMENTS
This work was supported by Kuaishou Technology, Natural Science Foundation China (U2342217), Fundamental and Interdisciplinary Disciplines Breakthrough Plan of the Ministry of Education of China UYB2025XDXM803) Beijing Scholar Program, National Engineering Research Center for Big Data Software.
REFERENCES
Abulikemu Abuduweili, Chenyang Yuan, Changliu Liu, and Frank Permenter. Enhancing sample generation of diffusion models using noise level correction. Transactions on Machine Learning Research, 2025.
Michael Albergo, Nicholas Boffi, and Eric Sunden-Eijnded Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797, 2023.
Fan Bao, Chongxuan Li, Jun Zhu, and Bo Zhang. Analytic-DPM: an analytic estimate of the optimal reverse variance in diffusion probabilistic models. In ICLR, 2022.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015.
Lukas Bossard, Matthieu Guillaumin, and Luc Van Gool. Food- 1old-in0g discriminative compo- nents with random forests. In ECCV, 2014.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In NeurIPS, 2021.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural networks, 107:3-11, 2018.
Kevin Frans, Danijar Hafner, Sergey Levine, and Picter Abbeel. One step diffusion via shortcut models. In ICLR, 2025.
Zhengyang Geng, Mingyang Deng, Xingjian Bai, J Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling. arXiv preprint arXiv:2505.13441 2025.
Gregory Griffin, Alex Holub, and Pietro Perona. Caltech-256 object category dataset. 2007.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. NeurlPS, 2017.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. In NeurIPS, 2021.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS. 2020.
Jonathan Ho. Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. NeurIPS, 2022.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion- based generative models. In NeurIPS, 2022.
11

---

Published as a conference paper at ICLR 2026
Tero Karras, Miika Aittala, Tuomas Kynkäänniemi, Jaakko Lehtinen, Timo Aila, and Samuli Laine. Guiding a diffusion model with a bad version of itself. In NeurlPS, 2024.
Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3d object representations for fine-grained categorization. In ICCV, 2013.
Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. NeurIPS. 2019.
Tuomas Kynkäänniemi, Miika Aittala, Tero Karras, Samuli Laine, Timo Aila, and Jaakko Lehtinen. Applying guidance in limited interval improves sample and distribution quality in diffusion models. In NeurIPS, 2024.
Mingxiao Li, Tingyu Qu, Ruicong Yao, Wei Sun, and Marie-Franche Moens. posure bias in diffusion models through sampling with shifted time steps. arXiv:2305.15583. 2023.
Alleviating arXiv preprint
Peiyuan Liao, Xiuyu Li, Xihui Liu, and Kurt Keutzer. The artbench dataset: Benchmarking generative models with artworks. arXiv preprint arXiv:2206.11400. 2022.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. In ICLR, 2023.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: fast ode solver for diffusion probabilistic model sampling in around 10 steps. NeurlPS, 2022.
Nanye Ma. Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden- and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In ECCV, 2024.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In ICML, 2021.
Mang Ning, Mingxiao Li, Jianlin Su, Albert Ali Salah, and Itir Onal Ertugrul. Elucidating the exposure bias in diffusion models. arXiv preprint arXiv:2308.15321, 2023a.
Mang Ning, Enver Sangineto, Angelo Porrello, Simone Calderara, and Rita Cucchiara. Input perturbation reduces exposure bias in diffusion models. arXiv preprint arXiv:2301.1100 2023b.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In ICCV, 2023.
Lukáš Picek, Milan Šule, Jiff Matas, Thomas S Jeppesen, Jacob Heilmann-Clane Clausen, Thomas Læsspe, and Tobias Frøslev. Danish fungi 2020-not just another image recognition dataset. In WACV, 2022.
Mare' Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojciech Zaremba. Sequence level training with recurrent neural networks. arXiv preprint arxi: 151 1.6632,, 2015.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In CVPR, 2022.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. NeurIPS, 2022.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. In ICLR, 2021.
Florian Schmidt. Generalization in generation: A closer look at exposure bias. arXiv preprint arXiw:1910.00292, 2019.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In NeurIPS, 2019.
12

---

Published as a conference paper at ICLR 2026
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ICLR, 2020.
Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In ICML, 2023.
Volker Stahl, Alexander Fischer, and Rolf Bippus. Quantile based noise estimation for spectral subtraction and wiener filterin. In CASSP, 2000.
Qiao Sun, Zhicheng Jiang, Hanhong Zhao, and Kaiming He. noise conditioning necessary for denoising generative models? In ICML, 2025.
Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie. The caltech-ucsd birds-200-2011 dataset. 2011.
Jianxiong Xiao, James Hays, Krista A Ehinger, Aude Oliva, and Antonio Torralba. Sun database: Large-scale scene recognition from abbey zoo. CVPR, 2010.
Junyu Zhang, Daochang Liu, Eunbyung Park, Shichao Zhang, and Chang Xu. Anti-exposure bias in diffusion models. In ICLR, 2025. URL http:://pppneeeeeeeeeeeerrreeeee
MtDo7rWok1
Jincheng Zhong, Xingzhuo Guo, Jiaxiang Dong, and Mingsheng Long. Diffusion tuning: Transferring diffusion models via chain of forgetting. In NeurIPS, 2024.
Jincheng Zhong. XiangCheng Zhang, Jianmin Wang, and Mingsheng Long. Domain guidance: simple transfer approach for pre-trained diffusion model. In ICLR, 2025.
13

---

Published as a conference paper at ICLR 2026
A DERIVATION OF STATEMENT I
We derive the expected noise shift 6 in the presence of additive Gaussian error.
Recall that the forward process is defined for noise level E [0, T
X X + oiC, where e~ (0,
(12)
Influence of error e. Consider an intermediate state perturbed by additive error:
X4 +
(13)
where RD assumed follow a zero-mean Gaussian distribution with unknown variance, ~ o2)
The perturbed state can be rewritten
X &x + (o +e)
(14)
Since and e are independent zero-mean Gaussians, their weighted sum is also Gaussian with variance
Var(orc + o, + I (0 + +
(15)
Thus, the distribution of &, is
x ~ N Naxxo (o +
(16)
The perturbed state &, can be expressed terms of the initial data Xp:
x; (a:Xo + ore) + a: + (ate + +
(17)
Definition of noise shift. This distribution coincides with that of intermediate state from the original forward process but evaluated at shifted noise level t' = t + By definition, ó satisfies
o 8 = of + o
(18)
and the noise shift is defined
= 1' - 1.
(19)
First-order approximation. Assume that or differentiable in and that the error variance of is small, so that 8 is also small. A first-order Taylor expansion of ae+ó around t gives
T o 5,
(20)
where o;
de.
de
By construction, o+tS Voi + + Substituting yields
at + àz ô R oo +o
(21)
Result. Solving for gives the following approximation for the noise shift:
8 o
(22)
de
B IMLEMENTATIONS
All experiments are conducted in PyTorch, based on the official DiT (Peebles & Xie 2023 and
SiT (Ma et al. 2024) codebases.
B. IMPLEMENTATION TO MAIN RESULTS
Architecture configurations. We follow the transformer architectures defined DiT, using four different configurations for various model sizes: Small (S), Base (B), Large (L), and XLarge (XL). All models employ patch size of 2, and latent states are obtained using the pre- trained Stable Diffusion tokenizer (Rombach al. 2022 The detailed model architectures provided in Table 3
14

---

Published as a conference paper at ICLR 2026
Table 3: Contigurations on DiTs and SiTs.
configs
S/2
B/2
L/2
XL/2
params (M)
33
130
458
676
FLOPs (G)
6.0
23.0
80.7
118.6
depth
12
12
28
hidden dim
384
768
1024
1152
heads
6
12
16
16
patch size
2x2
2x2
x 2
x 2
latent encoder
SD- SD-VA (Robacch al.
2022
Sampler. For DiT, we directly adopt the DDPM sampler from the official implementation For SiT, use the Maler-Maruyama sampler from its official implementation with the default setting Wt = in Equation 5 and the final step size set to 0.04.
Guidance weights. For all baselines with CFG, we keep the setting consistent with the original results, using Wefg 1.5. For all results of NAG without CFG, Wnag 3.0 by default. For NAG combined with CFG, we set Wafg 1.2 and Wnag 2.0 by default.
Training configurations. We retain most training configurations from DiT and SiT (Peebles & Xie, 2023 Ma et al. (2024), without modifying decay schedules, warmup schedules, AdamW W
hyperparameters, or applying additional data augmentation or gradient clipping. All results are reported using an exponential moving average (EMA) of model weights with decay of 0.9999. Our training setup includes two scenarios on ImageNet: (1) training from random initialization (Figure 3
and (2) fine-tuning off-the-shelf pre-trained models (1400 epochs) with an unconditional noise branch (Table []. Detailed configurations are summarized in Table 4
Table 4: Training Configurations on ImageNet
configs
from scratch (Figure
fine-tuning (Table
training iterations batch size
400K
50K
256
256
optimizer
AdamW
Adam W (0.9,0.999)
((1, (g)
(0.9,0.999)
noise dropout
10%
20%
learning rate
1 x 10
1 x 10
Fine-tuning with noise condition dropout on ImageNet. Compared to training from scratch, fine-tuning requires more careful handling to avoid catastrophic forgetting of learned generative capability. Following the strategy for glass-unconditional inputs, we introduce a pseudo noise level (i.e., 1001 for DiT, 1.001 for SiT) that remains consistent across imputs, rather than discarding noise embeddings directly. In addition, we reduce the learning rate to one tenth of the original value 10-5 instead of 10-4) and double the noise dropout ratio to 20%. When training from scratch, the choice of unconditional implementation has only minor effect on training dynamics.
Fine-tuning on new datasets. We strictly follow the setup in Domain Guidance (Zhong et al. 2025), using a constant learning rate of x 10-4 and batch size of 32 with the AdamW optimizer for 24,000 iterations across all datasets. For NAG, we apply 10% noise dropout.
FID calculation. For fair comparison across benchmarks, we strictly follow the FID calculation protocol used in the original implementation of each task. For ImageNet generation, compute
https:/ggthh..cm/ttcchoorrreeee https//ggthuh..mm/villiissmmaaa
15

---

Published as a conference paper at ICLR 2026
FID scores between generated images (10K or 50K) and all available real images in the ImageNet training set, using ADM's TensorFlow evaluation suite") (Dhariwal & Nichol 2021). For fine-tuning experiments downstream datasets, we observe small performance variations between different FID implementations. To ensure consistency with results reported in (Zhong et al.) 2025), we compute FID scores using Py'Torch implementation" comparing 10K generated images against all available images in the test set for each downstream task.
B.2 IMPLEMENTATION OF EMPIRICAL POSTERIOR ESTIMATOR Ié-
To empirically identify the noise shift issue, we rely on an external posterior estimator Gá. Here we describe the construction of the estimator de used in Section 3 and Section [..3) Code is publicly available at https://ithuub.com hwtpe//in///////iiiii www.ga/earchhhooss-aaaaaaaaa
guidance.
To reduce computational costs, we fine-tune the existing SiT-XL/2 checkpoint (the same model used for ImageNet generation) by replacing its final layer with noise level regressor. The regressor is implemented as two-layer MLP applied to the globally averaged token: the first layer projects the hidden state from 1152 576 dimensions with SiLU activation (Elfwing et al. 2018), and the second layer outputs the predicted noise level.
We inherit the training pipeline and hyperparameters from the noise-condition fine-tuning setup on ImageNet described in Section B.1 including a learning rate of 1 x 10-5, the same batch size, AdamW optimizer settings, and identical data preprocessing. The key difference is that the noise level is used the prediction target rather than as input condition. The model parameters are optimized by minimizing the L2 loss between the predicted and true noise levels, with the noise condition input masked by pseudo condition (set to 1.001 in practice).
The posterior model operates in the latent space obtained from the SD-VAE (Rombach et al., 2022) ) avoiding the need to transform noisy latent states back to image space. We train Is on ImageNet 256 x 256 for 40 epochs (approximately 200K iterations), reaching a training loss of 0.0002. No EMA is applied to gø-
All probability density functions in this paper are plotted using kernel density estimation (KDE) with 5,000 samples.
The samples are constructed in two steps. First, we randomly sample 5,000 images from ImageNet and generate 5,000 noise samples. We then linearly interpolate the images and noise following the linear schedule, producing 5,000 forward trajectories in which intermediate states share the same clean data point X. and noise e. Second, we generate 5,000 reverse trajectories using the Fuler-Maruyama SDE solver with 20 steps, incorporating the same class information, and save all intermediate states. both cases, intermediate states within the same trajectory are tied to the same clean data point and noise. Finally, we compute the densities via KDE for samples associated with the same prior and the same generation process.
C COMPARISON WITH REPRESENTATIVE BASELINES ON EXPOSURE BIAS
In this section, we compare NAG with several prior works that also aim to address training-inference misalignment. The baselines include two training-free free approaches, Epsilon Scaling (ES) (Ning et al. 2023a) and Time Shift (TS) (Li et al.) 2023), as well as one training-intensive approach, Input Perturbation (IP) (Ning et al. 2023b). The results are provided Table 5 NAG significantly outperforms all prior works, especially under the w/o CFG setup.
All these prior works rely on the assumption that model predictions during sampling tend exhibit larger variance compared to those during training, and they apply manually designed matching strategies to compensate for this. ES scales the model predictions by a fixed hyperparameter, which may not be effective across all sampled states. TS searches for nearby time steps that better align with the inherent variance of the intermediate states, but this approach is sensitive to the number of
16

---

Published as a conference paper at ICLR 2026
Table Full comparison on ImageNet 256 256 with DIT-XL/2 and SiT-XL/2. We fine-tune off-the-shelf DiT-XL/2 and SiT-XL/2 checkpoints for 10 additional epochs to support NAG sampling. with and without classifier-fre guidance (CFG), following the original setups (Peebles & Xie 2023 Ma et al. 2024). All metrics are reported on 50k generated images.
Training Epochs
Generation w/o CFG
Generation w/ CFG
Model
FID
Prec.
Rec.
FID
Prec.
Rec.
DIT-XL/2 (Peebles & Xie
2023)
1400
9.62
0.67
0.67
2.27
0.83
0.57
ES
(Ning et al. 2023a
1400
12.25
0.63
0.69
2.20
0.79
0.60
+ TS
(L1 ct al. 2023
1400
13.14
0.64
0.62
3.70
0.75
0.60
IP (Ning et al. 2023b
10+((1000*)
10.20
0.63
0.69
2.19
0.79
0.61
+ NAG (ours)
10+((1400))
2.55
0.79
0.60
2.14
0.80
0.61
SiT-XL/2 (Ma et al.
2024)
1400
8.61
0.68
0.67
2.06
0.82
0.59
ES (Ning et 2023a
1400
8.70
0.67
0.68
1.96
0.81
0.61
TS (Li et al. 2023
1400
8.65
0.67
0.68
1.94
0.81
0.61
IP (Ning 2023b
10+(1400")
8.06
0.68
0.67
1.95
0.81
0.59
NAG (ours)
10+{(1400))
2.26
0.75
0.66
1.72
0.77
0.66
sampling steps and the search window. IP introduces Lipschitz regularization by injecting additional noise into training samples, aiming to reduce the generalization error when becomes misaligned with the conditioning input However, accumulated errors arise from multiple factors, not solely from model prediction errors. We believe that NAG benefits from avoiding explicit assumptions about the error structure and from being naturally adaptive throughout the sampling process.
The implementations for DiT-XL/2 of these prior works are based on their official repositories. Since these methods are not directly available for flow models, the results on SiT-XL/2 are obtained by reproducing their basic discrete implementations and extending them continuous setups in our codebase.
D MORE VISUALIZATION RESULTS WITH KERNEL DENSITY ESTIMATION
In this section, we provide the full probability density results of the estimated posterior t, as an extension of Figure and Figure
E THE USE OF LARGE LANGUAGE MODELS (LLMs)
We large language models (LLMs) only as grammar checkers during paper writing.
17

---

Published as a conference paper at ICLR 2026
Mean
Mean
Mese
0.500
0.83
5.505
0.553
0.944
5.794
0,500
0.850
5.500
Mean
Mean
Moan
0.512
0313
0.526
0.559
0.455
5.407
0,500
0550
9.500
Mean
Mewn
Mean
Mean
0250
Mean
Mean
Mean
0.185
0.142
5.103
0.143
0.102
5.06)
0.150
0.100
0.50
Figure 7: More visualization of noise shift. The yellow curves indicate the estimated probability
density of the posterior (t ()) for sampled intermediate states &, while the orange curves indicate the posterior (t. x) for intermediate states x stochastically interpolated from training data Xo ~ Panaa (X)) on ImageNet. The black indicator is the pre-defined f.
18

---

Published as a conference paper at ICLR 2026
Mcan
Mcan
Mean
0.333
0202
0.232
0.314
02299
0.200
0.252
0.400
5.289
0.500
0.200
0.400
Moan
Nean
Mean
Mean
Mean
0.477
0,472
Mean
Mean
Mean
0.329
2770
0.227
0.311
226
0.204
0.252
240
0.180
0.500
250
5.200
0.180
0.150
Figure 8: Additional visualization of how NAG mitigates noise shift. The yellow curves represent the estimated probability density of the posterior P x) for sampled intermediate states *. The blue curve shows the density influenced by CFG, while the pale gold curve highlights the mitigating effect of NAG. The orange curves correspond to the posterior Pt x) for intermediate states stochastically interpolated from training data Xo Peata (Xo) on ImageNet. The black indicator denotes the pre-defined t.
19


Reviews:
[{"review_id": "vIBXJ4pIth", "reviewer": "Reviewer_JnRR", "summary": "This paper aims to solve a train-inference mismatch in diffusion models. Specifically, during the sampling procedure, the actual noise encoded in intermediate states drifts to larger noise levels than the intended schedule, which harms generation quality. To address this, the authors propose Noise Aware Guidance (NAG), which adds a guidance term based on the gradient of the posterior noise level to pull trajectories back toward the correct noise manifold; they provide both a classifier-based version using an external noise-level estimator and a classifier-free variant implemented by mixing conditional and unconditional scores via noise-condition dropout.", "questions": "N/A", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- **Well-written and easy to follow.** The paper clearly motivates the noise-shift train–inference mismatch and steadily builds evidence for it.\n- **Principled method.**: NAG is logically derived (using ($\\nabla_x \\log p(t\\mid x)$) to pull states back to the intended noise level), simple to plug in, and compatible with diffusion/flow samplers and with CFG/DoG.\n- **Consistent gains.** Across DiT/SiT backbones and multiple datasets, NAG improves FID and reduces the measured noise shift, with ablations that isolate its contribution and clarify when it helps (notably at higher SNR).", "weaknesses": "- My major concern is that there are some papers targeting the distribution mismatch between training and inference. In general, it is called exposure bias, and this is well explored in [1]. To summarize that paper, they derive how the variance of the backward diffusion model's process could deviate from the forward diffusion process, and give a solution scaling the estimated score to match the variance. Comparing this submission to [1], the flow is extremely similar: 1) finding the distribution mismatch between training and inference, and 2) solving this mismatch with NAG rather than using the score scaling method in [1]. Considering this similarity, the authors' contribution is that finding and addressing the distribution mismatch between training and inference no longer holds. Furthermore, completely ignoring these works makes me suspicious that the authors investigated related literature well. I recommend comparing [1] and the following papers about exposure bias, and adding these works to the related work section.\n\n- The method should be evaluated across various model sizes. And, I am also curious whether the distribution mismatch trend could be changed according to model size.\n\n\nReference \n\n[1] ELUCIDATING THE EXPOSURE BIAS IN DIFFUSION MODELS, ICLR 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AaLPVkDMpT", "reviewer": "Reviewer_wzAN", "summary": "This paper addresses the noise shift problem in diffusion models: sampling intermediate states has a mismatch between the pre-defined noise schedule and the actual noise levels. The authors propose Noise Awareness Guidance (NAG) to guide back to correct noise. Through empirical results on ImageNet generation and supervised fine-tuning tasks, the paper demonstrates that NAG mitigates noise shift and improves generative sample quality across mainstream diffusion model architectures.", "questions": "What is the computational overhead (parameters, runtime, memory) of adopting NAG in practical sampling/fine-tuning settings? How does this compare to other recalibration techniques?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper identifies the noise shift with strong empirical evidence (fig 1) , indicating that the actual posterior noise levels during sampling are systematically biased relative to the intended schedule.\n- The proposed NAG is simple and effective for diffusion and flow-matching models. And the compatibility of NAG with other guidance mechanisms (e.g., CFG and DoG) is empirically supported and presented in Table 2 and Figs. 4 and 5.\n- Clear writing, good ablations, and reproducible (with settings in Appendices)", "weaknesses": "The paper omits discussion against several closely related studies addressing noise misalignment and noise correction in diffusion models, for instance, \"Enhancing Sample Generation of Diffusion Models using Noise Level Correction\", https://arxiv.org/abs/2412.05488", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6wolAgGH4x", "reviewer": "Reviewer_zvzx", "summary": "This paper identifies a misalignment between the noise levels of the intermediate sampling states and those of the ground-truth data used during training. The authors demonstrate that this issue is prevalent in off-the-shelf diffusion models, highlighting a new direction for improving sampling performance. To address this misalignment, they propose NAG, a guidance technique analogous to classifier-free guidance, designed to correct the sampling trajectory. Empirical results validate the effectiveness of the proposed method.", "questions": "Please refer to the weaknesses box.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is generally well-written and easy to follow. \n2. Improving the sampling quality by mitigating the misalignment of the noise level between training and sampling appears new to me. \n3. The proposed method is simple to implement and can be easily applied to an off-the-shelf pre-trained model. \n4. According to the provided empirical results, the proposed method is very effective.", "weaknesses": "1. The paper lacks an ablation study on the choice of $w_\\text{nag}$, which makes it difficult for readers to gain a comprehensive understanding of the proposed method.\n2. It is known that classifier-free guidance (CFG) trades sample diversity for sample quality. Since the paper claims that the proposed method mitigates noise-level misalignment and provides an effect \"orthogonal\" to CFG, the results would be more convincing if an ablation study were included to analyze how $w_\\text{nag}$ influences sample diversity.\n3. (Minor) The work is mainly empirical and lacks theoretical support.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lABI3Ajkou", "reviewer": "Reviewer_Uw2P", "summary": "This paper addresses the issue of  noise shift, a discrepancy in denoising generative models where the actual noise level encoded in an intermediate state deviates systematically from the pre-defined noise schedule.  The authors provide empirical evidence for this shift using an external noise estimator and argue that it leads to sub-optimal generation. To mitigate this, they propose Noise Awareness Guidance (NAG), a training-based correction method that directs the sampling trajectory to remain consistent with the intended noise level.   The method is evaluated on DiT and SiT models for the ImageNet-256 dataset and various supervised fine-tuning tasks, demonstrating competitive improvements in FID.", "questions": "1. When the prior noise schedule changes (e.g., from uniform to logSNR), does the model need to be retrained to adapt to the new sampling conditions?\n\n2. Is addressing noise shift alone sufficient? How does the method distinguish between noise shift caused by the inherent properties of the sampling process and noise shift arising from external factors, such as approximations, numerical discretization, or model-specific biases?\n\n3. I'm curious whether NAG could help with the numerical instability issue that arises from the schedules in diffusion models, especially during the sampling process as $\\sigma_t$ approaches 0?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper discusses and provides some empirical evidence for the Noise Shift.\n\n2.  The approach is based on using the noise level $t$ as a conditioning signal for guidance, drawing an analogy to classifier guidance.\n\n3.  Empirical results indicate that NAG can complement existing guidance (e.g., CFG) and shows potential for improving FID scores on DiT/SiT models.", "weaknesses": "1. The paper lacks a thorough comparison with recent methods (e.g., [1]).  The authors incurred the cost of training a specialized component (NAG) and only validated it using basic samplers (DDPM and Euler-Maruyama), but failed to validate its benefit on advanced solvers. \n\n2. Experimental results suggest that the classifier-free NAG variant improves noise estimator accuracy, but it doubles the computational cost by requiring two forward passes per sampling step. Additionally, while the manuscript briefly covers the noise estimator’s training, it lacks a thorough analysis of its robustness, especially for out-of-distribution samples. \n\n3.  *Writing Quality*:  The manuscript requires significant proofreading. For instance,  in the main text, numerous typos include: \n     -  \"*pespetive*\" (Line 475) should be \"*perspective*\",\n\n     -  \"*apporach*\" (Line 480) should be \"*approach*\",\n\n     -   \"*varients*\"  (Line 480)  should be \"*variants*\",\n\n     -  \"*posible*\" (Line 483) should be \"*possible*\",\n\n     -  \"*nosie*\"  (Line 476) should be \"*noise*\",\n\n     -  \"*We anaylsis*\"  should be \"*We analyze*\", ....\n\n\n[1]. Abuduweili, A., et al., Enhancing Sample Generation of Diffusion Models using Noise Level Correction,  TMLR 2025.\n\n[2]. Ning, M., et al.,  Elucidating the Exposure Bias in Diffusion Models, ICLR 2024. \n\n[3]. Lin, S., et al.,  Common diffusion noise schedules and sample steps are flawed, WACV 2024 .\n\n[4]. Tang, Z.,  et al., Inference-Time Alignment of Diffusion Models with Direct Noise Optimization, ICML 2025.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)