PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jincheng Zhong, Boyuan Jiang, Xin Tao, Pengfei Wan, Kun Gai, Mingsheng Long

arXiv:2510.12497v1cs.LG

TL;DR

Denoising models can sample at states whose encoded noise differs from the predefined schedule, creating a training–inference mismatch and sub-optimal denoising. The paper introduces Noise Awareness Guidance and a classifier-free variant, reporting improved generation quality across ImageNet generation and supervised fine-tuning.

  • Problem

    Sampling errors can shift the noise level encoded in intermediate states away from the predefined schedule, causing training–inference mismatch and sub-optimal denoising.

    Source: §1 · pp. 1–2

  • Method

    Noise Awareness Guidance steers samples toward their predefined noise level, with a classifier-free variant combining conditional and unconditional score estimates trained using noise-condition dropout.

    Source: §1 · p. 2, §4 · pp. 5–6

  • Results

    Across DiT and SiT, NAG mitigates noise shift and improves generation quality in ImageNet conditional generation and supervised fine-tuning.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The evaluations indicate that foundation-model generation and downstream supervised fine-tuning can benefit from incorporating NAG.

    Source: §1 · p. 2

    Eliminating estimated noise shift is not sufficient for optimal generation because estimator inaccuracies or other complex factors may remain.

    Source: §5.3 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

During sampling, accumulated errors can shift the noise encoded in intermediate states away from the predefined schedule, creating training–inference misalignment. The paper proposes Noise Awareness Guidance (NAG) and reports improvements across ImageNet generation and supervised fine-tuning.

1 INTRODUCTION

  • Noise shift arises as accumulated sampling errors make the noise encoded in intermediate states differ from the predefined schedule, mismatching training and inference.The paper identifies imperfect network approximation, numerical discretization, and stochastic factors as error sources.Source: §1 · p. 1
  • ImageNet posterior estimates show a systematic drift toward larger noise levels than scheduled, with Figure 1 comparing priors t = 0.7, 0.5, and 0.3.The analysis compares estimated posteriors for sampled states with those for forward-process states.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • This mismatch can cause both out-of-distribution generalization errors and inaccurate denoising updates using predefined coefficients.The model is applied to shifted states rather than intended states, while the next state is computed using inaccurate coefficients.Source: §1 · p. 2
  • NAG uses an estimated noise-level signal to steer shifted samples toward the predefined noise level, while its classifier-free variant combines conditional and unconditional score estimates.Noise-condition dropout trains the two model behaviors jointly and removes dependence on an external classifier.Source: §1 · p. 2
  • Across DiT and SiT, NAG alleviates noise shift and improves generation quality in ImageNet conditional generation and supervised fine-tuning.The evaluations cover both foundation-model generation and small downstream datasets.Source: §1 · p. 2

2 PRELIMIARY

The preliminary section frames diffusion and flow-based generation as stochastic processes connecting prior noise and data. It reviews forward interpolation, probability-flow ODE and reverse-time SDE formulations, and conditional generation.

2 PRELIMIARY

  • The unified stochastic-interpolant framework describes diffusion and flow-based models as processes transforming a Gaussian prior sample into a data sample.The section adopts the notation of Ma et al. (2024).Source: §2 · p. 3
  • The forward process defines a continuous-time stochastic interpolant over t ∈ [0, T], with α_t and σ_t interpolating between clean data and pure noise.α_t decreases while σ_t increases.Source: §2 · p. 3
  • The probability-flow ODE describes the evolution through a velocity field given by the conditional expectation of the interpolant's time derivative.A neural network parameterizes the velocity, and samples are generated by integrating backward from Gaussian noise.Source: §2 · p. 3
  • The same marginals can be described by a reverse-time SDE whose dynamics use the score function and a diffusion coefficient.The section also relates the score to the velocity-field parameterization.Source: §2 · p. 3
  • For conditional generation, the conditional score adds a likelihood-gradient term to the unconditional score, enabling classifier guidance or classifier-free guidance.The section notes that classifier guidance models the conditional likelihood on noisy data.Source: §2 · p. 4
  • The analysis primarily uses a linear interpolant with T = 1, while the authors state that it extends to DDPM, alternative schedules, and other parameterizations.The stated linear schedule is α_t = 1 − t and σ_t = t.Source: §2 · p. 4

3 NOISE SHIFT ISSUE IN THE DENOISING PROCESS

The section defines noise shift as a mismatch between the intended noise schedule and the noise encoded in intermediate sampling states. It links accumulated errors to a systematic drift toward noisier states and reports empirical evidence of this effect, while noting limits on its diagnosis.

Noise shift issue in the denoising process

  • Noise shift is framed as a mismatch between the posterior noise inferred from perturbed states and the predefined prior noise schedule.The comparison concerns the intermediate distribution encountered during numerical SDE or ODE sampling versus the training distribution.Source: §3 · p. 4

Relation between noise shift and additive error

  • Accumulated errors can act like added Gaussian perturbation, shifting a state to an effective later noise level t′ = t + δ with δ > 0.The analysis is based on simplified assumptions but is described as qualitatively capturing noise shift in practical denoising.Source: §3 · p. 4

