Fit Like You Sample:

Sample-Efficient Generalized Score Matching from Fast Mixing Diffusions

Yilong Qin

1

Andrej Risteski

What is it about?

Our Question:

What makes Annealed Score Matching (diffusion model) so successful?

Specifically, why is Annealed Score Matching so much better than Vanilla Score Matching?

Our Approach:

We tackle this question from the perspective of asymptotic sample complexity.

We show that for multimodal distributions, the sample complexity of annealed score matching is polynomial in problem parameters.

The previous known bound is exponential for vanilla score matching (KHR 2022).

2

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

Yilong Qin, Andrej Risteski

Score Matching

[H2005] Score Matching objective

​

​

​

[L2012] Generalized Score Matching objective replaces the gradient with an arbitrary linear operator

​

[SE2019] Annealed Score Matching objective

​

is a weighted sum of score matching objectives at different temperatures beta, with each distribution tempered with different amount of noise.

3

Annealing process (Image Credit: Yang Song’s blog)

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

Yilong Qin, Andrej Risteski

Asymptotic Sample Complexity

[V2000] Under suitable conditions, the empirical estimator converges to the following asymptotic normal distribution by an extension to the Central Limit Theorem

Interpretation:

The “larger” the asymptotic covariance, the more sample is required to ensure the estimator is sufficiently close to the true parameter.

We will use the operator norm of the asymptotic covariance to quantify the sample complexity.

4

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

Yilong Qin, Andrej Risteski

Multimodality via Markov Chain

Intuition:

The multimodality of a distribution could be characterized by the mixing time of a Markov Chain (e.g. Langevin Dynamics) that converges to the distribution as its stationary distribution. For multimodal distributions, the chain suffers from slow mixing.

[Definition 3] Markov Semigroup of a Markov Process

​

[Definition 4] Infinitesimal Generator

​

[Definition 5] Poincaré constant is defined as the rate of convergence (i.e. mixing time) in chi-square divergence

​

Equivalently, it is the smallest C that satisfies the Poincaré Inequality for all g

​

where the Dirichlet form is given by

5

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

Yilong Qin, Andrej Risteski

Framework for Statistical Gap

We extend the results from KHR (2022) for exponential family: Under conditions (asymptotic normality, realizability) the operator norm of the asymptotic covariance of the Generalized Score Matching objective can be upper bounded by the following bound (Informal Theorem 3)

6

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

Yilong Qin, Andrej Risteski

Framework for Statistical Gap

We extend the results from KHR (2022) for exponential family: Under conditions (asymptotic normality, realizability) the operator norm of the asymptotic covariance of the Generalized Score Matching objective can be upper bounded by the following bound (Informal Theorem 3)

7

For a Markov Chain with some generator, its mixing time determines the sample complexity of a corresponding SM objective.

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

Yilong Qin, Andrej Risteski

Generalized Score and Markov Chains

For a Markov Chain with some generator, its mixing time determines the sample complexity of a corresponding SM objective.

[Theorem 3] For every continuous time Markov Chain, the corresponding SM objective is in fact a pre-conditioned score loss:

​

Key ingredients:

[MCF 2015] For every continuous time Markov Chain with stationary distribution, it admits the following Ito’s diffusion representation:

​

where D is PSD, Q is skew-symmetric, and

​

[Lemma 3] The above Markov Chain has the following Dirichlet form:

​

which is closely related to the Hessian of the loss.

8

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

Yilong Qin, Andrej Risteski

KHR (2022): Vanilla Score Matching - slow mixing

[Example 2 from KHR (2022).] For an exponential family distribution with sufficient statistic

​

The multiplicative gap is exponential w.r.t. a.

9

Asymptotic covariance of Generalized Score

Optimal efficiency via Cramer-Rao

Smoothness

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

Yilong Qin, Andrej Risteski

Ours: Annealed Score Matching - fast mixing

1. Annealed Score Matching (especially higher order) objective is very similar to Score Matching on an “lifted” state space augmented with temperature. We analyze Score Matching on this “lifted” space.
2. Under a Gaussian Mixture distribution (with shared covariance), the Score Matching loss (that arises from tempering dynamics) enjoys a poly(dimension, diameter of means, eigenvalues of covariance) bound under its natural parameterization.

10

Asymptotic covariance of Generalized Score

Optimal efficiency via Cramer-Rao

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

Yilong Qin, Andrej Risteski

Proof Outline

Polynomial mixing time bound:

• Langevin Dynamics on the “lifted” (x, beta) distribution corresponds to an Annealed Score Matching loss.
• Mixing time analysis uses Markov Chain Decomposition Theorem GLR (2018).

Polynomial Smoothness bound:

• Relate score of mixture with score of component via the convexity of perspective map.
• Single Gaussian derivative is given by the Hermite polynomial.

11

Asymptotic covariance of Generalized Score

Polynomial

Optimal efficiency via Cramer-Rao

Polynomial

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

Yilong Qin, Andrej Risteski

Continuously Tempered Langevin Dynamics

[Definition 7] We define Continuous Tempered Langevin Dynamics as the Langevin Dynamics over the joint (x, beta) state space, with reflection at the boundary of support of beta.

​

​

where

​

and

[Proposition 4] The Score Matching loss on the joint state space satisfies

​

Moreover,

12

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

Yilong Qin, Andrej Risteski

Lifted distribution

x: original state space

beta: temperature

i: mode index

13

beta

x

i

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

Yilong Qin, Andrej Risteski

Fast Mixing within mode across temperatures

[Lemma 7] In our mixture analysis, this is polynomial

14

beta

x

i

*VROOM VROOM*

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

Yilong Qin, Andrej Risteski

Fast Mixing within mode across temperatures

[Lemma 7] In our mixture analysis, this is polynomial

15

beta

x

i

*VROOM VROOM*

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

Yilong Qin, Andrej Risteski

Fast Mixing across modes at high temperature

[Lemma 8] In our mixture analysis, this is polynomial

16

beta

x

The projected chain on i becomes a random walk on complete graph

*VROOM VROOM*

i

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

Yilong Qin, Andrej Risteski

Putting it all together

The overall Poincaré constant is bounded by the two Poincaré constants.

​

[Theorem 4] In our mixture analysis, this is polynomial

17

beta

x

i

*VROOM VROOM*

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

Yilong Qin, Andrej Risteski

Takeaways

1. Sample complexity of score matching is governed by mixing time of Markov Chains. Given a (continuous time) Markov Chain, we can design score matching loss with different sample complexities.
2. Annealed score matching loss corresponds to Langevin Dynamics over the lifted (x, beta) distribution, which mixes much faster than Langevin on the original distribution over x.

18

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

Yilong Qin, Andrej Risteski

Thank you!

19