Enhancing Sample Generation of Diffusion Models using Noise Level Correction
Authors: Abulikemu Abuduweili
Chenyang Yuan
Changliu Liu
Frank Permenter
Venue: Accepted by TMLR ()

Abstract: The denoising process of diffusion models can be interpreted as an approximate projection of noisy samples onto the data manifold. Moreover, the noise level in these samples approximates their distance to the underlying manifold. Building on this insight, we propose a novel method to enhance sample generation by aligning the estimated noise level with the true distance of noisy samples to the manifold. Specifically, we introduce a noise level correction network, leveraging a pre-trained denoising network, to refine noise level estimates during the denoising process. Additionally, we extend this approach to various image restoration tasks by integrating task-specific constraints, including inpainting, deblurring, super-resolution, colorization, and compressed sensing. Experimental results demonstrate that our method significantly improves sample quality in both unconstrained and constrained generation scenarios. Notably, the proposed noise level correction framework is compatible with existing denoising schedulers (e.g., DDIM), offering additional performance improvements.

Published in Transactions on Machine Learning Research (6/2025)
Enhancing Sample Generation of Diffusion Models using Noise Level Correction
Abulikemu Abuduweili Carnegie Mellon University
abulkken@ondve..cmu.
Chenyang Yuan Toyota Research Institute
chenyang-yuaan www.rouurrriigllbbblooo
Changliu Liu Carnegie Mellon University
cliu@@cndrew.cmu.co edu
Frank Permenter
Toyota Research Institute
Reviewed on OpenReview: https:
// opwrrevee.......e.em net/ forum? id- y8V
Abstract
The denoising process of diffusion models can be interpreted as an approximate projection of noisy samples onto the data manifold. Moreover, the noise level these samples approxi- mates their distance the underlying manifold. Building on this insight, we propose novel method to enhance sample generation by aligning the estimated noisc level with the truc distance of noisy samples to the manifold. Specifically, we introduce noise level correction network, leveraging a pre-trained denoising network, to refine noise level estimates during the denoising process. Additionally, we extend this approach to various image restoration tasks by integrating task-specific constraints, including inpainting, deblurring, super- nuper-resolution, colorization, and compressive sensing. Experimental results demonstrate that our method significantly improves sample quality in both unconstrained and constrained generation sce- narios. Notably, the proposed noise level correction framework is compatible with cxisting denoising schedulers (e.g., DDIM), offering additional performance improvements.
1 Introduction
Generative models have significantly advanced capability of creating high-fidelity data samples across
various domains such images, audio, and text (Song & Ermon 2019 Sohl-Dickstein et al. 2015 Oussidi & Elhassouny 2018). Among these, diffusion models have emerged as one of the most powerful approaches due
to their superior performance in generating high-quality samples from complex distributions (Song & Ermon 2019 Sohl-Dickstein et al.) 2015 Li et al. 2024a). Unlike previous generative models, such as generative adversarial networks (GANs) (Goodfellow et al. 2014h and variational autoencoders (VAFs) (Kingma &
Welling 2014), diffusion models add multiple levels of noise to the data, and the original data recovered through a learned denoising process (Ho al., 2020 Song et al., 2023). This allows diffusion models to handle high-dimensional, complex data distributions, making them especially useful for tasks where sample quality and diversity are critical (Cao et al. (202)). Their capability to generate complex, high-resolution data has led to widespread applications across numerous tasks, from image generation in models like DALL E Ramesh et al. 2022) and Stable Diffusion (Rombach et al. 2022 to use in robotic path-planning and control
Janner et al. 2022 Chi et al. 2023), as well as text generation Li et al. 2022).
Permenter & Yuan 2024] have interpreted the denoising process in diffusion models approximate projection onlo the data. manifold, with the noise level a approximating
Previous studies (Rick Chang et al. 2017

---

Published in Transactions on Machine Learning Research (6/2025)
the distance between noisy samples and the data manifold. This perspective views the sampling process optimization problem, where the goal is to minimize the distance between noisy samples and the underlying data manifold using gradient descent. The gradient direction is approximated by the denoiser output with step size determined by the noise level schodule. However, the denoiser requires an estimate of the noise level as input. We claim that accurately estimating the noise level during the denoising process essentially the distance the data manifold- is crucial for convergence and accurate sampling.
Colorization
Super resolution
Compressive Sensing 12.5%
Inpainting
ICAM
Degraded
co I
DDNM
Ours
Ground
truth
Figure 1: Qualitative results of constrained image generation.
The expressive capabilities of diffusion models have also made them a compelling choice for image restoration tasks, where generating high-quality, detailed images is essential (Wang et al. 2024). Diffusion models can be used as an imago-prion for capturing the underlying structure of image manifold and have shown significant promise for constrained generation, such as image restoration (Song et al. 2022; Chung et al,, 2023a; Wu et al.) 2024). Plug-and-play methods were proposed to utilize pre-trained models without the need for extensive retraining or end-to-end optimization for linear inverse problems such as super-resolution, inpainting, and compressive sensing (Chung et al.) 2023a Dou & Song 2024 Wang et al. 2023). These methods can be interpreted as alternating taking gradient steps towards the constraint set (linear projection for linear inverse problems) and the image manifold (noise direction estimated by the Icarned denoiser) to find their intersection. However, they may sulfer from inconsistency issues if, after each projection step, the noise level no longer approximates distance (Figure 2). Thus, correcting the noise level after each step could increase the accuracy of image restoration tasks.
In this work, we propose a novel noise level correction method to refine the estimated noise level and enhance sample generation quality. Our approach introduces a noise level correction network that aligns the estimated noise level of noisy samples more closely with their true distance to the data manifold. By dynamically adjusting the sampling step size based on this corrected noise level estimation, our method improves the sample generation process, significantly enhancing the quality of generated data. Furthermore, our approach integrates seamlessly with existing denoising scheduling methods, such as DDPM (Denoising Diffusion Probabilistic Models) (Ho et al. 2020 DDIM (Denoising Diffusion Implicit Models) (Song et al., 2021a), and EDM (Karras et al. 2022). Furthermore,
extend the application of noise level correction to various image restoration tasks, showing its ability to improve the performance of diffusion-based models such as DDNM (Denoising Diffusion Nullspace Model) by (Wang el al. 2023). Our method achieves improved results across tasks including inpainting, deblurring, super-resolution, colorization, and compressive sensing, as illustrated in Figure [ Additionally, we introduce a parameter-free lookup table as an approximation of the noise level correction network, providing computationally efficient alternative for improving the performance of unconstrained diffusion models. In summary, our contributions are:
2

---

Published in Transactions on Machine Learning Research (6/2025)
We propose noise level correction network that improves sample generation quality by dynamically refining the estimated noise level during the denoising process.
We extend the proposed method to constrained tasks, achieving significant performance improvements various image restoration challenges.
We develop à parameter-free approximation of the noise level correction network, offering a computation- ally efficient. tool to improve diffusion models.
Through extensive experiments, we retenstrate that the proposed noise level correction method consis- tently provides additional performance gains when applied on top of various denoising methods. Our code is available at https: /etps://github... thub. httpp//..waaallllla////tftofaoo.
2 Background
2.1 Diffusion Models
Diffusion models represent a powerful class of latent variable generative models that treat datasets as samples from a probability distribution, typically assumed to lie on a low-dimensional manifold K. C IR" (Sohl- Dickstein et al. 2015 Ho et al. 2020). Given a data point 20 E K, diffusion models aim to learn a model distribution ion (=)) that can approximate this manifold and enable the generation of high-quality samples. The process involves gradually corrupting the data with noise during a forward diffusion process and incrementally denoising it to reconstruct the original data through reverse generative process.
Diffusion (forward) process. In the forward process, noise is added progressively to the data. Starting with a clean sample 20, the noisy version at step t, denoted is linear combination of zg and Gaussian noise c:
2/ = 20 + VI-a
(1)
where e ~ NO(, I) and the noise schedule a controls the amount of noise injected al each step. Typically, > 01 02 > or 0, ensuring that p(2r) ()) for large enough T (Ho et al.) 2020 For mathematical convenience, a reparameterization often employed, defining new variables It 2/ yan. which results in:
Jt- To + oie, where ~N(, I),
(2)
0t v
at
24
Tt
x0 = 20.
(3)
01
Where de denotes the noise level. Note that the diffusion process is originally presented in variable 2t, we 2 use the formulation x; you to simplify the forward and reverse denoising processes. A similar formulation
can be found in (Song et al. 2021b: Karras et al. 2022
Denoiser. Diffusion models are trained to estimate the noise vector added to a sample during the forward process. The learned denoiscr, denoted as co, is predict the noise voctor e from the noisy sample x, and the correspouding noise level dt- The denoiser is optimized using a loss function that minimizes the difference between the predicted and true noise vectors:
L Elleo(Z,, 0,) e Eelat,elles + oze, -el
Here, Zo is sampled from the data distribution, a, drawn from discrete predefined noise level schedule, and is drawn from a standard Gaussian distribution, (0, I). Training is typically performed using gradient descent, where randomly sampled triplets (To. e, a;) are used to update the denoiser's parameters 0. Once the denoiser is trained, we can apply a one-step estimation to approximate the clean sample onpt R øo:
Aojt = _t - Ozee(xt,
(5)
Denoising (sampling) process. The one-step estimation, eq. 5 may lack accuracy, in which case the trained denoiser is applied iteratively through the denoising process. This process aims to progressively
3

---

Published in Transactions on Machine Learning Research (6/2025)
denoise noisy sample ør and recover the original data zu- Sampling algorithms construct a sequence of intermediate estimates (IT, IT--; starting from an initial point 2r drawn from Gaussian distribu-
tion, 2T 1). One of the widely used samplers, the deterministic DDIM (Song et al. 2021ah, follows
the recursion:
Tt-1 = It. + (ot-1 (z)ea(11, 0t) = -Bo + oz- 100160())
(6)
2T
21
o 1 2T. 2T I),,1),
Vor
where eo is the predicted noise step f. This iterative process continues until zo is obtained, which represents denoised sample. On the other hand, the randomized DDPM (Ho ct al. 2020) follows the update rule:
2t-1 Zt (ov -r)co((2, 0z) + nwe fo.= co + nww
(8)
or = T = Vo- - oz, We ~ ~V0,11)
(9)
2.2 Additional Related Works
Diffusion models have gained significant attention for their ability to learn complex data distributions, excelling in diverse applications such as image generation (Ho et al.] 2020 Arechiga et al. 2023), audio synthesis (Kong et al. 2021), and robotics (Chi et al. 2023). On the theoretical side, significant research has explored the non-asymplotic convergence rates of various diffusion samplers, including DDPM (Benton et al. 2024) and DDIM (Li et al., 2024b), contributing to a deeper understanding of the optimization processes underlying certusion-basar models.
Few-step Sampling. Several methods have been proposed to improve the quality of generated samples while reducing the number of iterations. The first class of these techniques require modifications to the training process, or training of additional models. They include model distillation (Nichol & Dhariwal 2021), progressive distillation (Salimans & Ho. 2022) or using consistency models (Song et al. 2023 Kim
et al. 2023) to predict the endpoint of the sampling process. They can usually produce samples in fewer than 10 steps but are computationally expensive to train and could suffer from training instabilities. The second class of techniques do not require any additional training but instead modifies the sampling process. These include faster samplers known as EDM (Karras et al. 2022 DPM (Lu et al., 2022) modifications to the sampling noise schedule in (Sabour et al.] 2024 Hang et al., 2024). They do not need additional training and can be directly used on pre-trained models, but this also limils their performance compared Lo training- based methods.
Adaptive noise schedules. Our proposed mothod can be seen as a way of introducing additional compu- falion the sampling process to adaptively change the sampling noise schedule depending on input images. Noise schedules play an important role during sampling, and (Chen) (023) showed that different schedules are optimal for different data. To tailor the noise schedule to different inputs, (Sahoo et al. (024)) learns to adapt the noise schedule depending both the pixel location and input. AdaDiff (Tang al. 2024) dynamically adjusts the inference computation during sampling using learned uncertainty estimation module.
Diffusion models for inverse problems. Diffusion models have also demonstrated effectiveness in im- age restoration tasks, including super-resolution (Chung et al.] 2023a), inpainting (Kawar et al. 2022 deblurring, and compressive sensing (Wang et al., 2023). These tasks often roquire strict adherence to data consistency constraints, making diffusion models particularly suitable for addressing such challenges. Prior work such as DDRM (Kawar et al.) 2022) and DDNM (Wang et al.) 2023) pioneered the use of diffusion models to solve linear inverse problems by projecting noisy samples onto the subspace of solutions that satisfy lincar constraints. Other methods have employed alternative approaches, such as computing full projections by enforcing linear constraints or using the gradient of quadratically penalized constraints (Chung el al., 2022 2023b). Recent work has extended the application of diffusion models to more complex nonvconvex constraint functions. In these cases, itcrative methods such as gradient descent are utilized to guide the sam- pling process toward satisfying the constraints, as demonstrated in Universal Guidance (Bansal et al. 2023). Furthermore, a provably robust. framework for score-based diffusion models applied to image reconstruction
4

