Fit Like You Sample: Sample-Efficient Generalized Score Matching from Fast Mixing Diffusions

Thank you for your feedback! We are glad you find the problem well-motivated. We are happy to make changes to the presentation if you have concrete suggestions on how to make it more approachable.

Regarding your question: $\Gamma_{SM}$ and $\Gamma_{MLE}$ are the covariance matrices of the asymptotic limit for the score matching and maximum likelihood estimators respectively (i.e. the matrix $\Gamma$, s.t. for the estimator $\hat{\theta}_n$, we have $\sqrt{n}(\hat{\theta}_n - \theta^*) \to \mathcal{N}(0, \Gamma)$).

Dear reviewer BbBg,

thank you again for your feedback !

We'd like to remind you that the discussion period is ending tomorrow (11/22), and wanted to follow up and check if there are any outstanding concerns or questions we could answer. If our response (as well as the response to the other reviewers) answered your questions, we hope you'll consider raising your score.

Best,

--Authors

Thank you for your feedback and questions! We are glad you found our results new and non-trivial, and our writing clear. Hopefully the clarifications below ameliorate your concerns.

Contribution of Theorem 2: You are correct that our results are a generalization of Koehler et al. (2022). While the overall proof strategy is similar, Theorem 2 is important in two ways:

(1) Conceptually our result makes it clear that there is nothing special about Langevin diffusion (and its mixing time, as captured through the standard Poincare constant). In particular, for any diffusion of the form (7), one can connect the statistical complexity of the corresponding generalized score matching loss to the mixing time of the diffusion. This provides a “dictionary” to translate between fast-mixing diffusions and corresponding score losses with good statistical complexity. Furthermore, the results in Koehler et al. (2022) also only apply to exponential families—our proof shows this restriction is not needed.

(2) Technically the main new ingredient is Lemma 4 (along with the calculations in Lemmas 3, 11 and Lemma 13) which connects the Hessian of the generalized score matching loss to the Dirichlet form of the corresponding diffusion. Again, these calculations are completely new compared to Koehler et al. (2022).

Clarification on Theorem 5:

(Q1) Computational efficiency of training score loss, motivation: The general motivation behind score-matching methods is to avoid calculating a partition function, by fitting the score function through the integration by parts formula (eq (2)). In our case, as in Proposition 6, the CTLD score loss can also be rewritten by integration by parts as an expectation of expressions involving the fitted scores.

Note that even for basic score matching, the corresponding loss may not be convex, so all that’s guaranteed is that a gradient-based method can be efficiently run (which is not the case for maximum likelihood due to the partition function). Nevertheless, there is an algorithmic gain compared to maximum likelihood (the possibility of running a gradient-based algorithm efficiently)—so the question of understanding the "statistical cost” (how much worse it is compared to maximum likelihood) is meaningful.

(Q2) Dependence on K: the reason why we phrase the result as $\lVert\Gamma_{SM}\rVert_{OP} \leq C \lVert\Gamma_{MLE}\rVert^2_{OP}$ is that we are trying to compare how much (statistically) less efficient score matching is compared to MLE. We stress the point about $C$ not depending on $K$ for two reasons:

(1) While MLE could potentially depend on $K$—the dependence is likely to be mild (we are estimating $K$ vectors, so the dependence would be no worse than linear); if $C$ depends poorly (e.g. exponentially in $K$) it would render our result uninteresting for mixtures with large number of components. Moreover, since MLE is asymptotically statistically optimal, whatever dependence on $K$ it incurs is unavoidable.

(2) Theorem 2 proceeds via a decomposition result, similar as prior analyses of simulated tempering (Lee et al, Ge et al (2018)) — and all these results incur a dependence on $K$. This is because their setup is somewhat different — in particular, they analyze discrete tempering, where the annealing is done via temperature scaling (rather than Gaussian convolution).

