Disentangling Total-Variance and Signal-to-Noise-Ratio Improves Diffusion Models
Authors: Khaled Kahouli, Winfried Ripken, Stefan Gugler, Oliver T. Unke, Klaus-Robert Müller, Shinichi Nakajima
Organizations: BIFOLD – Berlin Institute for the Foundations of Learning and Data · Machine Learning Group, Technische Universität Berlin · Google DeepMind · Department of Artificial Intelligence, Korea University · Max-Planck Institute for Informatics · RIKEN Center for Advanced Intelligence Project
Diffusion models have been widely adopted as a powerful generative framework across numerous scientific domains, yet their inherently iterative sampling process remains a significant computational bottleneck. While this is typically mitigated by reducing the number of reverse diffusion steps, sample quality becomes highly sensitive to the chosen noise schedule, i.e. the manner in which noise is introduced and the signal is reduced at each step. Although prior work has improved on the original variance-preserving and variance-exploding schedules, these approaches passively adjust the total variance, without direct independent control over it. In this work, we propose a novel total-variance/signal-to-noise-ratio disentangled (TV/SNR) framework, where TV and SNR can be controlled independently. Our approach reveals that schedules with exponentially exploding TV can be consistently improved by adopting a constant TV schedule while preserving the underlying SNR schedule. Furthermore, we demonstrate that generalizing the SNR schedule of optimal transport flow matching significantly enhances generation efficiency. Crucially, these gains remain consistent across various reverse diffusion solvers and reduce the search space to two independent and interpretable parameters, enabling fast and robust sampling for complex scientific applications, such as molecular structure generation.
Figures & tables
Figure 1: Overview of our TV/SNR disentangled framework. Top: In standard diffusion models ( left ), the signal parameter a(t) and the noise parameter b(t) from equation 1 are mathematically coupled. Consequently, optimizing the signal-to-noise-ratio (SNR) schedule often forces a suboptimal Total Variance (TV), as defined in equation 9 , inducing high trajectory curvature. Our disentangled framework ( right ) breaks this dependency by controlling the diffusion process via two separate, independent schedules. This enables the design of an optimal Inverse Sigmoid SNR (ISSNR) schedule that can be independently paired with a constant TV schedule, resulting in our VP-ISSNR. Bottom: Numerical analysis of the simulated reverse ODE trajectories for a toy problem confirms this mechanism. Established methods, such as EDM (Variance-Exploding) and OTFM (Variance-Modulated), exhibit high trajectory curvature (green curves) near the data manifold ( t→0 ), where the support of the marginal distribution (red curves) is small. In contrast, our decoupled VP-ISSNR schedule achieves significantly lower trajectory curvature (total integrated curvature is shown inside the plots). This mitigates discretization errors and enables the generation of stable molecular structures in as few as 4 NFEs. A two-dimensional example with a mixture of 12 Gaussians forming two rings is given in figure 14 , and shows the same trends.
Figure 2: TV ( Left ) and SNR schedules ( middle ) of our ISSNR together with different established scheduling methods (see table 1 for the corresponding analytic expressions). Right: Our proposed inverse sigmoid SNR (ISSNR) schedule.
Figure 3: Stability rate ( higher is better ) as a function of the number of function evaluations (NFE) for molecular structure generation on the QM9 dataset. Top-right: Comparison between commonly established non-VP schedules (i.e., τ(t) non-constant) and their VP counterparts (i.e., τ(t)=1 ) in our framework. Bottom-left: Comparison of various baselines, including the VP analogs from the top-right plot, against our proposed VP-ISSNR schedule, with fixed parameters η=1.0 and κ=2.0 for all NFEs. Bottom-right: Same as the left plot but using the reverse SDE, instead of the reverse ODE. The best-performing ODE schedule, VP-ISSNR, is highlighted in black for reference. For reference, the top-left panel shows the sampled molecular structures obtained from the same prior sample using our approach, compared against the two widely used EDM and OTFM methods, as the number of NFEs increases. Each curve represents the mean across five independent sampling runs with different random seeds. The standard deviations are too small to be visually perceptible (median 0.6% ) and are provided in Table 5 .
Figure 4: FID score ( lower is better ) as a function of the number of function evaluations (NFE) for image generation on CIFAR-10 and FFHQ. Top-right: Comparison of existing non-VP models and their VP variants on CIFAR-10. Bottom-left: Comparison on CIFAR-10 between baseline methods and two variants of our proposed VP-ISSNR method, one using a scaled η in function of the NFE, and the other tuning η and κ for each NFE via Bayesian optimization, as detailed in section 3.4 . Bottom-right: Same comparison as in the left plot, but on higher-resolution FFHQ images. Both plots report the mean FID over five independent sampling runs with different random seeds. For reference, the top-left panel shows the generated images obtained from the same prior sample using our approach, compared against the two widely used EDM and OTFM methods, as the number of NFEs increases.
Schedule / NFE
3
7
15
31
63
CIFAR-10
EDM
468.3 ±0.1
99.66 ±0.25
4.55 ±0.06
2.02 ±0.03
2.05 ±0.04
OTFM
93.92 ±0.17
16.08 ±0.03
6.08 ±0.05
3.95 ±0.05
2.82 ±0.05
VP-ISSNR [η=1.5,κ=1.0] (ours)
322.7 ±0.2
130.2 ±0.3
4.67 ±0.05
2.11 ±0.03
2.02 ±0.03
VP-ISSNR [η scaled ] (ours)
94.39 ±0.28
74.56 ±0.30
7.31 ±0.06
2.07 ±0.03
2.03 ±0.03
VP-ISSNR [η,κ tuned ] (ours)
79.04 ±0.21
16.00 ±0.09
3.45 ±0.05
2.06 ±0.03
2.01 ±0.03
Table 2: FID ( lower is better ) on CIFAR-10 and FFHQ with Heun’s method as a function of the number of function evaluations (NFE). Each entry is the mean over five independent sampling runs with different random seeds, each computed from 50000 generated images, with ±1 standard deviation over the seeds as a subscript. Within each column, every value whose ±1 standard deviation interval overlaps that of the lowest FID is set in bold.
Appendix figures & tables17 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: The time discretization in equation 7 of the original EDM sampler [ 47 ] for various parameters ρ . At ρ=1 , linear time discretization is recovered, i.e. t(i)=i . For larger ρ , time steps get considerably shorter close to the data distribution. EDM uses ρ=7 .
Parameter
Distribution
Range
η
uniform
[0.5,4.0]
κ
uniform
[−3.0,3.0]
Appendix
Table 3: Search space for Bayesian Optimization of parameters η and κ
NFE 3
NFE 7
NFE 15
NFE 31
NFE 63
Dataset
Solver
η
κ
η
κ
η
κ
η
κ
η
κ
Bayesian optimization (VP-ISSNR [η,κ tuned ] )
CIFAR-10
Heun
1.04
−0.06
0.82
0.35
1.38
0.66
2.17
−0.47
2.96
2.61
CIFAR-10
Euler
2.15
1.05
0.97
0.26
1.00
0.00
0.88
0.13
0.92
0.29
CIFAR-10
DPM-Solver
2.24
−3.00
2.24
−2.35
3.86
−2.40
4.00
−0.22
3.95
2.97
FFHQ
Heun
1.76
−0.78
0.85
−0.05
1.31
0.68
1.76
0.01
1.79
−0.33
Appendix
Table 4: Hyperparameters of the ISSNR schedule for images. The upper block gives the values found by Bayesian optimization for each dataset, solver and NFE (VP-ISSNR [η,κ tuned ] ). The middle block gives the two untuned settings, which use the same values for every dataset. The lower separate block gives the values for the adaptive RK45 solver, whose NFE is set by the solver and is therefore reported as the average NFE over the generated samples.
Schedule
4
8
16
32
64
128
256
Reverse ODE
EDM
0.3 ±0.1
15.8 ±0.6
55.5 ±1.0
76.2 ±1.2
82.6 ±1.1
85.2 ±0.8
86.2 ±0.5
EDM-UT
0.0 ±0.0
0.0 ±0.0
43.0 ±0.9
81.3 ±0.5
86.1 ±0.8
87.6 ±0.7
87.8 ±0.6
SMLD
0.0 ±0.0
0.1 ±0.0
0.9 ±0.1
75.8 ±0.6
86.8 ±0.6
88.0 ±0.4
88.3 ±0.6
OTFM
2.5 ±0.7
28.3 ±0.5
65.3 ±0.8
79.9 ±0.6
84.9 ±0.6
86.3 ±0.9
86.9 ±0.9
VP-EDM-UT
0.0 ±0.0
42.9 ±0.6
72.4 ±0.7
81.1 ±0.6
84.6 ±1.0
86.1 ±0.5
86.7 ±0.5
Appendix
Table 5: Stability rate (%) on QM9 as a function of the NFE, for the reverse ODE and the reverse SDE with the Euler solver. Each entry is the mean over five independent sampling runs with different random seeds, with ±1 standard deviation over those seeds as a subscript. Within each column and ODE/SDE block, every value whose ±1 standard deviation interval overlaps that of the column’s largest value is set in bold. In the low NFE regime, our method is always outperforming the baselines. Binomial confidence intervals for these values are given in table 9 .
Schedule
4
8
16
32
64
128
256
Reverse ODE
EDM
0.0 ±0.0
7.4 ±0.5
46.7 ±0.9
69.2 ±1.3
76.8 ±1.1
79.7 ±1.0
81.0 ±0.7
EDM-UT
0.0 ±0.0
0.0 ±0.0
31.6 ±0.5
72.0 ±0.3
78.7 ±0.7
80.9 ±0.8
81.8 ±0.7
SMLD
0.0 ±0.0
0.0 ±0.0
0.5 ±0.1
64.3 ±0.6
78.1 ±0.6
80.7 ±0.5
81.8 ±0.5
OTFM
1.0 ±0.2
20.1 ±0.4
56.8 ±0.7
73.2 ±0.6
78.7 ±0.5
80.5 ±0.7
81.3 ±0.6
VP-EDM-UT
0.0 ±0.0
38.2 ±0.7
65.8 ±0.7
75.7 ±0.6
79.2 ±1.0
81.0 ±0.6
81.5 ±0.6
Appendix
Table 6: Validity (%) on QM9 as a function of the NFE, for the reverse ODE and the reverse SDE with the Euler solver. Each entry is the mean over five independent sampling runs with different random seeds, with ±1 standard deviation over those seeds as a subscript. Within each column and ODE/SDE block, every value whose ±1 standard deviation interval overlaps that of the column’s largest value is set in bold. In the low NFE regime, our method is always outperforming the baselines.
Schedule
4
8
16
32
64
128
256
Reverse ODE
EDM
—
99.9 ±0.2
99.6 ±0.3
99.3 ±0.1
99.4 ±0.1
99.3 ±0.1
99.3 ±0.0
EDM-UT
—
—
95.1 ±0.9
97.5 ±0.3
98.8 ±0.2
99.2 ±0.1
99.3 ±0.1
SMLD
—
—
—
93.9 ±0.6
98.3 ±0.3
99.1 ±0.3
99.2 ±0.1
OTFM
99.8 ±0.5
99.4 ±0.3
99.4 ±0.2
99.2 ±0.2
99.2 ±0.1
99.3 ±0.1
99.2 ±0.1
VP-EDM-UT
—
99.8 ±0.2
99.5 ±0.1
99.3 ±0.1
99.3 ±0.1
99.2 ±0.1
99.3 ±0.1
Appendix
Table 7: Uniqueness (%) on QM9 as a function of the NFE, for the reverse ODE and the reverse SDE with the Euler solver. Each entry is the mean over five independent sampling runs with different random seeds, with ±1 standard deviation over those seeds as a subscript. Within each column and ODE/SDE block, every value whose ±1 standard deviation interval overlaps that of the column’s largest value is set in bold. Entries with fewer than 50 stable structures per seed are excluded from the comparison.
Schedule
4
8
16
32
64
128
256
Reverse ODE
EDM
—
59.1 ±3.4
53.9 ±0.7
52.7 ±0.7
52.0 ±1.0
51.4 ±0.7
51.3 ±0.9
EDM-UT
—
—
58.2 ±1.6
53.4 ±1.4
52.7 ±0.4
52.7 ±1.4
51.7 ±1.2
SMLD
—
—
—
55.9 ±1.2
53.0 ±0.9
52.5 ±0.2
52.5 ±1.6
OTFM
63.6 ±4.8
52.0 ±2.1
51.6 ±0.8
51.4 ±1.2
51.2 ±1.2
51.2 ±0.5
51.1 ±0.2
VP-EDM-UT
—
56.1 ±1.3
54.1 ±0.7
52.8 ±0.3
52.3 ±0.5
51.8 ±0.8
51.3 ±0.8
Appendix
Table 8: Novelty (%) on QM9 as a function of the NFE, for the reverse ODE and the reverse SDE with the Euler solver. Each entry is the mean over five independent sampling runs with different random seeds, with ±1 standard deviation over those seeds as a subscript. Within each column and ODE/SDE block, every value whose ±1 standard deviation interval overlaps that of the column’s largest value is set in bold. Entries with fewer than 50 stable structures per seed are excluded from the comparison.
Schedule
4
8
16
32
64
128
256
Reverse ODE
EDM
[0.2, 0.4]
[15.2, 16.4]
[54.7, 56.4]
[75.5, 76.9]
[81.9, 83.2]
[84.6, 85.8]
[85.6, 86.8]
EDM-UT
[0.0, 0.0]
[0.0, 0.0]
[42.2, 43.9]
[80.6, 81.9]
[85.5, 86.7]
[87.0, 88.1]
[87.2, 88.3]
SMLD
[0.0, 0.0]
[0.0, 0.2]
[0.7, 1.1]
[75.0, 76.5]
[86.2, 87.3]
[87.4, 88.5]
[87.8, 88.9]
OTFM
[2.3, 2.8]
[27.5, 29.1]
[64.4, 66.1]
[79.2, 80.6]
[84.2, 85.5]
[85.7, 86.9]
[86.3, 87.4]
VP-EDM-UT
[0.0, 0.1]
[42.1, 43.8]
[71.6, 73.1]
[80.4, 81.8]
[84.0, 85.2]
[85.5, 86.7]
[86.1, 87.2]
Appendix
Table 9: Binomial 95% confidence intervals (Wilson score) of the stability rate (%) in table 5 .
Figure 6: Stability rate ( higher is better ) as a function of the number of function evaluations (NFE) for molecular structure generation on the QM9 dataset. Left: Second-order Heun sampler. Right: DPM solver [ 30 , 34 ] .
Figure 7: Root mean square deviation (RMSD, lower is better ) between the generated structures and reference structures obtained from geometry relaxations using DFT calculations at the B3LYP/6-31G(2df,p) level of theory (see D.4 for details), the same method used for generating the structures in QM9 [ 64 ] , which the model was trained on. We see a similar trend to the stability rate results in figure 3 , where our VP-ISSNR consistently outperforms other approaches at low NFE. This reveals that our method can generate physically plausible molecules that are structurally similar to ground truth reference structures.
Figure 8: Effect of the schedule used during training. Using the same schedule for both training and sampling does not enhance results for EDM, whereas the Cosine schedule emphasizes sampling in more challenging regions during training. Our VP-ISSNR sampling still outperforms other baselines, even when the model is trained using the EDM schedule.
Figure 9: Density of log-SNR values visited during training under the two training schedules considered in figure 8 . The cosine schedule concentrates the training budget in the intermediate log-SNR region, while the EDM schedule distributes it over a considerably wider range. Since training draws t uniformly on [0,1] , the density is ∣dlogγ/dt∣−1 , normalized to unit area.
Figure 10: FID score (lower is better) as a function of the number of function evaluations (NFE) in image generation on FFHQ, comparing existing non-VP schedules with their VP variants.
Figure 11: FID score (lower is better) as a function of the number of function evaluations (NFE) for image generation on CIFAR-10. Results are shown for combinations of solvers and noise schedules, including the VP-ISSNR variant with fixed hyperparameters in red. Solvers include Euler (first-order), Heun (second-order), RK45 (high-order), and the advanced DPM solver [ 30 , 34 ] . For the DPM solver, their proposed schedule is indicated in green. Due to the difficulty of precisely controlling the NFE for RK45, some data points may be missing or slightly shifted in the plot.
Figure 12: FID score (lower is better) as a function of the number of function evaluations (NFE) for image generation using different datasets with different image resolutions and the second-order solver Heun. For completeness, we add the results of using the fixed hyperparameters variant of our VP-ISSNR schedule in red.
Figure 13: ODE trajectories (black solid curves) and the marginal density path (orange shadows) when the data distribution is a mixture of three uniformly spaced delta peaks. The insets focus on the neighborhood of one of the peaks at x=23 and t≪1 , and the green and red curves indicate the local curvature and the support (relative to the one at t=1 ), respectively. On top of the green curve, we report the global curvature integrated over t .
Figure 14: Two-dimensional version of the toy analysis of figure 1 , where the target distribution pdata is a mixture of 12 isotropic Gaussians of equal weight, placed on two concentric rings (4 modes at radius 1 , 8 modes at radius 2.2 , all with standard deviation 0.03 ), and rescaled so that the mixture has zero mean and identity covariance. Left panel of each block: 96 trajectories of the reverse ODE (equation 4 with λ=0 ) started from the same prior samples. Right panel of each block: The position along the dashed ray, which passes through two picked modes, as a function of t , over the density of pt restricted to that ray, i.e. only the trajectories ending in the two circled modes are shown. The trajectories of VP-ISSNR enter the two modes early and along almost straight lines, those of OTFM contract towards the modes over the displayed interval, and those of EDM turn into the modes only at the very end. Bottom: Velocity deviation C(t) (green, equation 18 , summed over both dimensions) with its integral over t , and the support b(t)/b(1) of the marginal relative to its value at t=1 (red), as in figure 1 .
Diffusion model sampling studies how to generate high-quality samples efficiently by designing reverse processes, solvers, and inference-time controls.
We develop a principled framework for analyzing and designing noise schedules in diffusion models. We show that one can recast this design problem as an optimal control problem, whose state is the Fisher information of the diffusion process, which evolves according to an ODE, and whose control input is the noise schedule. The objective of the optimal control problem is a functional involving the Fisher information, which is shown to be an upper bound on the Kullback-Leibler sampling error. By solving this optimal control problem, we obtain sufficient conditions on noise schedules under which O~(d/n) sampling error is achievable, where d is the data dimension and n is the number of discretization steps. While existing theoretical works also prove that O~(d/n) sampling error bounds are achievable, these results hold for specific noise schedules that differ from those used in practice. Under additional assumptions on the data distribution, we derive a family of noise schedules that recovers the functional forms of schedules widely used in practice. We show that these schedules achieve the O~(d/n) error bound when their parameters are chosen according to our analysis. Our theory offers a principled approach to noise schedule design, yielding explicit functional forms and guidance for selecting their parameters.
Seo Taek Kong, Weina Wang, R. Srikant
ECE & CSL, University of Illinois Urbana-Champaign · Computer Science Department, Carnegie Mellon University · ECE, CSL & NCSA, University of Illinois Urbana-Champaign
Diffusion Probabilistic Models have demonstrated remarkable performance across a wide range of generative tasks. However, we have observed that these models often suffer from a Signal-to-Noise Ratio-timestep (SNR-t) bias. This bias refers to the misalignment between the SNR of the denoising sample and its corresponding timestep during the inference phase. Specifically, during training, the SNR of a sample is strictly coupled with its timestep. However, this correspondence is disrupted during inference, leading to error accumulation and impairing the generation quality. We provide comprehensive empirical evidence and theoretical analysis to substantiate this phenomenon and propose a simple yet effective differential correction method to mitigate the SNR-t bias. Recognizing that diffusion models typically reconstruct low-frequency components before focusing on high-frequency details during the reverse denoising process, we decompose samples into various frequency components and apply differential correction to each component individually. Extensive experiments show that our approach significantly improves the generation quality of various diffusion models (IDDPM, ADM, DDIM, A-DPM, EA-DPM, EDM, PFGM++, and FLUX) on datasets of various resolutions with negligible computational overhead. The code is at https://github.com/AMAP-ML/DCW.
Denoising diffusion models are widely used for high-quality image and video generation. Their performance depends on noise schedules, which define the distribution of noise levels applied during training and the sequence of noise levels traversed during sampling. Noise schedules are typically handcrafted and require manual tuning across different resolutions. In this work, we propose a principled way to design per-instance noise schedules for pixel diffusion, based on the image's spectral properties. By deriving theoretical bounds on the efficacy of minimum and maximum noise levels, we design ``tight'' noise schedules that eliminate redundant steps. During inference, we propose to conditionally sample such noise schedules. Experiments show that our noise schedules improve generative quality of single-stage pixel diffusion models, particularly in the low-step regime.