Disentangling Total-Variance and Signal-to-Noise-Ratio Improves Diffusion Models

arXiv:2502.08598 · cs.LG, stat.ML · Submitted 2025-02-12 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "Disentangling Total-Variance and Signal-to-Noise-Ratio Improves Diffusion Models".

Tom: The gist The authors propose a novel total-variance/signal-to-noise-ratio disentangled (TV/SNR) framework to independently control TV and SNR in diffusion models,

Jane: First, who's behind it and why it matters.

Title and authors: Tom: Moving on to the authors and their title: "Disentangling Total-Variance and Signal-to-Noise-Ratio Improves Diffusion Models." The authors are Khaled Kahouli, Winfried Ripken, Stefan Gugler, Oliver T. Unke, Klaus-Robert Müller, and Shinichi Nakajima.

Jane: These are researchers from various institutions like the Berlin Institute for the Foundations of Learning and Data and Google DeepMind. It shows a collaboration between different major AI groups working on this problem.

Lu: It’s interesting that we see expertise coming from so many places, which usually means they're tackling this from multiple angles, not just one perspective.

Meng: Having such a broad team suggests they're not just focusing on one niche; they are looking at the fundamental mechanisms of how these models sample data.

Tom: They frame the problem as moving beyond passive adjustments of total variance and introducing direct control over both total variance and signal-to-noise ratio, which is what their TV/SNR framework aims to do.

Jane: That’s a key distinction because most prior work has only adjusted the total variance without directly managing that whole picture.

Lu: They are essentially saying we can separate the process into two controllable parts, which is a novel structural approach to noise scheduling in these models.

Meng: So this suggests that the bottleneck isn't just about making the model learn better, but also about how efficiently we sample from its learned distribution.

Tom: That’s right; it shifts the focus onto making the sampling process itself more efficient, which is a major area of research in diffusion models.

The paper's summary: Jane: So summarizing the paper, they introduce this new TV/SNR disentangled framework where you can control TV and SNR independently to guide the reverse process through distinct drift and diffusion terms.

Lu: They reformulate the perturbation kernel into a specific mathematical form where TV is defined as τ2(t) equals a2(t) plus b2(t), and SNR is γ(t) which is just a divided a over b, allowing them to control both of these quantities separately in the SDE.

Meng: That means they’ve derived specific drift and diffusion coefficients that depend on how those TV and SNR terms change over time. It's like they’ve written the rules for the movement of the noise at every single step.

Tom: They then show that this formulation leads to specific Stochastic Differential Equations with drift f(t) equals τ˙(t) over τ t plus γ˙(t) over γ (one plus γ2(t)) and diffusion coefficient g(t) equals s−2τ2(t) dot γ / γ (one plus γ2(t)) <ref:2502.08598#pg1>.

Jane: Those equations give them the mathematical tools to actually solve the reverse process using these new dynamics, which is why this is a major step forward in theory.

Lu: The real power here is that they then use these equations to find their specific schedules for TV and SNR, like the VP-ISSNR schedule.

Meng: They don't just give you a generic equation; they provide a concrete schedule, which means we can actually implement something immediately and test it out.

Tom: And this concrete schedule is what allows them to then demonstrate the performance improvements in their experimental results across different tasks.

The paper's improvements: Jane: Regarding the specific improvements they suggest, they show that for schedules where the total variance grows exponentially, adopting a constant TV schedule, or τ(t) equals one consistently improves performance <ref:2502.08598#pg1>.

Lu: This means instead of trying to find a super complex exponential growth schedule for TV that might not actually be necessary for performance gains.

Tom: They suggest that focusing on optimizing the SNR schedule alone is sufficient if you keep the total variance constant at one, which they call the VP trajectory.

Meng: That simplifies things because it means we don't have to hunt for a complicated schedule for TV; we can focus all our effort on tuning that single SNR function.

Jane: They propose this specific exponential inverse sigmoid or ISSNR function for scheduling the SNR, which they show is a generalization of optimal transport flow matching when you set η equals one point zero and κ equals zero to recover the original OTFM schedule <ref:2502.08598#pg2>.

Lu: By generalizing OTFM with a constant TV schedule, they gain more flexibility in controlling the probability flow trajectory, which is a significant theoretical win for understanding these dynamics.

Tom: In practice, this translates into concrete results like stability rates surpassing seventy-four percent with only eight steps on the QM9 dataset using their VP-ISSNR schedule.

Meng: Those stability rates are impressive because they show that this method can generate physically plausible molecules much faster than methods like EDM or OTFM in those specific benchmarks.

Conclusion: Jane: So to wrap up, the main point of "Disentangling Total-Variance and Signal-to-Noise-Ratio Improves Diffusion Models" is that separating these two concepts gives us independent control over the variance and SNR.