(Q3) Comparison with prior work on Gaussian mixture models: We are happy to expand on the discussion in Appendix I. Mostly though, these results are incomparable: they focus on end-to-end provable algorithms for learning mixtures in various regimes—typically trying to deal with the case of mixtures with not-well-separated means. Our focus is on understanding the statistical complexity of score matching losses—not on providing new end-to-end provable algorithms for Gaussian mixtures. (Note, the CTLD loss is not convex, and we don’t analyze what gradient descent or related algorithms would converge to.) Gaussian mixtures are just chosen as a very natural toy model for multimodal data, and as a “proof of concept” that annealed score matching is statistically well-behaved for many distributions for which vanilla score matching is not (via the lower bound in Koehler et al (2022)).

(Q4) Gaussians with different covariance: it is quite likely that our results apply to mixtures with covariances which are $1+O(1/d)$ away from each other, but not larger (were d is the ambient dimension). The reason for this is that annealing “distorts” exponentially the relative volumes of the Gausians (technically, what breaks is the chi-squared divergence bound in Lemma 21). However, we note that for universal approximation using mixtures, one can use a mixture of (exponentially many) Gaussians with covariance $\epsilon I$ to $\epsilon$-approximate the distribution of choice.

Dear reviewer ufaE,

thank you again for your feedback and questions !

We'd like to remind you that the discussion period is ending tomorrow (11/22), and wanted to follow up and check if there are any outstanding concerns or questions we could answer. If our response clarified your original concerns, since you otherwise thought our results are clear, new and non-trivial, we hope you'll consider raising your score.

Best,

--Authors

Thank you for your encouraging review! We are glad you found our framework novel and interesting, and our paper well-written.

(Q1) Strength of Assumptions 1 and 2: Assumption 1 is meant to capture a broad class of multimodal distributions. Moreover, by the universal approximation property of mixtures, any distribution can be approximated by a mixture of sufficiently many Gaussians. Assumption 2 is meant to check that for the “canonical” parametrization of a mixture, the smoothness term in the asymptotic sample complexity (Theorem 4) can be effectively bounded. Some kind of weight-tieing in the parametrization of the score at different temperatures has to be assumed to bound the smoothness term (else, the scores the model estimates at different temperatures need not have anything to do with each other!) It would be a great direction for future research to understand what assumption suffices to have a good bound on this term (we certainly don't think Assumption 2 is necessary). Our goal was to check that the “canonical” parametrization results in a good bound.

(Q2) What is O for CTLD: $\mathcal{O}$ is $\nabla_x$ for CTLD (the Markov Chain is however over the augmented space that includes the temperature random variable).

(Q3) Theorem 5, relation to Koehler et al (2022): The results in Koehler et al (2022) are for the basic version of score matching (eq (2)), and they incur a bound that depends on the Poincare constant of the distribution, which can be quite bad for well-separated mixtures. For instance, for a mixture of even two Gaussians with means $-\mu$ and $\mu$ and variance $\sigma^2$, the Poincare constant scales exponentially in $\mu^2/\sigma^2$ [1]. By contrast, our results scale polynomially in the norms of the means, the ambient dimension, and minimum and maximum eigenvalues of the covariance matrices.

(Q4) Using CTLD for diffusion models: The CTLD loss can plausibly be used for better diffusion model training: the primary obstacle is the fact that it involves derivatives of the score, which is a $d \times d$ matrix. As we note after Proposition 3, similar “high-order” analogues of score losses have been proposed in some works (e.g. [2]), with some success. Some preconditioning strategies for sampling from diffusion models have also been recently proposed [3], again with some success.

[1] Chafai, Djalil, and Florent Malrieu. "On fine properties of mixtures with respect to concentration of measure and Sobolev type inequalities." In Annales de l'IHP Probabilités et statistiques, vol. 46, no. 1, pp. 72-96. 2010.

[2] Meng, Chenlin, Yang Song, Wenzhe Li, and Stefano Ermon. "Estimating high order gradients of the data distribution by denoising." Advances in Neural Information Processing Systems 34 (2021): 25359-25369.

[3] Zhang, Qinsheng, and Yongxin Chen. "Fast Sampling of Diffusion Models with Exponential Integrator." In The Eleventh International Conference on Learning Representations. 2022.