Empirical analysis

  • On ImageNet 256 × 256 samples from pre-trained SiT-XL/2, posterior estimates show noise shifting toward larger values during denoising.The comparison uses a learned noise estimator and forward-process samples as a reference for its accuracy on ground-truth intermediate states.Source: §3 · p. 4

The effect of noise shift δ

  • At mid-level noise, the estimator substantially overestimates noise, but the diagnosis is constrained by estimator accuracy and δ is sufficient, not necessary, evidence of sub-optimal denoising.The observed shift does not establish that every sub-optimal denoising behavior involves noise shift.Source: §3 · p. 5
  • Noise shift can cause out-of-distribution errors in the learned model and inaccurate reverse-time SDE coefficients, making denoising sub-optimal.The model processes shifted intermediate states rather than the states intended for the predefined noise level.Source: §3 · p. 5

4 NOISE AWARENESS GUIDANCE

Noise Awareness Guidance (NAG) reinforces conditioning on the intended noise level to steer shifted intermediate states back toward the predefined schedule. Its classifier-free variant uses conditional and unconditional scores, while its noise-level control complements class-conditional guidance.

Noise Awareness Guidance

  • NAG reinforces conditioning on intended noise t to steer intermediate states toward consistency with the predefined noise schedule.The method targets the mismatch between a sampled state and its intended noise condition.Source: §4 · p. 5

Noise awareness guidance

  • Classifier-based NAG guides sampling with ∇log gϕ(t | x), using an external posterior estimator to represent the noise level encoded in an intermediate state.The estimator supplies the guidance signal for the conditional reverse-time sampling trajectory.Source: §4 · p. 5

Classifier-free noise awareness guidance

  • Classifier-free NAG combines conditional and unconditional scores and trains a shared model by randomly dropping noise condition t, avoiding a separate noise-level predictor.This variant addresses the external estimator's computational cost, pipeline complexity, and risk of adversarial-like behavior.Source: §4 · p. 6

Discussion and relation to CFG

  • NAG's noise-level conditioning axis is orthogonal to CFG's class-conditioning axis, providing complementary control; CFG may mitigate noise shift indirectly, while NAG targets δ directly.CFG's indirect effect is attributed to its bias toward lower-temperature regions where models are more confident.Source: §4 · p. 6

5 EXPERIMENTS

The experiments test NAG in ImageNet generation and supervised fine-tuning, then examine how it changes noise shift during sampling. Across these settings, the authors report improved generation quality and reduced noise shift.

5 EXPERIMENTS

  • The evaluation covers ImageNet generation with DiTs and SiTs and supervised fine-tuning on small, fine-grained datasets.The ImageNet setup uses 256 × 256 resolution; the reported fine-tuning FIDs use 10,000 samples.Source: §5 · p. 6, §5.1 · pp. 6–7, §5.2 · pp. 7–8

5.1 NAG FOR IMAGENET GENERATION

  • After 80 epochs without CFG, NAG substantially improves DiT and SiT results; with extensive pre-training, it approaches CFG quality alone and adds improvements when combined with CFG.The converged ImageNet comparisons fine-tune DiT-XL/2 and SiT-XL/2 checkpoints for 10 additional epochs to support NAG sampling.Source: Figure 3 · p. 7, §5.1 · p. 7, Table 1 · p. 7

5.2 NAG FOR SUPERVISED FINE-TUNING

  • NAG improves FID across supervised fine-tuning tasks and baselines, including CFG and Domain Guidance (DoG).The authors also report that NAG alone performs comparably to CFG-guided sampling.Source: Table 2 · p. 8, §5.2 · p. 8

5.3 EMPIRICAL OBSERVATIONS OF NOISE SHIFT WITH NAG

  • NAG reduces noise shift most when the signal-to-noise ratio exceeds 1, while its effect is less pronounced early in denoising; it also reduces residual shift under CFG.The observed shift grows and then plateaus during sampling, with estimates at very low noise potentially overstated because of estimator error.Source: §5.3 · pp. 8–9, Figure 5 · p. 9

6 RELATED WORK

Prior work addresses denoising errors through schedules, objectives, architectures, or faster integration methods. This paper frames noise shift as a distinct issue and positions NAG as guidance using noise level itself as a signal.

Denoising generative models

  • Prior denoising work has targeted schedule design, training objectives, model capacity, or numerical errors from discretized integration.Faster methods include high-order solvers and improved interval modeling.Source: §6 · p. 9
  • The paper contrasts its focus on noise shift with prior approaches that target particular sources of error in denoising generation.It argues that addressing noise shift alleviates persistent sub-optimality in generation.Source: §6 · p. 9

Guidance techniques for condition generations

  • NAG is presented as the first method to use noise level itself as a guidance signal to align sampling with the intended noise condition.The paper distinguishes this target from guidance techniques focused on alignment with conditioning information.Source: §6 · p. 9

7 CONCLUSION