Lu: This leads to a practical schedule like VP-ISSNR, which offers significant improvements in stability for molecular generation and competitive results in image generation across various modalities.

Tom: The big implication here is that improving the sampling process isn't just one thing; it’s about finding ways to optimize the noise schedule orthogonally to other parts of the model training.

Meng: It also shows that this scheduling approach is robust, because it performs well across different solvers and architectures, suggesting we can combine these scheduling tricks with other sampler improvements.

Lu: So they've laid a path for controlling these complex dynamics in a new way using TV/SNR disentanglement that has broad applicability.

Tom: We’ve seen how this paper, "Disentangling Total-Variance and Signal-to-Noise-Ratio Improves Diffusion Models," offers a concrete method, the VP-ISSNR schedule, to significantly boost efficiency in diffusion models.

BIFOLD – Berlin Institute for the Foundations of Learning and Data · Machine Learning Group, Technische Universität Berlin · Department of Artificial Intelligence, Korea University · Max-Planck Institute for Informatics · Google DeepMind · RIKEN Center for Advanced Intelligence Project

cs.LG, stat.ML

Submitted: 2025-02-12

Updated: 2026-10-09

Code: https://github.com/khaledkah/tv-snr-diffusion

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 89/100

The gist: The gist The authors propose a novel total-variance/signal-to-noise-ratio disentangled (TV/SNR) framework to independently control TV and SNR in diffusion models, revealing that constant TV schedules

Key concepts

Total Variance (TV)
TV measures the total amount of noise present in the diffusion process, defined by a combination of two terms. The authors found that keeping this value constant often leads to better generation quality, especially when the variance naturally tends to increase over time.
Signal-to-Noise Ratio (SNR)
SNR represents the ratio between the signal and noise in the diffusion process. The authors developed a specific schedule for SNR, called VP-ISSNR, which they found significantly enhances sampling efficiency and achieves superior results compared to existing methods.
Disentangled Framework
This approach separates the control of TV and SNR into independent variables. By controlling them separately, the researchers could test how each parameter affects generation quality without being constrained by the other, leading to more robust and interpretable scheduling strategies.

Terminology

Summary

The gist The authors propose a novel total-variance/signal-to-noise-ratio disentangled (TV/SNR) framework to independently control TV and SNR in diffusion models, revealing that constant TV schedules improve performance for exponentially growing TV schedules, and they introduce an exponential inverse sigmoid (ISSNR) schedule for SNR that achieves state-of-the-art sampling efficiency.

How it works

The proposed framework reformulates the perturbation kernel as p(x(t)x(0)) = N x(t); s τ2(t)γ2(t) 1 + γ2(t) I, where TV is defined as τ2(t) = a2(t) + b2(t), and SNR is defined as γ(t) = a(t)/b(t), allowing them to be controlled independently <ref:2502.08598#pg14>. This formulation leads to SDEs with drift f(t) = τ˙(t) / τ (t) + γ˙(t) / γ (t) (1 + γ2(t)) and diffusion coefficient g(t) = s−2τ2(t) ˙γ(t)/γ (1 + γ2(t)) <ref:2502.08598#pg20>.

VP Variants of Established Non-VP Schedules

The authors empirically show that for common existing schedules where the TV grows exponentially, their constant TV (i.e., VP) variants consistently improve performance. Specifically, they demonstrate that the VP versions lead to substantial enhancements for SMLD and EDM-UT, where the TV increases exponentially, whereas OTFM with a smooth TV modulation performs comparably to its VP analog. This suggests that exponentially increasing TV can be detrimental, thereby validating their choice to adopt a constant TV schedule.

Variance-preserving Inverse Sigmoid SNR (VP-ISSNR) Schedule

The authors propose to schedule TV and SNR with the following functions: τ2(t) = 1, and γ2(t) = exp [2η log 1 / t(tmax − tmin) + tmin − 1 + 2κ]. They show that the proposed SNR schedule is a generalization of the optimal transport flow matching in Eq. (14) with a constant TV schedule, where setting η = 1.0 and κ = 0 recovers the SNR schedule of OTFM. This generalization allows for more control over the generated probability flow.

Experimental Results and Discussion

In molecular structure generation on the QM9 dataset, the VP-ISSNR schedule leads to stability rates that surpass 74% with only 8 steps and reaches nearly 87% with 128 steps, outperforming all other schedules. In image generation on CIFAR-10 and FFHQ, the VP-ISSNR schedule achieves competitive results across various solvers and different modalities. Furthermore, they observe that the VP-ISSNR schedule leads to trajectories with smaller curvature than both EDM and OTFM in toy problems where data is a mixture of three delta peaks.