---

Published in Transactions on Machine Learning Research (6/2025)
introduced by Xu & Chi 2024 offering robust performance in handling nonlinear inverse problems while ensuring consistency with the observed data.
3 Methods
3.1
Noise Level as Distance from Noisy Samples to the Manifold
This work builds on the insight that denoising in diffusion models can be interpreted as an approximate projection onto the support of the training-set distribution. Previous studies (Rick Chang et al. 2017 Permenter & Yuan 2024) have established this connection. The distance function of set C R", denoted as distx (z), defined as the minimum distance from point æ to any point zg K:
distx(a)) == inf [|x zoll zo E K.).
(10)
The projection onto IC, denoted projx(2); refers to the point (or points) on K that achieves this minimum distance. Assuming the projection is unique, we can express as:
(r)) == (ar E K: (istx(())= || zall
(11)
Permenter & Yuan 2024) introduce the following relative-error model, inspired by the manifold hypothesis
Bengio et al., 2013; Fefferman et al. 2016) as well as from the fact that the idcal/optinn denoiser is the
gradient of smoothed distance function:
Assumption 3.1 (Informal) The learned denoiser e" as well as the sampling noise schedule {oelem sat- isfies the following for all x, during sampling:
1. oe decreases slowly enough to ensure that vnoi zo distx((1) after each iteration
If o; approrimates distx((r), then Dt - oie (1) projx( with constant relative error with respect to dislx(xx))
Leveraging this assumption, we outline the following theorem informally, which summarizes results from Permenter & Yuan 2024h (for formal proofs refer to [Permenter & Yuan 2024)))
Theorem 3.1 (Informal) Suppose Assumption 3.1 holds and the initial distance satisfies distx((z))
VnoT. Then the DDIM sampler generates the sequence (IT, .0) by performing gradient descent on the objective function f(x)= == (i)) with a step-size of Be == 01-1/01:
I1-1 = It Br ((xx) = I, - 3r dislx(xt) v (dis) dist = Vnoi, V (dist(((x))) = Ez o.)/vn
(12)
(13)
Theorem 3.1 demonstrates that the estimated clean sample generated by the denoising process Zoj: (Te ot) serves as an approximation of the projection of the noisy sample onto the manifold IC. It also establishes that throughout the sampling process, the distance of a noisy sample 2, onto manifold KC, can be approximated by the noise level vna,- Given that a, and 69 satisfies Assumption 3.1 the DDIM process generates a sequence de with monotonically decrcasing distance the manifold by gradient descent the squared-distance function to the manifold. Specifically, the distance of each noisy sample from the manifold is determined by the noise level, dxs) = Vño,, while the gradient V ff(z), which is the direction of projection of z4 onto the manifold, is estimated by the noise vector (x. R 6j (21, ot)/
Theorem 3.1 exposes two key assumptions that should be satisfied by the denoising process in order to obtain good samples on XC: (1) a denoiser eø (2,, ot) satisfying the relative projection-error assumption, and (2) a noise schedule that docrease slowly enough so that for all f, we have distx;(a;) x /no,, For sample generation, cumulative errors introduced during the denoising process can lead to biases due to the imperfections of the denoiser (Ning al. 2023 These errors become particularly significant. by the final steps, where deviations
5

---

Published in Transactions on Machine Learning Research (6/2025)
disty (te) + oz
X4
disty (tt)
Adj
8 distg (xp)
disty (4-1) e 06-1
01-1
Comitralnn
disty (Re-s) o Rbos
Contreant
Rott-1
Roit
Rytc
Image manifold
Image manifold
Image manifold K
(a) Denoising Process
(b) Constrained Denoising
(c) Denoising with NLC
Figure 2: (a) Standard denoising using DDIM. (b) Constrained denoising interpreted as an alternative projection. In both (a) and (b), the estimated noise level o, may not align with the distance ( (c) Constrained denoising with noise level correction (NLC), where replacing or with more accurate estimate o, yields projections that better align with the manifold K.
in the denoising process result in a large mismatch between the truc distance (xs) and the estimated noise level y as illustrated in Figure 2a As observed, a mismatch between the true distance and the estimated noise level can lead to the clean image estimation using eq. 5 falling outside the manifold I. In constrained sample generation tasks, such as image restoration, this issue becomes more pronounced. As illustrated in Figure 2b] deviations introduced by guidance terms projection steps onto the constraint C can amplify the mismatch between the estimated noise level Vna, and the actual distance dis(x(x;). This discrepancy is especially significant in the later denoising stages and can hinder the reconstructed sample from accurately aligning with the data manifold K. Additional details are provided in Appendix A
3.2 Noise Level Correction
To address the issue of inaccurate distance estimation caused by relying on a predefined noise level de as discussed in section 3.1 we propose a method called noise level correction to better align the corrected noise level with the true distance. This approach replaces the predefined noise level scheduler oz with a corrected noise level 6, during the denoising process, enabling more accurate distance estimation and improved sample quality. As shown in Figure 2c] using the corrected noise level brings the estimated clean sample tojr closer to the data manifold K. compared to the naive denoising process that relies solely de The corrected noise level is defined as o, == a, + fr] & distx(())/ Vn, where År represents the residual and can be modeled using either a neural network or a non parametric function. This residual alignment approach is effective because the residual f₂ is stable across noise levels, making it casier to modol, while the noise level itself may become unbounded during large diffusion time steps.
For (2,), calculating the ground-truth distance to the manifold is generally not feasible. Instead, we approximate it using the distance between the noisy sample and its clean counterpart in the forward diffusion process. Specifically, given the noisy samples xa generated by the diffusion process from Zo E K. in eq. ( we estimate the distance as distx(a;) R [x; zol. This approximation is reasonable when the projection of 2, onto IC satisfies (x)) proje (2o + o,e) x zq. As illustrated in (Permenter & Yuan 2024), ; this occurs when o, is small (due to the manifold hypothesis, the random noise e is orthogonal to the manifold IC), and when oe is much larger than the diameter of K.
We introduce neural network to learn f. ro(a,, o,) for noise level correction. To minimize computational costs, design the noise level correction network ro(+) to be small and efficient. It leverages the encoder module of the denoiser's UNet architecture, followed by compact layers that fully utilize the pre-trained denoiser's capabilities. As shown in Figure B the denoiscr network uses a UNet structure to estimate the noise vector (denoising direction) based on the noisy image H, and ot. Meanwhile, the noise level correction network utilizes the shared encoder, followed by additional neural network blocks, to predict the residual noise level. We train the noise level correction network re(+) alongside a fixed, pro-trained denoiser co(:), ensuring coordinated improvement in denoising accuracy. In training, to further enhance the noise level correction network, we use randomly sampled scaling factor to perturb the true noise level Øs, expanding
6

---

Published in Transactions on Machine Learning Research (6/2025)
Noisy sample Xy.07
U-Nel
EB
clean sample
Xg
denoising
Xp-1
Te
6. 011 + +r(()))))
Figure 3: Architecture of the noise level correction network ro(+) and the denoiser co())-
the input-output space re( ot). The objective function for noise level correction is defined as:
Lr, == Euo.hei [[ + re - f = Io + ofde, e~ N),,1), À U(1 6, +
(14)
(15)
The scaling factor À is sampled from a uniform distribution U(1 ( 6,11,,1100 8), with 8 = 0.5 in our experiment,
to control the level of variation introduced into the noise level correction. By allowing À to vary, e.g., U(1 - 6,1 + 6), the true noise level becomes approximately λoιy and ro(3,, adapts to capture this variation. This approach covers broader output space compared to the naive setting = 1, leading better generalization. For further discussion and an ablation study on different values, please refer appendix D.3
3.3 Enhancing Sample Generation with Noise Level Correction
The trained noise level correction network rо be integrated into various existing sampling algorithms to improve sample quality. For algorithms that include an initial sample estimate Roje eq. we reformulate the one-step estimation as follows:
fpj: - -,
(16)
6, o.1l ra o.),, & vn co(x,, 0t)
(17)
[To( 0:)1
We normalize the noise vector ca(-) during the sampling in eq. 17 process to decouple noise level (magnitude) correction o, from direction estimation & Empirical experiments show that normalizing ea(-) and using &, to account for magnitude yields better results compared to not normalizing co(-). In the training loss function eq. [4]), normalization of the noise vector is unnecessary because the randomly sampled noise é naturally concentrates around its expected norm, approximately vn. However, the neural network-estimated noise vector (e(-) does not necessarily maintain a constant norm.
Using eq. 17 we integrate noise level correction into the DDIM and DDPM sampling algorithms, illustrated in Algorithm [ with the modifications from the original ODIM/DDPH algorithms highlighted in blue. In this algorithm, lines 3 and 4 represent the current and next-step noise level corrections, respectively, while line 5 provides the normalized noise vector. Lines 6 through 8 follow the steps of the original DDIM and DDPM algorithms. Similarly, noise level correction can be integrated into the EDM sampling algorithm, as shown in Algorithm 3 Note that EDM with noise level correction, we do not normalize the noise vector cg(x,, since EDM employs a socond-order Heun solver to improve noise vector estimation. By incorporating noise level correction, these algorithms produce higher-quality samples with improved accuracy by considering both the direction and distance to the data manifold during the denoising process.

---

Published in Transactions on Machine Learning Research (6/2025)
Algorithm DDIM/DDPM with Noise Level Correction (DDIM/DDPM-NLC)
Input: Denoiser eg and noise level corrector rg
Input: Noise scheduler (ax)/=1s randomness scale 1) (1 0 for deterministic DDIM and "f 1 for DDPM
Output: samples To E K
1: oT Voj + 27, 2T ~ No0, T),
2: for = T,T 1,.. do do
3: ô; = =o,[1 + r-ro t-(o
- = 0: 111
de
5: & = Vncol ( 0t)
6: noose = 11 d1
0c-1
0, 07-1
7: ssgnnll 1-- o²
noise
8: Tt-1 = 3t + siignal de) + Noisehll, where wi ~ N((0,) 7)
end for
3.4 Constrained Sample Generation
The noise level correction method can also improve performance in constrained sample generation tasks, such image restoration. Let IC denote the data manifold, and let represent specific constraints, such masked pixel matching in inpainting tasks. Constrained sample generation aims to generale samples x that satisfy both the manifold and constraint requirements, meaning æ E K. E C. Similar to DDIM-NIC sampling approach in Algorithm [ generation methods, such as DDNM (Wang et al. 2023 a DDIM-based image restoration algorithm. An example of this shown Algorithm
noise level correction can be incorporated into existing constrained sample
To further enhance constrained sample generation, we propose a flexible iterative projection algorithm in- spired by the alternating projection technique (Bauschke & Borwein 1996 Lewis & Malick 2008). This approach iteratively projects samples onto each constraint set to approximate solution that lies in the interscction of K and C. The itcrative projection process can be expressed as follows:
(18) Where dojk and Tolk represent the k-th estimates of points satisfying x E K. and x E C, respectively. The iterative rule ensures that zоja approximates point in Knc. C. 2k+1 introduces small noise term é, which helps avoid convergence local minima in non-flat regions. This noise term is analogous the one used in DDPM (eq. (8)) and facilitates iterative refinement toward the final clean samples. The noise term ē can be gradually reduced over iterations or set to zero once salisfactory iteration k = Kmar is achieved. At this point, the algorithm returns a final estimate such that KK C.c.
fajk = (=))) Zoja = proje(fojs), I(k-1) = Tojs + 2, , = 0,
Theorem demonstrates that the projection operator projx(+) can be approximately computed using the denoising diffusion models. Starting from an initial random point To = maxt, the projection onto the manifold I. can be iteratively refined using eq. (16) and eq. [7), as follows:
(ro)x(((())) Pook T()) (())
(19)
For the additional constraint projection, (x)) the specific calculation depends the nature of the constraint. The constraints for many image restoration tasks are linear, including inpainting, colorization, super-resolution, deblurring, and compressive sensing. For tasks with linear constraints, the projection can be computed directly or optimized using gradient descent. Consider an image restoration task formulated as y = A where Ho represents the ground-truth image, y is the degraded observation, A is the linear degradation operator. Given the degraded image y and the current estimate ooJa. the projection onto the constraint can be computed as:
Tojk (rojk = A y + (I A'A)aojk
(20)
8

---

Published in Transactions on Machine Learning Research (6/2025)
Where Al is the pseudo-inverse of A. In this work, adopt the values of A for image restoration tasks provided in [Wang et al. 2023). Here, Pojs is the t-th th estimate satisfying zo E EnnC This iterative estimation requires a predefined noise scheduler 01,-- (Kk.. to generate Tojk and T(k+1)- To allow flexible and potentially unlimited refinement steps, we define the noise schedule (R) with a maximum noise level (() = maz and a minimum level min: decaying by a predefined factor 1/ < 1. If the noise level reaches miin the process can either stop with returning zojk or restart from drestart- This strategy permits an arbitrary number of refinement steps, stopping either at desired loss threshold or continuing indefinitely. Since (A) represents the distance of noisy samples from the manifold, this decaying schedule incrementally
reduces T()) is distance from the manifold.
Algorithm Constrained Sample Generation with Noise Level Correction (IterProj-NI)
Input: Denoiser eo and noise level corrector Tg
Input: Constraint C, distance decay a, noise scale
Input: mmaz; Omin, and Trestart; and maximum iterations Kmar
Output: samples ToJK E KnC
1: z(() = (o)), E ~ I), ((o) = maz
2: for k = 0, 1, .... do
3: σ() = d(n) ro(2()))
(() (T
folk I()) - -(())((k)
Tolk (rojc((o)kk (For image restoration lasks, refer eq. 20
((++1) 1) (K))
8:
if (K+1) < mmin then
9:
(k-+1) restart
10:
end
Ê(K) = / - (p(() + ne, where E ~~ 1),
11:
12:
I(k+1) = Foj + ((++))((K)
13: if l Kmaz or | Fojk-1ll 1|| is small enough then
14: return Zojk
15: end
16: end for
Algorithm 2 presents the proposed constrained generation approach, termed IterProj-NLC (Iterative Projec- tion with Noise Level Correction). Lines 3 to 5 project onto the data manifold IN, while Line 6 projects onto the constraint set C. Lines 7 to 10 update the noise level, and Lines 11 and 12 compute the next noisy sample ((++1) from the current clean estimate Tyjo- This step also acts as convex combination of the current clean estimate Tajo and the previous noisy sample I(k).
4 Experiments
4.1 Toy Experiments
We conducted a toy experiment to demonstrate the effectiveness of noise level correction sample generation for diffusion models. The goal was to generate samples on d-sphere manifold. The toy training datasct was sampled from d-dimensional sphere manifold embedded within an v-diienssional data space, where d < Detailed experimental design information is available in Appendix C.1
After training, we applied the proposed 10-step DDIM with Noise Level Correction (DDIM-NLC), as de- scribed in Algorithm [ and compared it against the standard 10-step DDIM baseline. Sample quality was assessed by measuring the Euclidean distance from generated samples to the ground-truth manifold IC, with lower values indicating better quality. As discussed in Thcorem 3.1. the noise level /no, approximates the distance to the manifold. To further evaluate this approximation, we examined the alignment between the imaimdded noise level or and the true distance (
9

---

Published in Transactions on Machine Learning Research (6/2025)
Figure 4a shows the sample-to-manifold distances over time. The dashed curve indicates the predefined noise level Vioi, while the solid curves represent the measured distances from intermediate samples to the manifold. In the DDIM baseline, the actual distance deviates from especially in later steps, highlighting the inaccuracy of predefined noise level assumptions (see also Figure 2 In contrast, the DDIM- NLC consistently produces samples that are closer to the manifold, demonstrating improved alignment through noise level correction. Figure quantifies this improvement using the relative distance estimation bias, defined
Distance Estimation Bias distx vnō
Vnoa,
Where o, or for DDIM, and o, oi[1 + re(āt, az)l for DDIM-NLC. Results show that DDIM-NLC significantly reduces estimation bias, especially in later steps when samples approach the manifold. These findings support the effectiveness of the proposed noise level correction strategy. Results for constrained sample generation can be found in Appendix C.3 The results show that DDIM-NLC achieves significantly smaller distance estimation bias than DDIM, particularly in the later steps as samples approach the manifold. As illustrated in fig. 2 by correcting noise level cstimation d,, DDIM-NLC with more accurately estimates the distance. The results of constrained sample generation are shown in Appendix C.3
10
distg(x)) of DDIM distg(x,) of DDIM-NLC
0.0
Distance Estimation Blas
10
vnnd
-0.2
Value
-0.4
10
-0.6
10
DDIM DDIM-NLC
0.8
6
8
10
6
10
Sampling Step
Sampling Step
(a) Comparison between (Rs) and de
(b) Distance Estimation Blas
Figure 4: Results from the toy experiment. (a) Sample-to-manifold distance time for DDIM and DDIM-
NLC. The dashed line shows the predefined noise level vno,, (b) Relative bias in distance estimation.
4.2 Unconstrained Image Generation
We conducted experiments evaluate the elfectiveness of noise level correction in unconstrained image generation tasks. The noise level correction network was trained on top of pre-trained denoiser network. Notably, the noise level correction network is approximately ten times smaller than the denoiser network. Additional details the experimental setup are provided in Appendix D.1 We used the FID (Fréchet Inception Distance) score as the evaluation metric to assess the quality of generated samples, where lower scores indicate better quality, following standard practice in image generation tasks (Heusel et al. 2017). The experimental results for the DDIM/DDPM framework on the CIFAR-10 dataset are shown in Table As observed, DDPM-NICC and DDIM-NIC, which incorporate noise level correction, outperform the original DDPM and DDIM models across all sampling steps. Specifically, our proposed noise level correction approach improves DDIM performance by 32%, 31%, 22%, 33%, and 28% for 10, 20, 50, 100, and 300 sampling steps, respectively.
We evaluate the effectiveness of noise level correction with EDM Karras et al. 2022), as it achieves state- of-the-art sampling quality with few sampling steps. Specifically, we assess the impact of NLC using both a first-order Euler ODE solver and a second-order Heun ODE solver within the EDM framework. As shown in Table / the proposed NLC also enhances the performance of EDM-basod sampling methods. Notably, as a robust sampling technique, noise level correction improves the performance of the Heun sampler by 10% with just 13 sampling steps.
10

---

Published in Transactions on Machine Learning Research (6/2025)
Table FID on DDIM/DDPM sampling on CIFAR-10 with and without noise level correction.
Table 2: FID on EDM sampling on CIFAR-10 with and w/o NLC.
Sttp
1000
300
100
50
20
10
Method Step
35
21
13
DDPM
2.99
2.95
3.37
4.43
10.41
23.19
Euler
3.81
6.29
12.28
DDPM-NLC
2.35
2.21
2.39
2.74
6.44
19.27
Euler-NLC
2.79
4.21
8.17
DDIM
4.29
4.32
4.66
5.17
8.25
14.21
Heun
1.98
2.33
7.22
DDIM-NLC
3.11
3.11
3.12
4.04
5.66
9.61
Heun-NLC
1.95
2.22
6.56
4.3 Image Restoration
In this section, we evaluate the effectiveness of noise level correction on five common image restoration tasks, 4x super-resolution (SR) using bicubic downsampling, deblurring with a Gaussian blur kernel, colorization using an average grayscale operator, compressed sensing (CS) with Walsh-Hadammrr sampling matrix a 0.25 compression ratio, and inpainting with text masks. These experiments are conducted on the ImageNet (Deng et al.] 2009) and CelebA-HQ (Karras et al.) 2018) datasets. We compare
method
with recent state-of-the-ars diffusion- based image restoration methods, including ILVR (Choi al. 2021 RePaint (Lugmayr et al. 2022), DDRM (Kawar et al. 2022). and DDNM (Wang et al. 2023 For a fair comparison, all diffusiion based methods utilize the same pretrained denoising networks with the same 100- step denoising process (100 number inference steps). following the experimental setup (Wang et al. 2023 To evaluate sample quality, we use FID, PSNR (Peak Signal-to-Noise Ratio), and SSTM (Structural Similarity Index Mcasure). For colorization, where PSNR and SSIM are less effective metrics (Wang al. 2023), we additionally
Consistency metric, denoted as "Cons' and calculated as || -yll As baseline, also include the inverse solution for each image restoration task, given by f A A which achieves zero constraint violation but lacks the data manifold information.
The results on the ImageNet dataset are summarized in Table 3 while those for CelebA-HQ are shown in Table Tasks not supported by certain methods are marked as "N/A." As the results indicate, integrating noise level correction (as in Algorithm 4 enhances sample generation performance for DDNM. Furthermore, the proposed IterProj-NLa method (Algorithm 2 achieves the best performance across all benchmarks. For instance, IterProj-NI.C outperforms the baseline DDNM in FID score by 6%, 59%, 14%, and 50% 4x SR, Deblurring, CS 25%, and Inpainting tasks, respectively. It also improves Consistency in colorization by 9%. Qualitative comparisons are shown in Figure with additional results in Appendix
Table Comparative results of five image restoration tasks on ImageNet.
ImageNet. Method
SR
Deblurring PSNRt/SSIM//FII
Colorization ConsL/FIDJ
25%
Inpainting
PSNR1
PSNRt/SSIMY/ /SSS///STTIIII
PSNR /SSTM///IIIIII
PSNR1 /STT///ST///TT. /FTD.
Aly
24.26 / 0.684
134.4
18.46 0.6616 55.42
0.0 / 43.37
15.65
0.510 277.4
14.52/ 0.799 72.71
IVR
27.40
0.870
43.66
N/A
N/A
N/A
N/A
RePaint
N/A
N/A
N/A
N/A
31.87
0.968
13.43
DDRM
27.38
0.869
43.15
43.01
0.992
1.48
260.4 36.56
19.95
0.704
97.99
31.73
0.966
10.82
DDNM
27.45
0.870/
39.56
44.93
0.993
1.17
42.32
36.32
21.62
0.748
64.68
31.60
0.946
9.79
DDNM-NLC
27.50
0.872
37.82
46.20
0.995
41.60
35.89
21.27
0.769
58.96
32.51
0.957
IttpPro-NNCC
27.56 / 0.873 37.48
48.24
0.997
0.48
38.30/
35.66
22.27
0.771
55.69
33.58
0.966
4.90
Table 4: Comparative results of five image restoration tasks on Celeba-HQ.
Celeba-HQ Method
4 SR
Deblurring
Colorization ConsL
CS 25%
Inpsinting
PSNRt/SS /St///SSIMMIIII
PSNRt/SSIMt /FIDJ
PSNRT/SSIMF /FID
PSNR PSNRt/SSIMt/ /FIM
Aly
27.27
0.782 103.3
18.85 / 0.741 / 54.31 N/A
0.0 / 68.81
15.09 0.583 / 337.7
15.57 / 0.809 181.56
ILVR
31.50
/ 0.945 29.82
N/A
N/A
N/A
RePaint
N/A
N/A
N/A
N/A
35.20
0.981
/18.21
DDRM
31.63
0.945
31.04
43.07
0.993 / 6.24
455.9 / 31.26
24.86
0.876
46.77
34.79
0.978
/16.35
DDNM
31.63
0.945
22.50
46.72
0.996
1.42
26.25 / 26.78
27.52
0.909
28.
36.64
/ 0.979
12.21
DDNM-NLC
31.78
0.947
22.10
0.997
1.36
24.92
25.81
27.63
0.914
24.72
36.48
0.980
/ 11.60
ProCPro-NLCC
31.93
0.949
21.96
46.97
0.997
1.29
24.65
/ 25.30
27.78
0.916
23.45
36.57
0.981
11.07
11

---

Published in Transactions on Machine Learning Research (6/2025)
4.4 Lookup Table for Noise Level Correction
In this section, we explore the statistical properties of the noise level correction network and demonstrate how these propertics can be leveraged create a lookup table for correcting noise levels without neural network inference. The lookup table for noise level correction defined 6, = fr] where 6 is non-parametric function that approximates the actual distance to the data manifold. As illustrated in the toy experiment shown in Figure 4b distance estimation error using noise levels is lower in the initial sampling steps and increases in later stages when the true distance to the manifold decreases. This trend is expected: in the early stages, noisy samples are farther from the manifold, making approximate projections easier and reducing relative distance estimation error. More specifically, at the initial steps, the true distance (x;) is slightly larger than the estimation from the noise level vnooe as supported by oq. (22). In later steps, however, d (at) decreases rapidly than Vnor..
0.2
0.2
Colorization
0.1
Imagenet
0.1
Inpainting
-0.1
(ooo))
-0.1
-0.2
-0.2
0.3
0.4
005 10a
-0.
10"
10"
10'
10'
dt
Gt
(a) Unconstrained sample generation
(b) Constrained sample generation
Figure 5: Plot of ro())) versus d, in the unconstrained DDIM-NLC denoising process and constrained DDNM- NT. denoising process. The curve represents the average over samples, with shaded regions indicating the standard deviation. The larger variance (right) illustrates that the corrections applied by ro(a)) are too complex for a simple look-up table in the context of constrained generation.
Table 5: FID DDIM sampling CIFAR-10 with lookup table noise level correction.
Table 6: FID on EDM sampling CIFAR-10 with and with LT-NLC.
Method Step
1000
300
100
50
20
10
Method's Steep
35
21
13
DDIM DDIM-LT-NLC
4.29
4.32
4.66
5.17
8.25
14.21
Heun Heun-LT-NLC
1.98
2.33
7.22
4.01
3.97
3.83
4.37
6.54
11.21
1.97
2.27
6.84
DDIM-NLC'
3.11
3.11
3.12
4.04
5.66
9.61
Heun-NLC
1.95
2.22
6.56
We conducted an experiment to analyze the statistical behavior of the neural network-based noise level cor- rector (-() for unconstrained sample generation on the CIFAR-10 and ImageNet datasets. Figure 5airee5aap presents the relationship between ra(a,) and ot, averaged over the samples 2, during the DDIM-NLC denoising pro- the final denoising steps (higher
cess. As scen, ro(a,) values are negative for smaller a, corresponding time steps t), and increase as 0, increases. This trend aligns with the observation in the toy experiment Figure 4b] indicating that distance decreases in the final steps and thus requires reducing ô, for accurate distance representation. Moreover, similar trend is observed across different datasets, such as CIFAR-II 10 and ImageNet. We further analyzed the statistical behavior of ro(-) in constrained sample generation tasks on the ImageNet dataset. These tasks rntuee additional variability due to constraint projections eq. (20) resulting in higher variance in re(.) across samples. Notably, even within the same dataset, constraints such as colorization and inpainting exhibit distinct trends during the final denoising steps (i.c., small noise levels). Moreover, the variance small noise levels is substantially higher in constrained tasks compared unconstrained scenarios.
12

---

Published in Transactions on Machine Learning Research (6/2025)
Using the values of re(ar) recorded in the average value curve of Figure 5 created lookup table-based noise level correction (I.T-NI.C) search ft to estimate ro(at). We evaluated the effectiveness of L.-NNCC in unconstrained sample generation tasks. The experimental results for LT-NLC applied to the DDIM framework on the CIFAR-10 dataset are shown in Table 5 As expected, the trained noise level correction (NI.C) achieves the best performance. However, LT-NIIC also significantly improves the original DDIM, enhancing performance by 14%, 20%, and 15% for 10, 20, and 50 sampling steps, respectively. The results for LT-NLC applied to the EDM framework on the CIFAR-10 dataset are presented in Table f Similar to the DDIM results, LT-NLC improves the performance of EDM-baser sampling methods, demonstrating its effectiveness a. network inference-free enhancement. The results for constrained generation can be found in Appendix D.4 As illustrated in Figure 5 the variance of ro(o)) in constrained generation tasks, such as image restoration, is significanlly higher. Consequently, the performance improvements achieved by LT-NIC are smaller compared to those of the neural network-based NLC, as LT-NLC applies the same correction across all samples. Therefore, constrained image generation tasks, the ncural network-based NLC remains essential for achieving optimal performance.
The noise level correction (NLC) network is significantly smaller than the denoiser, resulting in minimal additional computational overhead. Detailed comparisons of the training and inference times for the proposed NLC method are provided in Appendix D.2
5 Conclusions
In this work, we explore the relationship between noise levels in diffusion models and the distance of noisy samples from the underlying data manifold. Building on this insight, propose novel noise level correction method, utilizing a neural network to align the corrected noise level with the true distance of noisy samples to the data manifold. This alignment significantly improves sample generation quality. We further extend this approach to constrained sample generation tasks, such as image restoration, within an alternating projection framework. Extensive experiments on both unconstrained and constrained image generation tasks validate the effectiveness of the proposed noise level correction network. Additionally, we introduce lookup table- based approximation for noise level correction. This paramcter-frve method effectively enhances performance in various unconstrained sample generation tasks, offering a compulationally efficient alternative to the neural network-based approach.
References
Nikos Arechiga, Frank Permenter, Binyang Song, and Chenyang Yuan. Drag-guided diffusion models for vehicle image generation. arXiv preprint Xrxv::2306.00 099355, 2023.
Arpit Bansal, Mong-Min Chu, Avi Schwarzschild, Sounyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 843 843-852, 2023.
Heinz H Bauschke and Jonathan M Borwein. On projection algorithms for solving convex feasibility problems. SIAM review, 38(3):367 426, 1996.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspec- tives. IEEE transactions on pattern analysis and machine intelligence, 35(8):1798-1888 2013.
Joe Benton, Valentin De Bortoli, Arnand Doucet, and George Deligiannidis. Nearly Sdn-linear convergence bounds for diffusion models via stochastic localization. In The Twelfth International Conference on Learn- ing Representations, 2024.
Hanqun Cao, Cheng Tan, Zhangyang Gao, Yilun Xu, Guangyong Chen, Pheng- Ann Heng, and Stan Z Li. A survey on generative diffusion models. IEEE Tronsactions on Knowledge and Data Engineering, 2024.
Ting Chen. On the importance of noise scheduling for diffusion models. arxiv preprint arXiv: 10972, 2023.
13

---

Published in Transactions on Machine Learning Research (6/2025)
Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Robotics Research, pp. 027836492411736668 2023.
Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Tlv:: Conditioning method for denoising diffusion probabilistic models. In 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pp. 14347 14356. IEEE, 2021.
Hyungjin Chung, Bycongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems, 35:25683-25696 2022.
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffu- sion posterior sampling for general noisy inverse problems. In The Fleventh International Conference on Learning Representations, 2023a.
Hyungjin Chung, Suhyeon Lee, and Jong Chul Ye. Fast diffusion sampler for inverse problems by geometric docomposition. arXiv preprint arxiv 03223099057555 05754,77444 2023b.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248 255. leee, 2009.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Aduances in neural information processing systems, 34:8780 8794, 2021.
Zehao Dou and Yang Song. Diffusion posterior sampling for lincar inverse problem solving: A filtering perspective. In The Twelfth International Conference on Learning Representations, 2024.
Charles Fefferman, Sanjoy Mitter, and Hariharan Narayanan. Testing the manifold hypothesis. Journal of the American Mathematical Society, 29(4):983-1049, 2016.
Ian Goodfellow, Jean Pouget-Abadic, Mchdi Mirza, Bing Xu, David Warde-Farloy, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural informalion processing systems, 27, 2014.
Tiankai Hang, Shuyang Gu, Xin Geng, and Baining Guo. Improved noise schedule for diffusion training. arXiv preprint ar Xrx:22407 03302977 2024.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to local nash equilibrium. Advances in newral infor- mation processing systems, 30, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Newral Information Processing Systems, 33:6840-6851, 2020.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible
behavior synthesis. In International Conference on Machine Learning, pp. 9902 9915. PMLR, 2022.
Tero Karras, Timo Aila, Samuli Laine; and Jaakko Lehtinen. Progressive growing of GANs for improved quality, stability, and variation, 2018.
Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35:26565-26577, 2022.
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models.
Advances in Newral Information Processing Systems, 35:23593 23606, 2022.
Dongjun Kim, Chich-Hsin Lai, Wei-Hsiang Liao, Naoki Murata, Yuhta Takida, Toshimitsu Uesaka, Yutong He, Yuki Mitsufuji, and Stefano Ermon. Consistency trajectory models: Learning probability flow ode trajectory of diffusion. ar Xíe preprint xrx::2310000222 2023.
14

---

Published in Transactions on Machine Learning Research (6/2025)
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. 2014.
Zhileng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations, 2021.
Adrian S Lewis and Jérôme Malick. Alternating projections on manifolds. Mathematics of Operations Research, 33(1):216-223) 2241, 2008.
Dongxu Li, Jumnan Li, and Steven Hoi. Blip-diffusion: Pre-trained subject representation for controllable
text-to-image generation and editing. Advances in Neural Information Processing Systems, 36, 2024a.
Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards non-asymptotic convergence for diffusion-based generative models. 2024b.
Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion-lm improves controllable text generation. Advances in Newral Information Processing Systems, 2022.
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling around 10 steps. Advances in Neural Information Processing Systems, 35:5775- -77777 2022.
Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Lue Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CV conference compuler vision and pallern recognilion, pp. 11461 11471, 2022.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pp. 8162-877 8171. PMTR, 2021.
Mang Ning, Enver Sangineto, Angelo Porrello, Simone Calderara, and Rita Cucchiara. Input perturbation reduces exposure bias in diffusion models. In Proceedings of the 40th International Conference on Machine Learning, pp. 26245 26265,, 2023.
Achraf Oussidi and Azeddine Elhassouny. Deep generative models: Survey, In 2018 International conference on intelligent systems and computer vision (ISCV), pp. 1-8. IEEE, 2018.
Frank Permenter and Chenyang Yuan. Interpreting and improving diffusion models from optimization perspective. In Forty-first International Conference on Machine Learning, 2024.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.00 0612222,0 1(2):3 2022.
JH Rick Chang, Chun-Liang Li, Barnabas Poczos, BVK Vijaya Kumar, and Aswin C Sankaranarayanan. One network to solve them all-solving linear inverse problems using deep projection models. In Proceedings of the IEEE International Conference on Computer Vision, pp. 5888 1888-5897, 2017.
Robin Rombach, Andreas Blattmann, Dominik Torenz, Patrick Faser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 10684-10695, 2022.
Amirmojtaba Sabour, Sanja Fidler, and Karsten Kreis. Align your steps: Optimizing sampling schedules in diffusion models. arXio preprint Xrx:222444 14507, 2024.
Subham Sahoo, Aaron Gokaslan, Christopher M De Sa, and Volodymyr Kuleshov. Diffusion models with learned adaptive noise. Advances in Neural Information Processing Systems, 37:105730-105779, 2024.
Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models, In International Conference Learning Representations, 2022.
15

---

Published in Transactions on Machine Learning Research (6/2025)
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pp- 2256-2266 PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference Learning Representations, 2021a.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b.
Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. In International Conference on Learning Representations, 2022.
Yang Song, Prafulla Dhariwal, Mark Chen, and llya Sutskever. Consistency models. In International Conference Machine Learning, pp. 32211-32252. PMLR, 2023.
Shengkun Tang, Yaqing Wang, Caiwen Ding, Yi Liang, Yao Li, and Dongkuan Xu. Adadiff: Accelerating diffusion models through step-wise adaptive computation. In European Conference on Computer Vision, pp. 73 90. Springer, 2024.
Jianyi Wang, Zongsheng Yue, Shangchen Zhou, Kelvin CK Chan, and Chen Change Loy. Exploiting diftusion prior for world image resoutiino. International Journal of Computer Vision, pp. 1-21, 2024.
Yinhuai Wang, Jiwen Yu, and Jian Zhang. Zero-shot image restoration using denoising diffusion null-space model. In The Eleventh International Conference on Learning Representations, 2023.
Luhuan Wu, Brian Trippe, Christian Naesseth, David Blei, and John P Cunningham. Practical and asymp- totically exact conditional sampling in diffusion models. Advances in Newral Information Processing Sys- tems, 36, 2024.
Xingyu Xu and Yuejie Chi. Provably robust score-based diffusson posterior sampling for plug-and-pla, image reconstruction. Xr preprint orXiv:2403.1100. 2024.
16

---

Published in Transactions on Machine Learning Research (6/2025)
A Insights for the Noise Level Correction
A.1 Discrepancy between Vno, and the Distance (x4)
Previous works (Ho et al., 2020 Karras et al. 2022) have primarily focused on improving the estimation of the noise vector (3 Fe) 6 eq.
( aligning with the first assumption of zero-error denoiser. However, these approaches typically rely on predefined noise level scheduler o, for distance estimation, often without explicit validation. This reliance can lead to inaccuracies, even at the initial step, where dislx(xT) # vnoT.
Consider the DDIM as an example. The initial noisy point or is sampled as follows:
JT
2T Vot =
07
11+11. .T, 2T ~~~ 1),
(21)
Let 2g = Proje( (2xe) E K. denotes the projection of Ie onto the manifold K. In the context of image manifold, we have ||اا|| > O. If (ET, mg) s O, the expectation of squared distance is given by:
E [[|ET - 26l1 = E |ly ... 2 - =
(o7 + 1| + | 2 0f 1(T,
z {o} + 1) 22 + E[[=ū|
= (o) + 1)m + =||+=ا|? > oin
(22)
where the final equality uses the fact that E E |2] = for for 2r NN (mx)). Equation (22 implies that, with high probability, distx( VnoT. Consequently, at any step T, .., 0, these deviations may result in (xs # Vnoi, potentially causing the final sample to deviate from the manifold K. These errors arise from imperfections in the denoiser and become particularly significant in the final steps. The toy experiments on the discrepancy between noise level and true distance to the manifold can be found in appendix C.2 It is worth noting that the EDM sampling method initializes with a random step as TT = 0T 2T. However, leads to the same conclusion: for (2T, 2τ) I O, we have distx((T) # vnor.
A.2 Training Objective
To address the issuc of inaccurate distance estimation causod by relying on a predefined noise level de, we propose method called noise level correction (NLC) to better align the estimated noise level â with the true distance distxx Specifically, train neural network minimize the following objective
(initial form):
L L: -iiorr
(23)
Inspired by residual learning, instead of directly predicting the distance, we propose to learn the residual between the true distance and the predefined noise level. We define:
8, = o. [1 + ro(ox, distx(xx)
(24)
Vn
where (0(2),))) 0.) is the residual correction network to be trained. This residual formulation is advantageous because re(:) is typically more stable across noise levels, while oe may vary widely during large diffusion time steps.
Since computing the true distance (xs is impractical, we approximate it by the distance between a noisy sample and its clean counterpart obtained from the forward diffusion process. Given ot = Zo + dtC, with zo IC and e ~ No I), we approximale:
dastx((x;) R || zoll-
(25)
17

---

Published in Transactions on Machine Learning Research (6/2025)
This approximation is reasonable under the manifold hypothesis, assuming that the noise vector c is orthog- onal to the manifold, and that projx(#t)) R s0. To further improve robustness, we introduce perturbation scaling factor À and expand the input-output domain of ra(-) via:
24 + orde, c~,(0,,), A~ - -=
(26)
where å controls the perturbation range (e.g., 8 = recovers the original diffusion process). This leads to the approximation:
(1.) Re | zoll = arA||e|l-
(27)
Combining eq. eq.
23
24
and eq.
(7 arrive at the final training objective, shown in eq. 14
Lr, == Emosara [(Vno,[1 + ro(a,, ox
(28)
8, 3o oode, N),, A u(1 +)
(29)
During training, we sample zo from the training set, t from a uniform diffusion time range (e.g., t [1, 1000)), and draw ~ NO, I) and 6,,16, 1+)). The loss computed according oq. (14) and used to optimize the noise level correction network ro(-)-
B Sampling with Noise Level Correction
The proposed noise level correction (NLC) method is broadly compatible with wide range of existing diffusion sampling algorithms and can be seamlessly integrated to enhance sample quality. In many diffusion samplers, each denoising step consists of two substeps: (1) estimating the original clean sample ōoj given noisy obscrvation It, and (2) adding noise to propagate the sample to the next timestep. Our NLC method fits naturally into this structure by correcting the noise level during the clean sample estimation step, as described in eqs. [6) and 17)- In sections 3.3 and 3.4 we demonstrated the integration of NLC with DDIM/DDPM and constrained sampling tasks. In this section, extend our discussion to other samplers and show how the NLC method can be incorporated. All modifications to the original algorithms are highlighted in blue.
Algorithm 3 shows how NLC is integrated into the EDM sampler (Karras et al. 2022 Unlike DDIM-NLC in algorithm 1 EDM-NLC does not normalize the noise vector c0(x,, ot), since EDM based on sccond-order ODE solver. Nevertheless, the NLC method be applied by modifying the noise level estimation step,
shown in blue.
Algorithm EDM with Noise Level Correction (EDM-NLC)
Input: Denoiser eo and noise level corrector To
Input: Noise scheduler {o,]T=
Output: samples Zo K.
Fr ore, é (0, I),
for t= == T, - 1,... do
θ: 0 + -=(
4: ô = o, d=1
di
5: & = ( σt)
6: It-1 = z, + (ō G.)&E
7:
if / > then
Ê c0( x(-1)
à 0.54, + 0.5
10: It-1 = + (b,-1 6.)
11: end
12: end for
18

---

Published in Transactions on Machine Learning Research (6/2025)
Algorithm | presents the incorporation of NLC into DDNM, constrained image generation algorithm for linear inverse problems of the form Az (Wang et al. 2023 This demonstrates that the NLC method can also enhance constrained sampling tasks.
Algorithm DDNM with Noise Level Correction (DDNM-NLC)
Input: Denoiser eo and noise level corrector To
Input:
Noise scheduler {or]í=1, randomness scale 7/,
Input: Linear dogradation operator A, and pseudo-inverse operator A for image restoration constraint C,
Input: Degraded image y,
Output: samples Zo KnC
1: JT Voz + 1- 2T, 2T I),
2: for t= T- 1..., do
3: à; = o + ro(
4: - = न: d:l1
ai
5: & = Veea (2,, (00///
6: fo It δgÊt
Toj A (I
A AtA) S0/
8: nooie = 11 de
d-1
o7 6.-1
9: Signaal = v 0,-1 onoiee
10: It-1 = 20/e + O signaates + Tnoseely, where We ~ I))
11: end for
Furthermore, the NLC method can be extended to other advanced samplers, such as DPM-Solver (ODE solver for Diffusion probabilistic models) (Lu et al. 2022). In algorithm 5 we present the integration of NLC into the second-order DPM-Solver. As with EDM-NLC, the primary modification involves adjusting the noise level using the NLC module. The DPM-Solver uses noise schedules {ar)== and {ar)== and defines the auxiliary paramcter == log(ar/o.). The inverse mapping from to denoted by tx("), and we define the interval h; A d These examples collectively demonstrate the flexibility of the proposed noise level correction strategy, making it broadly applicable enhancement across various diffusion model sampling framcworks.
Algorithm 5 DPM-SSllvr with Noise Level Correction (DPM-NT.C)
Input: Denoiser Eg and noise level corrector Te
Input: Noise schoduler {ar)T and {or)T=1
Output: samples To K.
1: FT OTC, 6 (0, T),
for = T,T 1,... do
δ. + σi [1 1 (z 01)]
Ac-1tAe
4
5: d, = dt d.
& = c0(It, σt)
at Zt 8z c -1 de
h2
C, colu,ō+ 0z)
1t 01-1 (e - -1)
I1-1 ar
10: end for
19

---

Published in Transactions on Machine Learning Research (6/2025)
DDNM-NLC: distance
100
10a
DDNM-NLC :: consistency
058
DDNM: consistency
10'
100
056
10a
-1.0
6.35
0.25
0.38
160.00
sampling steps
(a) Sphere datasets toy experiments
(b) Sample generation with lincar constraint
Figure 6: (a) Sphere datasets in toy experiments; (b) shows the results of sample generation with lincar
constraint Az = b.
C
Toy Experiments
C.1 Experimental design for sphere manifold.
The dataset consists of samples from a d-dimensional sphere manifold, denoted as 8 St embedded within an n-dimensional data space. To create training samples x we apply m different lincar projections (rotations) to the original sphere and add a small amount of Gaussian noise Inoise the d-sphere signal Tsignal- Let K.
represent the d-sphere manifold, such that ssignal E K. Each training sample IR" is generated according to the following equations:
= Zeignal + Tnoise, Enoise ~ ~N(0010-01
(30)
Øsignal R&s RR, where K U((,, 21,11,,,.... ,m,), R& (nxx)
(31)
d+1
8 ~ St where
= 1, 8442 8643 = 8n =
(32)
Where N),,E) Σ) denotes the normal distribution, u({-}) denotes the uniform distribution. S refers to d -dimensional sphere. The matrices R1, R2,... "e,,.,,,,,, Rm are fixed random orthogonal matrices that are utilized to 'hide" zeros in certain coordinates. For our experiments, we focus primarily on a dataset with parameters n 100, 1, and m = 4, resulting a 100-dimensional dataset composed of circles. For illustration, also generate 3-dimensional dataset with parameters (n 3, 1, m 4), as shown in Figure (aa A total of 10,000 samples were used to train the model. The denoiscr co(-)is implemented with fully connected network containing 5 layers, each with hidden dimension of 128. For noise level correction, ra(-) is implemented with a 2-layer fully connected network, also using a hidden dimension of 128.
C.2 Misalignment between noise level and true distance to the manifold
As discussed in section 3.1 diffusion-based sample generation methods, such DDIM and DDPM, often suffer from cumulative errors introduced during the denoising process. These errors arise from imperfections in the denoiser and become particularly significant in the final steps. As a result, there can be a substantial mismatch between the true distance to the data manifold, d and the estimated noise level, Vnoa. To visualize this phenomenon, we conduct a toy experiment. Figure 7 shows the deviation between the true distance (á,) and the estimated noise level Vno, during a 10-step DDIM sampling process. As observed, the actual distance increasingly deviates from the predefined noise level, especially in the later steps. This result aligns with the noise estimation inaccuracy illustrated earlier in fig.
20

---

Published in Transactions on Machine Learning Research (6/2025)
10a
10a
Value
100
distx(X))
vnoa
2
6
8
10
Sampling Step
Figure 7: Deviation between the true distance (ās) and the estimated noise level vno, in DDIM sampling.
C.3 Results of constrained sample generation.
In Section 4.1 demonstrate the effectiveness of noise level correction in unconstrained sample generation. Here, we extend the evaluation to constrained generation, using a linear constraint Ar = b with A E R lxx as a random variable and b = 0. We applied the proposed 10-step DDNM with Noise Level Correction (DDNM-NLC), as detailed Algorithm to generate samples and compared this method to the 10-step DDNM baseline. We report the distance (lo measure sample quality) and Consistency error ||Ar - bll (lo measure constraint satisfaction) as shown in Figure 6b The proposed method shows superior performance in both metrics, generating high-quality samples that satisfy the constraint more effectively than the baseline.
D Image Generation Experiments
D.1 Experimental design
Table 7: Architecture of the noise level correction network r-(-)
Network
Hyperparameters
CIFAR-10
ImageNet
Celeba-HQ
Input Channel (feature channel)
256
1024
512
Input Size (feature size)
4
8
8
Network Architecture for ro(-)
Residual Blocks
2
2
Attention Blocks
1
1
Atlention Heads
4
ro())
14 M
234 M
59 M
Parameters
c+(-)
(Prozen, not trained)
218 M
2109 M
434 M
Table 8: Hyperparameters used for the training
Sottings
CIFAR-10
ImageNet
Celeba-HQ
Batch size Learning rate
128
64
Hyperparameters
0.0003
0.0003
0.0003
Iterations
150 K
300 K
300 K
Training ro(-)
19.2 M
19.2 M
19.2 M
# Samples
Training ep(.)
200 M
2500 M
~ 1000 M
Implementation. In this exporiment, we used pretrained denoiser cp(-), and trained the noise level correction network ro(-) enhance the denoising process. For unconstrained image generation experiments, employed the pretrained 32 32 denoising network from (Song et al.) 2021a) in DDIM-based experiments
21

---

Published in Transactions on Machine Learning Research (6/2025)
the CIFAR-10 dataset, and the pretrained 32 32 denoising network from (Krrras et al. 2022 in EDM- based experiments on CIFAR-10. For image restoration experiments, we utilized the pretrained 256 x 255 denoising network from (Dhariwal & Nichol 2021 for ImageNet experiments and the pretrained 256 x 255 denoising network from Lugmayr et al. 2022 for CelebA-HQ experiments.
Network architecture. The noise correction network ro(+) is designed to be significantly smaller than the denoiser eØ(-) while still incorporating residual and attention blocks similar to those used in the original DDPM denoiser (Ho et al.) 2020). Table 7 outlines the architecture of the noise correction network for CIFAR-10, ImageNet, and CelebA-HQ datasets. For all datasets, employ 2 residual blocks, 1 attention block, and 4 attention heads. The primary architectural difference lies in the varying feature sizes (input channels and input dimensions) generated by the denoiser's encoder. For comparison, the ImageNet denoiser contains 9 attention blocks in the encoder and 13 in the decoder. Consequently, the noise correction network is approximately ten times smaller than the denoiser network. The last two rows of Table provide a parameter count comparison between the noise correction network re and the denoiser dg.
Training. The noise correction network is trained following the original DDPM training procedure. Key hyperparameters for training are listed in Table Notably, the noise correction network is trained until 19.2 million samples have been drawn from the training set. In comparison, training the denoiser requires 200 million samples for CIFAR-10 and over 2000 million samples for ImageNet.
D.2 Time and Memory Cost
Table Inference Time Overhead of the Proposed NL.C Method.
model
CIFAR10 - DDPM CIFARI0-EDM CelebA DDNM ImageNet-DDDNM
Baseline Inference (s)
0.035
0.022
0.049
0.102
+NL.C (s)
0.037
0.023
0.050
0.104
Time Overhead Ratio(%)
5.7
4.3
2.0
2.0
Table 10: Inference Time Comparison Between the Baseline and the Proposed NLC Method. For the CIFAR-10 experiment, the baseline is DDPM and the proposed method is DDPM+NLC. For the ImageNet experiment, the bascline is DDNM and the proposed method is DDNM+NLC (inpainting experiments).
model
CIFARIO-DDPM (Song et al. 2021a ImageNet DDNM Dhariwal & Nichol 2021
baseline
0.32
0.93
bascline + NLC
0.34
0.95
Training time. The training time is shown in the last two rows of Table g faster training for the noise level correction network. For example, training the ImageNet denoiser on 8 Tesla V100 GPUs takes approximately two weeks, while training the noise correction network requires only about one day. This efficiency due the smaller size of the noise level correction network and its use of the pretrained denoiser.
This results in significantly
Inference time. The proposed noise level correction (NLC) method requires only one additional inference step for the NLC network re(-) per overall inference. This extra step is independent of the specific sampling algorithm, applicable to both unconstrained and constrained approaches, as shown in line 3 of algorithms 1 2 Band i Table summariies the additional time cost in different settings, indicating that each inference with the proposed NLC incurs an overhead of less than 6%. Furthermore, Table I0 presents the inference times for a real 10-step sample generation with a batch size of 1 for CIFAR-10 10 unconstrained sample generation with DDPM and ImageNet constrained (inpainting) generation with DDNM, confirming that incorporating the noise level correction results in modest increase of approximately 6% in inference time.
22

---

Published in Transactions on Machine Learning Research (6/2025)
Table 11: FID score on the scaling factor in training r ot). Here, = corresponds to the naive
setting (A = 1). CIFAR-10 experiments use 100-step sampling with DDIM-NLC, and ImageNet inpainting
experiments use 100-step sampling with DDNM-NLC.
0
0.2
0.5
1.0
CIFAR-10 (DDIM-NLC)
3.52
3.23
3.12
3.14
ImageNet Inpainting (DDNM-NLC)
7.87
7.59
7.20
7.24
D.3 Ablation Study on the Scaling Factor
In the training objective for the NLC nelwork re(-), we introduce the scaling factor À adjust the input- output range, Specifically, we use the predicted noise level vno,, level (i.c., the distance from the noisy sample to the clean sample) ||æ, zoll- This adjustment addresses the mismatch between the noise scales in the forward (diffusion) and reverse (denoising) processes. Noisy samples are generated as f To o,Ye,, where ~ NO0 1) and typically ||<ell ~ vn. Without scaling (i.e., setting À = 1), the true noise level becomes concentrated around oz o,, leading to ro( o,) R 0. However, the denoiser network 60 o0( does not necessarily predict a noise level with such a concentrated norm. This norm mismatch, also noted in (Permenter Yuan 2024 motivates our introduction of A.
[1+ [1
+ r ot)] to estimate the true noise
By allowing À to vary, e.g., À U(1 ro(f,, o.) reflects the variation in À (i.e., roughly in [-ô,ō]). To validate this design, we conduct an ablation study comparing different values of using 100-step sampling in terms of FID score, shown in table [ For CIFAR-10, the DDIM-NLC method is used for unconstrained sampling, while for ImageNet inpainting experiments the DDNM-NLC method is employed. As can be seen, = 0.5 (the standard setting of our experiment) outperforms the naive setting of 6 0.
+ 6), the true noise level becomes approximately λory and
D.4
Lookup table method for image restoration
Table 12: Comparative results of image restoration tasks on ImageNet for lookup-table noise level correction.
ImageNet
x
Deblurring
Colorization Cons
CS 25%
Inpainting
Method
PSNR1/ /SSRR//IISM//III
PSNR1 /SSIM" FID
PSNR1 /SSIM//IIDMII /FID
PSNR1 /SS////SI/////
DDNM
27.45
0.870/ 39.56
44.93
0.993
1.17
42.32
36.32
21.62
0.748
64.68
31.60
0.946 9.79
DDNM-LT-NLC
27.47 / 0.870/ 39.03
45.14
0.993
1.02
42.12
36.07
21.48
0.751
/
62.64
31.84 / 0.951 9.19
DDNM-NLC
27.50
0.872 37.82
46.20
0.995
0.79
41.60 / 35.89
21.27
0.769
58.96
32.51
0.957 7.20
We evaluate the performance of the lookup table noise level correction (LT-NLC) in image restoration tasks. The results the ImageNet dataset are summarized in table 12 As shown, the DDNM image restoration method achieves additional performance gains with LT-NLC. However, these improvements are notably smaller compared to those achieved with the neural network- NLC method.
D.5
Additional experiments integrating NLC with DPM
Table 13: FID scores on CIFAR-10 using DPM sampling with and without NLC, across varying numbers of function evaluations (NFE).
Method NFE NFF
12
18
24
30
36
DPM
5.28
3.43
3.02
2.85
2.78
DPM-NCC
4.83
3.22
2.97
2.81
2.77
As described in algorithm we intograte the proposed Noise Level Correction (NLC) method into the widely used DPM-Solver sampling algorithm. To evaluate the effectiveness of this integration, we conduct experiments on the CIFAR-10 image generation task. Table 13 presents the FID scores for DPM with
23

---

Published in Transactions on Machine Learning Research (6/2025)
and without NLC under different numbers of function evaluations (NFE). As shown, incorporating NLC consistently improves sample quality across all settings. For instance, with NFE 12, NLC reduces the FID score from 5.28 to 4.83, corresponding to a 9% improvement. These results demonstrate that NLC can effectively enhance the performance DPM-based sampling methods.
E
Qualitative study
E.1
Image Restoration
We present qualitative comparisons between the proposed method, TterProj-NT.C, and the baseline, DDNM, across various image restoration tasks. These tasks include compressive scnsing, shown in fig. g colorization, shown in fig. g inpainting, shown in fig. T0. and super-resolution, shown in fig. [1 The comparisons demonstrate the effectiveness of IterProj-NLC in producing visually superior results over the baseline.
Compressive sensing 25%
Compressive sensing 12.5%
Degraded
DDNM
Figure Qualitative results compressive sensing.
Colorization
Degraded
DDNM
Ours
Figure 9: Qualitative results of colorization.
24

---

Published in Transactions on Machine Learning Research (6/2025)
Inpainting
Degraded
DDNM
Ours
Figure 10: Qualitative results of inpainting.
Super resolution x8
Degraded
DDNM
Ours
Figure 11: Qualitative results of super resolution.
25

---

Published in Transactions on Machine Learning Research (6/2025)
E.2 Unconstrained Image Generation
The example results of CIFAR-10 generated using 100-step sampling of DDIM-NLC are presented in fig. 12
Figure 12: Example results of CIFAR-10 generated using DDIM-NT.C.
26


Reviews:
[{"review_id": "2LhDiis8Vm", "reviewer": "Reviewer_iHw6", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. The presented technique is intuitive and can be readily applied to existing samplers in diffusion models with minimal implementation overhead.\n\nWeaknesses / Requested Changes/Questions:\n\n1. The authors start with Assumption 3.1 but claim this assumption might not hold in diffusion models due to error propagation during the denoising process. However, except for the visually intuitive illustration in Fig.2, no empirical evidence is provided to support this claim. I think demonstrating the fallacy of this assumption in the context of a toy example could be very useful (let's say comparing ground truth distance dist(x_t) with its approximation \\sqrt{n}\\sigma_t). I see the experiment in Section 4.1, but to my understanding, the presentation in Section 4.1 relatively compares DDIM with DDIM + NLC.\n2. The motivation behind including the lambda parameter is unclear from the explanation in the main text. Is it primarily due to empirical convenience?\n3. Re. the results on unconditional generation in Table 2, the difference between the Heuns solver in EDM and the Heun’s + NLC solver is marginal at best. Does this imply that using higher-order solvers like DPM-Solver and its variants would not benefit much from this technique?\n4. In my opinion, the discussion of related work in Section 2.2 needs significant revision. The authors overlook a lot of work on improving noise schedules in diffusion models. Here are some references I found from a quick search [1,2,3], and the authors should do their own due diligence when updating this section, which would improve the paper quite a bit.\n[1] Diffusion Models With Learned Adaptive Noise, Sahoo et al.\n[2] Improved Noise Schedule for Diffusion Training, Hang et al.\n[3] On the Importance of Noise Scheduling for Diffusion Models, Ting Chen\n\nMinor Comments:\n\n1. The original image row is missing from Fig. 1. Including this would make this figure better.\n\n2. There seems to be a typo below the statement of Theorem 3.1. Notably the Tweedie’s estimate should be equal to x_0t = x_t - \\sigma_t e_\\theta(x_t,  t)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper leverages theoretical results on viewing diffusion model sampling as an optimization problem where the distance of a noisy sample from the data manifold can be approximated using the magnitude of the noise at any time t. The authors propose to correct this noise using a residual correction scheme that utilizes training an extra neural network (in addition to a pretrained denoiser) for modeling these residuals. Empirical results are demonstrated in the context of unconditional and conditional generation.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See the Weaknesses Section above", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "7MtYcqwcgV", "reviewer": "Reviewer_FJeZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n- The NLC network is designed as a modular, plug-and-play component, which, while incurring additional computational costs, integrates seamlessly into existing diffusion frameworks.\n\n- The use of a lookup table approximation offers an efficiency-oriented alternative. Although this method shows a slight performance trade-off compared to the full NLC network, it remains a practical option in scenarios where computational resources are limited.\n\n- Figure 2 clearly illustrates the advantages of the proposed noise level correction strategy. Figure 5 provides compelling evidence of the improvements during the denoising process. Experiments on several schedulers have been done to show the improvements from the NLC network.\n\n- The method not only enhances quantitative performance in unconstrained image generation but also benefits various image restoration tasks, including super-resolution, deblurring, and colorization.\n\n- The paper is well-organized and clearly written.\n\nCons:\n\n- A question regarding the adaptability of the NLC: it remains unclear whether the noise level correction approach can be extended to consistency models or integrated with other popular schedulers such as DPM. Addressing this point could provide further insights and broaden the applicability of the method.\n\nOverall, I think it is a solid paper to benefit the diffusion-based research works.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents a simple yet effective method to enhance the denoising process by introducing a compact noise level correction (NLC) network. Despite a modest increase in training and inference costs, the proposed approach significantly improves sample generation quality.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Since this work is a modification to the existing diffusion schedulers, it is good to show the adaptability of NLC to various schedulers and models.", "broader_impact_concerns": "No concerns on the ethical implications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ZQ1FWYiUvu", "reviewer": "Reviewer_A2wi", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The studied problem is relevant and quite interesting.\n- The paper is quite well written overall, most things are explained well.\n- The proposed approach makes sense overall, it is conceptually quite simple.\n- The proposed approach achieves at least small performance gains across all experiments.\n\n\n\n\nWeaknesses:\n- It is not entirely clear to me exactly how the noise level correction network is trained, I think this should be clarified a bit. In eq. (14), what is the expectation over? It is over $x_0, t, \\epsilon$ as in eq. (4)? I also think it could be clarified a bit further how the objective in eq. (14) is derived.\n- In particular, it is not clear to me why you need to introduce the scaling factor $\\lambda$ in eq. (14)? There are no ablation results showing how this affects the performance either? Is this $\\lambda$ crucial for the observed performance gains of the proposed approach? \n- The additional computational cost at training and test-time of the proposed approach is only briefly mentioned at the very end of Section 4, I think this should be discussed in some more detail. In Appendix D.2, test-time cost for DDIM-NLC vs DDIM is shown, but there is no comparison of the image restoration methods in Table 3/4. How does IterProj-NLC compare with DDNM-NLC and DDNM?\n\n\n\n\n\n\nMinor things:\n- Section 1, \"robotic path-planningand control\", typo.\n- Section 1, \"Previous studies Rick Chang et al. (2017); Permenter & Yuan (2024) have\", incorrect citation formatting.\n- Section 1: \"This perspective views the sampling process an optimization problem\" --> \"This perspective views the sampling process as an optimization problem\"?\n- Figure 1, \"Compressive Sensing\" --> \"Compressed Sensing\"?\n- Section 1, \"such as DDNM (Denoising Diffusion Nullspace Model) in Wang et al. (2023)\" --> \"such as DDNM (Denoising Diffusion Nullspace Model) by Wang et al. (2023)\"?\n- Section 2.1, \"can be found in Song et al. (2021b); Karras et al. (2022)\", citation formatting.\n- Section 2.1, \"the randomized DDPM Ho et al. (2020) follows the update rule\", citation formatting.\n- Section 3.1, \"Previous studies Rick Chang et al. (2017); Permenter & Yuan (2024) have established this connection\", citation formatting.\n- Section 3.1, \"Permenter & Yuan (2024) introduces\" --> \"Permenter & Yuan (2024) introduce\"?\n- First line after eq. (13), the expression of $\\hat{x}_{0|t}$ should match eq. (5)?\n- Section 3.1, \"the imperfections of the denoiser Ning et al. (2023)\", citation formatting.\n- Section 3.2, \"As illustrated in Permenter & Yuan (2024), this occurs\": I would do either \"As illustrated by Permenter & Yuan (2024), this occurs\" or \"As illustrated in (Permenter & Yuan, 2024), this occurs\".\n- Section 4.1, \"Theorom 1\", typo. \"as shown in ??\", missing reference.\n- Figure 4 caption, \"(b) Distance Estimation Error\" --> \"(b) Distance Estimation Bias\"?\n- 4.2, \"following standard practice in image generation tasks Heusel et al. (2017)\", citation formatting.\n- 4.3: \"RePaint ?\", missing reference.\n- Table 7, \"Heper\" --> \"Hyper\".", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors aim to improve the quality of generated samples by diffusion models.\n\nFor this, they train a noise level correction (NLC) network which is used to replace $\\sigma_t$ (from the predefined noise schedule) with a corrected $\\hat{\\sigma}_t$ in the denoising process. The NLC network utilizes the encoder of a pretrained denoiser network $\\epsilon$.\n\nThey get -NLC versions of the DDIM, DDPM and EDM sampling algorithms for unconditional image generation, for which they conduct experiments on CIFAR10.\n\nThey also incorporate their NLC approach in algorithms for constrained sampling, specifically applied to image restoration. For this, they do experiments on super resolution, deblurring, colorization, compressed sensing and inpainting, on both ImageNet and CelebA-HQ.\n\nTheir -NLC versions of the sampling algorithms consistently improve the generated sample quality at least to some extent across all experiments.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The studied problem is relevant and quite interesting. This could definitely be relevant for the TMLR audience.\n\nHowever, I think that some clarifications are required in the current version of the paper, see \"Weaknesses\" above.", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper focuses on improving sample generation in diffusion models. Motivated by the insight that the denoising process can be viewed as projecting noise samples onto the data manifold, the authors propose a noise level correction network to enhance denoising and demonstrate its benefits across a range of downstream tasks. Reviewers agree that the work has a clear and intuitive motivation (iHw6, FJeZ), shows improved performance across a broad range of tasks (A2wi, FJeZ), and can be readily applied to existing sampling algorithms in diffusion models (iHw6). Meanwhile, reviewers also raised concerns regarding the empirical evidence (iHw6), limited analysis of adaptability (FJeZ), and insufficient details on training settings (A2wi). The authors addressed these concerns thoroughly during the rebuttal phase, and all reviewers acknowledged that their concerns were resolved. Given the overall consensus, the AE recommends acceptance of this paper.
Decision: Accept (poster)