The paper identifies noise shift as a widespread drift toward larger noise levels and presents Noise Awareness Guidance and a classifier-free variant to mitigate it. It argues that this misalignment can make denoising generations sub-optimal.

7 CONCLUSION

  • The paper reports that noise shift tends toward larger noise levels and that iterative denoising under shifts produces sub-optimal generations.It presents NAG and a classifier-free variant to reduce the noise-shift gap.Source: §7 · pp. 9–10

A DERIVATION OF STATEMENT 1

The derivation models an intermediate state with additive Gaussian error as a forward-process state at a shifted noise level. Under differentiability and small-error assumptions, it approximates the expected shift; implementations use the official DiT and SiT codebases.

A DERIVATION OF STATEMENT 1

  • The derivation assumes additive error e is zero-mean Gaussian with unknown variance, independent of the forward-process Gaussian noise.Their weighted sum is Gaussian, allowing the perturbed state to be expressed through a shifted noise level.Source: A DERIVATION OF STATEMENT 1 · p. 12
  • Noise shift δ is defined by matching the perturbed-state distribution to a forward-process intermediate state at t′ = t + δ.The derivation then solves for an approximation to δ in the presence of additive Gaussian error.Source: A DERIVATION OF STATEMENT 1 · p. 12
  • A first-order approximation assumes σ_t is differentiable and the error variance is small, so the resulting shift δ is also small.The section solves for an approximation to δ under these conditions.Source: A DERIVATION OF STATEMENT 1 · p. 12

B IMLEMENTATIONS

  • Experiments use the official DiT and SiT codebases.The work reports its derivation alongside this implementation subsection.Source: B IMLEMENTATIONS · p. 12

B.1 IMPLEMENTATION TO MAIN RESULTS

The implementation follows DiT and SiT architectures and samplers, with specified guidance weights and training setups for ImageNet and downstream fine-tuning. FID evaluation follows the original task protocols.

B.1 IMPLEMENTATION TO MAIN RESULTS

  • DiT and SiT experiments use the respective DDPM and Euler–Maruyama samplers, with default NAG and CFG guidance weights specified for each setup.DiT uses DDPM; SiT uses Euler–Maruyama with w_t = σ_t and a final step size of 0.04.Source: §B.1 · p. 13
  • ImageNet experiments train from scratch or fine-tune pre-trained models for 1,400 epochs with an unconditional noise branch, retaining most original DiT and SiT training configurations.Fine-tuning uses a pseudo noise level, a learning rate of 1 × 10^-5, and 20% noise dropout.Source: §B.1 · p. 13
  • Downstream fine-tuning follows Domain Guidance settings for 24,000 iterations, using a learning rate of 1 × 10^-4, batch size 32, and 10% noise dropout for NAG.The setup uses AdamW across datasets.Source: §B.1 · p. 13
  • FID evaluation follows each task’s original protocol, using generated samples compared with ImageNet training images or downstream test images.ImageNet uses 10K or 50K generated images; downstream tasks use 10K.Source: §B.1 · pp. 13–14

B.2 IMPLEMENTATION OF EMPIRICAL POSTERIOR ESTIMATOR gϕ.

The empirical posterior estimator gϕ is trained as a noise-level regressor on ImageNet latent states. The analysis constructs forward and reverse trajectories and estimates their densities using KDE.

B.2 IMPLEMENTATION OF EMPIRICAL POSTERIOR ESTIMATOR gϕ.

  • The posterior estimator gϕ replaces the final layer of a pre-trained SiT-XL/2 checkpoint with a two-layer MLP noise-level regressor.Its hidden-state input is globally averaged before prediction.Source: §B.2 · p. 14
  • The estimator is trained on SD-VAE latent states for 40 epochs, minimizing L2 prediction error with the noise condition masked by a pseudo condition.Training reaches a loss of 0.0002, with no EMA applied.Source: §B.2 · p. 14
  • The density analysis builds 5,000 forward and 5,000 reverse trajectories, then estimates densities with KDE for samples grouped by prior noise level and generation process.Reverse trajectories use a 20-step Euler–Maruyama SDE solver.Source: §B.2 · p. 14

C MORE VISUALIZATION RESULTS WITH KERNEL DENSITY ESTIMATION

This section extends earlier figures with posterior-density visualizations and examines the sensitivity of NAG to its guidance weight and sampling steps.

C MORE VISUALIZATION RESULTS WITH KERNEL DENSITY ESTIMATION

  • Posterior-density results extend Figures 1 and 5 with sampled-state and forward-process estimates shown against the predefined noise level.Figure 6 uses yellow for sampled intermediate states, orange for forward-process states, and a black indicator for predefined t.Source: C MORE VISUALIZATION RESULTS WITH… · p. 14, Figure 6 · p. 15
  • Figure 7 contrasts sampled-state density under CFG and NAG with forward-process density and the predefined noise level.The caption identifies blue as CFG-influenced and pale gold as the mitigating effect of NAG.Source: Figure 7 · p. 16

D SENSITIVITY TO HYPERPARAMETERS OF NAG

Related papers

Loading 2510.12497v1…