Paper: Fit Like You Sample: Sample-Efficient Generalized Score Matching from Fast Mixing Diffusions (ICLR-2024)

Reviewer comment (Writing):
The presentation is fairly technical and may pose challenges for readers who are not an expert in this particular area.

Author response:
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.

Reviewer comment (Writing):
Theorem 2: What are $\Gamma_{SM}$ and $\Gamma_{MLE}$?

Author response:
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)$).

Reviewer comment (Theory):
Regarding Theorem 5: The estimator that you use here doesn't seem to be efficiently computable, is that correct? The score matching estimator for exponential families from Koehler et al. (2022) is efficiently computable (please correct me if I'm wrong), so their motivation to study it and compare with MLE is clear to me. What is the motivation of usage of your estimator for this problem if it is not efficiently computable?

Author response:
*(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.

Reviewer comment (Theory):
Regarding Theorem 5: As I understood, you are interested in the regime \(K \gg d\), so the fact that \(C\) from your bound \(\Vert \Gamma_{SM} \Vert_{OP} \le C \Vert \Gamma_{MLE}\Vert_{OP}^2\) does not depend on \(K\) is important and nice. However, as you said in the footnote, \(\Vert \Gamma_{MLE}\Vert_{OP}\) may depend on \(K\), so \(K\) can appear in the end in the error, i.e. \(\Vert \hat{\mu_i} - \mu^*_i \Vert\) may depend on \(K\) even when the corresponding error for MLE doesn't. So it is not clear to me why this dependence on \(K\) was important from the very beginning. If it was really important, then it would make sense to bound not \(\Vert \Gamma_{SM} \Vert_{OP}\), but the largest diagonal entry of \(\Gamma_{SM}\) in terms of the largest diagonal entry of \(\Gamma_{MLE}\). In this case, if there is such a bound with a factor that does not depend on \(K\), then it should imply a bound on \(\Vert \hat{\mu_i} - \mu^*_i \Vert\) that does not depend on \(K\) (as long as corresponding errors of MLE do not depend on \(K\)). Is it possible to derive such a bound?

Author response:
*(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).

Reviewer comment (Novelty):
Regarding Theorem 5: There is no comparison with prior works on Gaussian mixtures. While you refer to some of these works in the paper, it is not immediately clear how your result is comparable with them. I think it makes sense to add such a comparison.

Author response:
*(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)).

Reviewer comment (Theory):
Can your approach be generalized to more general mixtures of Gaussians, when not all of them have the same covariance (but, say, when all covariances have condition number bounded by \(O(1)\))?

Author response:
*(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.

Reviewer comment (Theory):
Several assumptions in the paper seem a bit strong, and it would be good if the authors can elaborate a bit more on them.

Author response:
*(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.

Reviewer comment (Theory):
Assumption 1 and 2 of Theorem 2 seem pretty strong. Could the authors give an example where these assumptions hold when $\mathcal{O} \neq \nabla_x$ and not from CTLD? In general, how do we validate these assumptions?

Author response:
*(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.

Reviewer comment (Writing):
What is $\mathcal{O}$ for CTLD?

Author response:
*(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).

Reviewer comment (Theory):
Maybe the authors can comment a bit on how Theorem 5 compares with the previous results. In particular, can the results of Koehler 2022 also be applied to get an upper bound?

Author response:
*(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.