Denoising Efficiency and Lines Matching Models

Abstract:

In this paper, we analyze the denoising loss used by key denoising models and identify an inefficiency that stems from the random pairing which they employ between samples from the source and target distributions. Regressing the denoiser under these non-deterministic conditions causes its predictions to collapse toward the mean of the source or target distributions. We show that this degeneracy creates false basins of attraction, distorting the denoising trajectories and ultimately increasing the number of steps required to sample these models.

We also analyze the alternative strategy of deriving the pairing from an Optimal Transport between the two distributions, and show that while this approach can alleviate this degeneracy, it suffers from a curse of dimensionality, where the pairing set size must scale exponentially with the signal's dimension.

In order to empirically validate and utilize these theoretical observations, we design a new training approach that circumvents these pitfalls by leveraging the deterministic ODE-based samplers, offered by certain denoising diffusion and score-matching models. These deterministic samplers establish a well-defined change-of-variables between the source and target distributions. We use this correspondence to construct a new probability flow model, the Lines Matching Model (LMM), which matches globally straight lines interpolating between the two distributions. We show that the flow fields produced by the LMM exhibit notable temporal consistency, resulting in trajectories with excellent straightness scores, and allow us to exceed the quality of distilling the input correspondence.

The LMM flow formulation allows us to further enhance the fidelity of the generated samples beyond the input correspondence by integrating domain-specific reconstruction and adversarial losses. Overall, the LMM achieves state-of-the-art FID scores with minimal NFEs on established benchmark datasets: 1.57/1.39 (NFE=1/2) on CIFAR-10, 1.47/1.17 on ImageNet, and 2.68/1.54 on AFHQ.

Submission Type: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

We are uploading a revised version of our paper addressing the reviewers' comments. Specifically, this revision contains:

1) in blue reviewer's 3us6 comments regarding the nonlinearity of the data distribution and the models ability to tie it with straight lines from the source distribution.

2) in green reviewer's DB4m comments regarding the focus of our paper, the overloaded mathematical notation, and clarified that our contribution to optimal transport is at the theoretical level.

3) in magenta reviewer's LsYS comments regarding the connection between the theoretical observations made, and their validation through the LMM, and the theoretical question regarding the ability to achieve parameterization using lines model and its practical impact. As in (2) we also addressed the reviewers comment in regard to our contribution to OT.

One last comment:

As we were revising our paper according to the reviewers' comments, we noticed that the comment by reviewer DB4m, saying "For instance, (4)
is the same as Karras 2022's Eq. (2,3)," is inaccurate as Eq (3) in Kerras states that

\grad log p (x,sigma) = (D(x,sigma)-x) / 2sigma^2 in general for arbitrary sigma.

In our paper, Eq 4 focuses on the specific value (maximal) of sigma (t=T), where the denoiser can be computed explicitly by

D(y,sigma_T) = E_x (\grad log p(y|x,T)), for x~p_data.

Besides that these equations are part of different discussions, these equations make different mathematical statements.

We thank all the reviewers for their time and effort,
The Authors

Assigned Action Editor: ~Julius_Berner1

Submission Number: 6287