Conclusion

In conclusion, reformulating diffusion processes within their TV/SNR framework enables a new way to improve the efficiency of diffusion models, leading to significant progress in domains like molecular structure generation. The authors find that the ISSNR schedule leads to drastic performance improvements (up to 30-fold) in molecular structure generation: Stable molecules can be generated after only 4 steps (much less than the previous SOTA). This work demonstrates that their noise schedule optimization is orthogonal to other methods that make sampling of diffusion models more efficient. The TV/SNR framework is robust and effective across domains, as evidenced by its strong performance in both molecular structure and image generation tasks. The authors conclude that their TV/SNR framework enables a new way to improve the efficiency of diffusion models, leading to significant progress in domains like molecular structure generation.

How it works

The authors show that the stability rate increases significantly to 93.16% with 32 NFEs and 95.82% with 64 NFEs when solving the reverse SDE using their VP-ISSNR schedule. They also find that their method can generate physically plausible molecules that are structurally similar to ground truth reference structures. The results further support the hypothesis that improvements in the sampler/solver are orthogonal to improvements in the noise schedule and can be combined for even greater overall improvement. The authors also observe similar trends across different solvers, supporting this orthogonality.

How it works

The authors show that the model trained on the cosine schedule achieves consistently better results than the model trained on the EDM schedule, even when using the EDM schedule during sampling. This suggests that the cosine schedule samples more points on the relevant SNR region. The authors also observe that their VP-ISSNR sampling still outperforms other baselines, even when the model is trained using the EDM schedule. The results further support the hypothesis that improvements in the sampler/solver are orthogonal to improvements in the noise schedule and can be combined for even greater overall improvement.

How it works

The authors observe that their VP-ISSNR sampling still outperforms other baselines, even when the model is trained using the EDM schedule. The results further support the hypothesis that improvements in the sampler/solver are orthogonal to improvements in the noise schedule and can be combined for even greater overall improvement. The authors also observe similar trends across different solvers, supporting this orthogonality. The authors conclude that their TV/SNR framework is robust and effective across domains, as evidenced by its strong performance in both molecular structure and image generation tasks. The authors also observe consistent quality improvements across different solvers, suggesting that their scheduling approach is orthogonal to other methods that make sampling of diffusion models more efficient.

How it works

The authors show that the VP-ISSNR schedule leads to low overall curvature, while the local curvature is only high in regions with large support. This suggests that their hypothesis would explain why the VP-ISSNR schedule and VP variants of existing VE schedules improve sample quality. The authors also find that OTFM, their VP-OTFM, and their VP-ISSNR lead to straight trajectories around t ≈ 0, which is consistent with their experimental results on molecule and image generation. The reason why EDM performs so well for image generation despite the high curvature of trajectories around t ≈ 0 is likely because the non-uniform time grid assigns most integration steps precisely in this region.

How it works

The authors observe that their VP-ISSNR schedule and VP variants of existing VE schedules improve sample quality. The authors also find that OTFM, their VP-OTFM, and their VP-ISSNR lead to straight trajectories around t ≈ 0, which is consistent with their experimental results on molecule and image generation. The reason why EDM performs so well for image generation despite the high curvature of trajectories around t ≈ 0 is likely because the non-uniform time grid assigns most integration steps precisely in this region.

How it works

The authors observe that their VP-ISSNR schedule and VP variants of existing VE schedules improve sample quality.

Improvements for AI systems

  1. Improved noise schedule control via TV/SNR framework: By introducing a total-variance/signal-to-noise-ratio disentangled (TV/SNR) framework, where TV and SNR can be controlled independently, AI systems can generate samples with significantly reduced sampling time, as the paper shows that schedules where the TV explodes exponentially can often be improved by adopting a constant-TV schedule while preserving the same SNR schedule.

  2. Enhanced molecular structure generation fidelity: The proposed VP-ISSNR schedule is shown to lead to drastic performance improvements (up to 30-fold) in molecular structure generation, enabling the creation of physically plausible molecules that are structurally similar to ground truth reference structures when generating equilibrium states from a given composition.

  3. Optimized image generation efficiency: The ISSNR schedule, which is a generalization of optimal transport flow matching, achieves state-of-the-art (SOTA) sampling efficiency in molecular structure generation and performs comparably to SOTA methods in image generation tasks across various modalities.

  4. Robustness across solvers and architectures: The TV/SNR framework demonstrates that the scheduling approach is orthogonal to other methods that make sampling of diffusion models more efficient, leading to consistent quality improvements across different solvers, suggesting that our scheduling approach is orthogonal to other methods.

Sources

Related papers