Diffusion models recover accurate mixture weights despite score function insensitivity
July 16, 2026
Abstract
Score-based generative models exhibit a puzzling behavior: they often appear to cover all modes of a target multimodal distribution and yet may fail to learn the correct relative mode amplitudes, which can be interpreted as mixture weights. We resolve this apparent paradox by relating the diffusion score matching (DSM) loss to the error in estimating mixture weights from generated samples. We show that, even when the target score is insensitive to mixture weights, generated samples can recover the weights accurately if the scores at intermediate noise levels are informative about the weights. Accordingly, we define the diffusion score sensitivity index (DSSI) as the variation in the DSM loss relative to changes in a parameter. We then show that the DSSI governs the accuracy with which the parameter of the target distribution can be estimated from generated samples. For Gaussian mixtures in arbitrary dimensions, we prove that the mixture weight estimation errors are on the same order as the DSM loss under mild conditions. Empirically, we show the emergence of sensitivity during the noising process of benchmark data distributions under typical noise schedules, and that these sensitivity values predict how well a well-trained model recovers mixture weights. Furthermore, we show that the choice of noise schedule can reduce diffusion sensitivity, leading to mode amplification. Although we focus on mixture weights, the proposed sensitivity framework governs the recovery of any qualitative parameter of the target distribution.
1 Introduction↩︎
Generative models are powerful tools that are increasingly used in scientific research [1]–[3], medicine [4]–[6], the justice system [7], [8], and more. In these and other settings, users rely on generative models to accurately reflect the qualitative properties of the true target distribution. Past work notes that learning the score of a target distribution may be insufficient for accurate sampling [9]–[11], particularly for canonical multimodal distributions such as Gaussian mixtures [12], [13]. Generative models that incorrectly reflect the amplitudes or variances of modes [14]–[16] may hinder uncertainty quantification and the recovery of qualitative parameters [17]–[19]. Our work proposes an explicit evaluation of generative models through the lens of parameter recovery. Such an evaluation augments sample-quality metrics, such as the Frechet inception distance (FID) [20], which may not explicitly test for target-specific structure, such as multimodality [21]. In this work, we present an information-theoretic mechanism that determines when such key parameters may be represented accurately by diffusion models. In the Gaussian mixture setting, even when the score matching loss is insensitive to mixture weights, learning a diffusion model anneals the score matching loss across steps of the diffusion process [22]; this annealing leads to sensitivity to mixture weights, resulting in accurate estimates that were unanticipated by past analyses that focused only on the (unannealed) score matching loss [9].
This idea is illustrated in 1. The top row plots the score functions for two Gaussian mixture distributions that are identical except for the mixture weight \(\gamma\) and shows that they are nearly identical except in a region with very low probability mass. We are thus unlikely to have sufficient training samples for the loss to reliably indicate the value of \(\gamma\) that best represents the training data. However, for larger values of \(t\) in the diffusion process (\(t=0.25\) or \(t=0.50\) in the figure), the corresponding distributions are less well separated, and the score functions are clearly distinguishable over a higher-probability region. Since training diffusion models uses the score across multiple \(t \in [0,1]\), learned models hone in on the correct value of \(\gamma\).
To quantify this phenomenon, we measure the sensitivity of the diffusion score matching loss \(\ell_\mathrm{DSM}\) 7 to changes in parameters of interest. Specifically, let \(\{p_{}^{(\gamma)}\}_{\gamma \in \boldsymbol{\Gamma}}\) be a set of parameterized distributions. Let \(p_{}^{(\gamma^*)}\) be a target distribution with target parameter \(\gamma^*\). The diffusion score sensitivity index (DSSI) of the target parameter \(\gamma^*\) is defined as, \[\label{eq:32DSSI} L(\gamma^*) = \inf_{\gamma \in \boldsymbol{\Gamma}}~ \frac{\ell_\mathrm{DSM}(p_{}^{(\gamma)},p_{}^{(\gamma^*)})}{(\gamma-\gamma^*)^2}\tag{1}\]
We note that both \(\ell_\mathrm{DSM}\) and \(L(\gamma^*)\) depend on our choice of noise schedule. If the target data distribution has mixture weight \(\gamma^*\), and we learn a distribution with mixture parameter \(\gamma\), then the squared error of mixture weight estimation \(\gamma\) is bounded by the diffusion score matching loss divided by the DSSI. Thus, when the DSSI is small, our learned distribution may be highly misrepresentative of \(\gamma^*\) even when the diffusion score matching loss is low. In contrast, when \(L(\gamma^*)\) is large, we know that a small diffusion score matching loss indicates low error in the parameter \(\gamma\). Our main contributions are as follows:
For general bimodal mixtures, we prove a strict minimax separation in mixture-weight recovery between estimators that use only the target-score data and estimators with access to score data at an intermediate, mode-overlapping noise level. (Theorem 1)
For Gaussian mixtures, we prove a lower bound on the DSSI that holds in any dimension; thus, parameter estimation error is controlled by the DSM loss (Theorem 2).
We further empirically evaluate the DSSI on a real-world dataset. These experiments illustrate how design decisions relating to diffusion model training and inference-time sampling schedules can affect the DSSI and the fidelity of generated samples to the training data distribution.
2 Preliminaries: score-based diffusion models↩︎
Score-based diffusion models sample from a target distribution on \(\mathbb{R}^d\) by tracing out a one-parameter family of distributions \(\{p_t\}_{t \in [0,1]}\) indexed by a diffusion time \(t\), with \(p_0\) being the target data distribution and \(p_1 \approx \mathcal{N}(\mathbf{0}, \mathbf{I})\) the standard normal distribution. A forward process progressively transforms samples from \(p_0\) into samples from \(p_1\) along this family. A reverse process inverts the densities along the forward process, starting from \(\mathcal{N}(0,\mathbf{I})\), using learned approximations of the scores of the intermediate marginals \(p_t\). Thus, score-based diffusions have 3 components: (a) the forward diffusion process and its corresponding reverse process, (b) a training objective for learning the marginal score functions \(\nabla \log p_t\) of the intermediate distributions for \(t\in [0,1]\), and (c) a sampling procedure used to approximately generate samples from the target data distribution at inference time. In this section, we provide a brief description of each component.
2.0.0.1 Noising process.
Given a data distribution \(p_0\) on \(\mathbb{R}^d\), the noising process we consider is the variance-preserving forward SDE (VP-SDE) proposed in [23]. Give a noise schedule \(\set{\beta_t}\), this takes the form \[\mathrm{d}{\mathbf{x}}_t \;=\; -\tfrac{1}{2} \beta_t \, {\mathbf{x}}_t\, \mathrm{d}t \;+\; \sqrt{\beta_t}\,\mathrm{d}\mathbf{w}_t, \qquad {\mathbf{x}}_0 \sim p_0,\quad t \in [0,1], \label{eq:forward-sde}\tag{2}\] where \(\mathbf{w}_t\) is a standard \(d\)-dimensional Wiener process. We denote by \(p_t\) the marginal distribution of \({\mathbf{x}}_t\). The forward conditional law of \({\mathbf{x}}_t\) given \({\mathbf{x}}_0\) under 2 is Gaussian, \[p_t({\mathbf{x}}_t \mid {\mathbf{x}}_0) \;=\; \mathcal{N}\!\left( {\mathbf{x}}_t;\; \sqrt{\alpha_t}\,{\mathbf{x}}_0,\; (1-\alpha_t)\,\mathbf{I} \right), \label{eq:forward-kernel}\tag{3}\] with conditional mean \(\sqrt{\alpha_t} {\mathbf{x}}_0\) and conditional covariance \((1-\alpha_t) \mathbf{I}\), where \(\alpha_t=e^{-\int_0^t\beta_s \mathrm{d}s}\).1
2.0.0.2 Denoising process.
[24] established that the time-reversal of 2 is itself an SDE: \[\mathrm{d}{\mathbf{x}}_t \;=\; \left[-\tfrac{1}{2}{\beta_t}\,{\mathbf{x}}_t \;-\; {\beta_t}\,\nabla_{{\mathbf{x}}} \log p_t({\mathbf{x}}_t)\right] \mathrm{d}t \;+\; \sqrt{{\beta_t}}\,\mathrm{d}\bar{\mathbf{w}}_t, \qquad {\mathbf{x}}_1 \sim p_1, \label{eq:reverse-sde}\tag{4}\] where \(\bar{\mathbf{w}}_t\) is the reverse-time Wiener process. Starting from \({\mathbf{x}}_1 \sim p_1\) and integrating the reverse SDE with the exact score function from time \(t=1\) to \(t=0\) produces samples \({\mathbf{x}}_0 \sim p_0\). In practice, this reverse SDE is simulated starting from \({\mathbf{x}}_1 \sim \mathcal{N}(0,\mathbf{I}) \approx p_1\) using the learned approximation \(\mathbf{s}_{{\boldsymbol{\theta}}}: \mathbb{R}^d \times [0,1] \rightarrow \mathbb{R}^d\) with parameters \({\boldsymbol{\theta}}\in \boldsymbol{\Theta}\) trained such that \(\mathbf{s}_{{\boldsymbol{\theta}}}({\mathbf{x}},t) \approx \nabla \log p_t({\mathbf{x}})\) simultaneously for all noise levels to produce samples approximately distributed as the target \(p_0\).
For a given trained score model \(\{\mathbf{s}_{{\boldsymbol{\theta}}}(\cdot, t)\}_{t\in [0,1]}\) (or the true score function \(\{\nabla \log p_t\}_{t\in [0,1]}\)), generation of samples is an inference task: one may simulate the reverse SDE in 4 with any choice of consistent numerical integration schemes [25]. The overall process has two distinct sources of error: the approximation error incurred in learning the score and the sampling error incurred by the choice of sampling algorithm (e.g., discretization error, finite timesteps, etc.).
3 Minimax bounds for parameter estimation from diffusion models↩︎
In this section, we show that when the score along the noising process exhibits sensitivity to \(\gamma\), the information-theoretic barrier that disallows its recovery from the target score alone is removed. We prove that when the score is sensitive at some noise level, we obtain a lower parameter estimation error when compared to only using the target score. While stated explicitly for mixture weight as the parameter, the results in this section apply to any parameter. We consider only bimodal mixtures for simplicity, but make no distributional assumptions about the components.
3.0.0.1 Mixture Family.
Let \(p_{}^{(0)}\) and \(p_{}^{(1)}\) be two fixed probability distributions on \(\mathbb{R}^d\). A probability density \(p_{}^{(\gamma)}, \gamma \in \boldsymbol{\Gamma} \subseteq (0,1),\) is a member of a family of two-component mixtures if it is of the form \(p_{}^{(\gamma)}\mathrel{\vcenter{:}}= (1-\gamma)\: p_{}^{(0)} + \gamma \:p_{}^{(1)}\). The mixture parameter \(\gamma\) controls the relative weights of the two component distributions, \(p_{}^{(0)}\) and \(p_{}^{(1)}\). We now consider the evolution of a mixture, \(p_{}^{(\gamma)},\) under a noising process or forward process in a score-based diffusion model. Specifically, when \(p_0 \equiv p_{}^{(\gamma)},\) we denote by \(p_{t}^{(\gamma)}\) the marginals of \(\mathbf{x}_t\) following the process 2 . Let \(\mathbf{s}_{}^{(\gamma)}\) denote the score of the density \(p_{}^{(\gamma)}\), and similar for marginal densities of \(\mathbf{x}_t\): \[\label{eq:score-and-noisyscore} \mathbf{s}_{}^{(\gamma)} := \nabla \log p_{}^{(\gamma)} \qquad \text{and} \qquad \mathbf{s}_{t}^{(\gamma)} := \nabla \log p_{t}^{(\gamma)}.\tag{5}\] The relative weights of the mixture are preserved under the forward process: the noisy marginal \(p_{t}^{(\gamma)}\) is itself a \(\gamma\)-mixture of the noisy component marginals \(p_{t}^{(0)}\) and \(p_{t}^{(1)}\) (see 3). As a consequence the score of the noisy mixture density \(p_{t}^{(\gamma)}\) can be decomposed as follows, \[\begin{align} \mathbf{s}_{t}^{(\gamma)}({\mathbf{x}}) = w_t^{(\gamma,0)}({\mathbf{x}})\,\mathbf{s}_{t}^{(0)}({\mathbf{x}}) \;+\; w_t^{(\gamma,1)}({\mathbf{x}})\,\mathbf{s}_{t}^{(1)}({\mathbf{x}}) \label{eq:mix32score} \end{align}\tag{6}\] with component weights \(w_t^{(\gamma,0)}({\mathbf{x}}) = (1-\gamma)\,p_{t}^{(0)}({\mathbf{x}})/p_{t}^{(\gamma)}({\mathbf{x}})\) and \(w_t^{(\gamma,1)}({\mathbf{x}}) = \gamma\,p_{t}^{(1)}({\mathbf{x}})/p_{t}^{(\gamma)}({\mathbf{x}})\). This decomposition enables a simplified theoretical analysis, as we see in the next section.
3.0.0.2 Classical score matching loss/relative Fisher information.
The classical score matching (SM) loss measures the relative Fisher information between distributions [26]. For a pair of mixtures, \[\begin{align} \ell_\mathrm{SM}(p_{}^{(\gamma')},p_{}^{(\gamma)}) := \mathbb{E}_{\mathbf{x}\sim p_{}^{(\gamma)}} \|\mathbf{s}_{}^{(\gamma')}(\mathbf{x})-\mathbf{s}_{}^{(\gamma)}(\mathbf{x}) \|^2. \end{align}\] The relative Fisher information between a pair of mixture distributions with identical components but different mixture weights is small when the components are well-separated. In 1, the scores of Gaussian mixtures may appear nearly identical outside a low-probability region, an effect strengthened with increasing separation between mixture components as shown in 2 (d). Any sampling algorithm that relies solely on the score, such as Langevin sampling, exhibits near-identical behavior across distinct mixture distributions; consequently, the generated samples may not accurately reflect the target mixture weights.
3.0.0.3 Diffusion score matching loss for mixtures.
The diffusion score matching (DSM) loss anneals the score matching loss across the noising process 2. For a pair of mixtures, \[\begin{align} \ell_\mathrm{DSM}( p_{}^{(\gamma')},p_{}^{(\gamma)}) &:= \mathbb{E}_{\substack{t \sim \mathrm{Unif}(0,1)\\\mathbf{x}\sim p_{t}^{(\gamma)}}}[\lambda_t \|\mathbf{s}_{t}^{(\gamma')}(\mathbf{x}) - \mathbf{s}_{t}^{(\gamma)}(\mathbf{x})\|^2] = \mathbb{E}_{t \sim \mathrm{Unif}(0,1)}[\lambda_t \: \ell_\mathrm{SM}(p_{t}^{(\gamma')},p_{t}^{(\gamma)})]. \label{eq:diffScoreMatchingLoss} \end{align}\tag{7}\] Here, \(\lambda_t\) is a weighting factor set to \(\lambda_t = 1-\alpha_t\) following [23]. The DSM loss measures the relative Fisher information over the entire sequence of noised densities generated by the process 2 . Recall that the marginals \(p_t\)’s and scores \(\mathbf{s}_t\)’s depend on the noise schedule \(\{\alpha_t\}_t\), as shown in 3 and 5 , and this way, the diffusion score matching loss \(\ell_\mathrm{DSM}\) is dependent on the choice of noise schedule. In practice, we may train a neural network \(\mathbf{s}_{{\boldsymbol{\theta}}}\) to approximate the score function. With some abuse of notation, we denote \(p_{}^{({\boldsymbol{\theta}})}\) the distribution2 of samples generated using the score functions \(\{\mathbf{s}_{{\boldsymbol{\theta}}}(\cdot, t)\}_{t\in (0,1)}\) for a fixed noise schedule as per 4 . Since the true score function \(\mathbf{s}_{t}^{(\gamma)}\) and indeed the target distribution \(p_{}^{(\gamma)}\) are unknown, calculating the diffusion score matching loss is intractable. Instead, we minimize the objective \[\label{eqn:32neural32network32loss} \ell_{\mathrm{NN}}(p_{}^{({\boldsymbol{\theta}})}, p_{}^{(\gamma)}):=\mathbb{E}_{\substack{t \sim \mathrm{Unif}(0,1)\\\mathbf{x}_0\sim p_{}^{(\gamma)}, \mathbf{x}_t\sim p_{t}^{(\gamma)}(\cdot | \mathbf{x}_0)}}[\lambda_t \|\mathbf{\mathbf{s}}_{{\boldsymbol{\theta}}}(\mathbf{x}_t, t) - \nabla \log p_{t}^{(\gamma)}(\mathbf{x}_t|\mathbf{x}_0)\|^2].\tag{8}\]
As per [{lem:nn-dsm-ordering}], there exists a constant \(C\geq 0\), depending only on \(p_{}^{(\gamma)}\), such that \(\ell_{\mathrm{NN}}(p_{}^{({\boldsymbol{\theta}})}, p_{}^{(\gamma)})=\ell_\mathrm{DSM}(p_{}^{({\boldsymbol{\theta}})}, p_{t}^{(\gamma)}) + C\). Thus, \(\ell_\mathrm{DSM}(p_{}^{({\boldsymbol{\theta}})}, p_{}^{(\gamma)}) \leq \ell_{\mathrm{NN}}(p_{}^{({\boldsymbol{\theta}})}, p_{}^{(\gamma)})\).
3.0.0.4 Minimax theory for mixture weight recovery.
From classical minimax theory (see e.g., [27] for a textbook exposition), we know that the error in any estimator of a parameter is lower-bounded by the informativeness of the data distribution about the parameter. Now, the data distribution loses information about the parameter in the noising process. How, then, would annealing benefit arameter recovery, if the original mixture is well-separated and hence insensitive to parameter changes? The answer lies in the fact that diffusion models learn to approximate the scores at various noise levels. By minimizing \(\ell_\mathrm{DSM},\) diffusion models obtain information about the parameter through the scores even if such information is not present in the samples. Our main result in this section shows the existence of a noise scale at which the score-data law is sensitive, provably enabling parameter recovery rates that are unattainable from target-score data alone.
Definition 1 (Score oracle and score-restricted estimators). We consider a Gaussian White Noise (GWN) oracle from nonparametric regression [28] of the form: \(\hat{\mathbf{s}}_{t}^{(\gamma)} = \mathbf{s}_{t}^{(\gamma)}\: + \sigma_t \: W,\) where \(\sigma_t > 0\) is a scalar and \(W\) is an isonormal field indexed by \(L^2(p_{t}^{(\gamma)})\). The statistical experiment is to estimate \(\gamma\) from the observation, \(\hat{\mathbf{s}}^{(\gamma)}_t.\) Let \(q_{t}^{(\gamma)}\) denote the distribution of the observation, which is a Gaussian measure with mean \(\mathbf{s}_{t}^{(\gamma)}.\) We define the set, \(\Gamma_{\mathrm{score},t},\) of all measurable functions of the GWN observations to \((0,1)\) to be the class of “score-restricted estimators”. Further, we use \(\mathbb{E}_{\gamma^*,t}\) as a shorthand for expectations with respect to the score oracle model, \(q_{t}^{(\gamma)}.\) See also Remark 4 for more details on the estimator classes.
We call \(\Gamma_{\mathrm{score},t}\) score-restricted estimators since they access the score vector field (corrupt by Gaussian noise) and not the target data samples. In other words, our class of estimators \(\Gamma_{\mathrm{score},t}\) models parameter recovery achievable from having oracle access to a corrupt score vector field at time \(t,\) reflecting the denoising/generation phase in diffusion models after scores have been trained. Since the observation model, \(q_{t}^{(\gamma)},\) is Gaussian, we can compute, fixing the indexing Hilbert space for \(W\) to be \(L^2(p_{t}^{(\gamma)}),\) that \(\mathrm{KL}(q_{t}^{(\gamma)}||q_{t}^{(\gamma')}) = \ell_\mathrm{SM}(p_{t}^{(\gamma')}, p_{t}^{(\gamma)})/(2\sigma_t^2).\) See Remark 5 for this computation.
Assumption 1 (Stable score-matching at time \(t\)). We assume that at some time \(t > 0,\) the score observation law, \(q_{t}^{((\cdot))},\) satisfies \((\gamma - \gamma')^2 \leq A_t^2 \: \mathrm{KL}(q_{t}^{(\gamma)}||q_{t}^{(\gamma')})\) for some \(A_t >0\) and all \(\gamma, \gamma' \in (0,1).\)
Since \(\mathrm{KL}(q_{t}^{(\gamma)}||q_{t}^{(\gamma')}) = \ell_\mathrm{SM}(p_{t}^{(\gamma')}, p_{t}^{(\gamma)})/(2\sigma_t^2),\) the above assumption implies that when the score matching loss is bounded above by some \(C,\) then, the squared difference in the corresponding parameters, \((\gamma - \gamma')^2 \leq \:A_t^2\: C / (2\sigma_t^2).\) Intuitively, this means that, for an \(O(\sigma_t^2)\) constant \(A_t,\) the score observations predict distances on the space of weight parameters. Our next definition allows us to make this setting rigorous by comparing score observations at time \(t\) to a different behavior of the scores at time \(0:\) the target score observations, \(\hat{\mathbf{s}}^{(\gamma)},\) are insensitive in law to changes in \(\gamma\) and therefore, do not predict distances on parameter space. In the next section, note that we obtain explicit estimates on \(A_t\) for Gaussian mixtures.
Definition 2 (\((\epsilon,\delta)\)-diffusion sensitivity). We say that a mixture family \(\{p_{}^{(\gamma)}\}\) exhibits an \(\epsilon\)-diffusion sensitivity if for every \(\gamma^* \in (0,1),\) there exist \(\epsilon\)-packing sets \(S(\gamma^*) \subset (0,1)\) containing \(N(\gamma^*,\epsilon) \geq 16\) elements with \(\mathrm{KL}(q_{}^{(\gamma)}||q_{}^{(\gamma')}) = \ell_\mathrm{SM}(p_{}^{(\gamma')}, p_{}^{(\gamma)})/(2\sigma_0^2) \leq \delta(\epsilon,N) \leq \epsilon^2/4,\) for all \(\gamma, \gamma' \in S(\gamma^*)\) satisfying \(|\gamma - \gamma'| \geq 4\sqrt{2}\: A_t \epsilon\).
Note that \(A_t\) here is as defined in Assumption 1.
Definition 3 (\(\epsilon\)-diffusion sensitivity). We say that a mixture family \(\{p_{}^{(\gamma)}\}\) exhibits \(\epsilon\)-diffusion sensitivity if for every \(\gamma^* \in (0,1),\) there exist \(\epsilon\)-packing sets \(S(\gamma^*) \subset (0,1)\) containing \(N(\gamma^*,\epsilon) \geq 16\) elements with \(\mathrm{KL}(q_{}^{(\gamma)}||q_{}^{(\gamma')}) = \ell_\mathrm{SM}(p_{}^{(\gamma')}, p_{}^{(\gamma)})/(2\sigma_0^2) \leq \epsilon^2/4,\) for all \(\gamma, \gamma' \in S(\gamma^*)\) satisfying \(|\gamma - \gamma'| \geq 4\sqrt{2}\: A_t \epsilon\). Note that \(A_t\) here is as defined in Assumption 1.
Theorem 1. Fix \(t>0\) satifying Assumption 1. Let \(\epsilon_{t} > 0\) be a covering radius that solves \(\epsilon_{t}:= \sigma_t\: \sqrt{V(\epsilon_{t})},\) where \(V(\epsilon_{t})\) is the \(\epsilon_{t}\)-covering entropy3 with respect to square root KL divergence of the class of Gaussian measures \(\{q_{t}^{(\gamma)}\}.\) Suppose the bimodal mixture family satisfies 1 and exhibits an \(\epsilon\)-diffusion sensitivity with \(\epsilon = \epsilon_{t}.\) Then, we have the following minimax upper bound on the score-restricted estimators at time \(t\), \[\begin{align} \min_{\hat{\gamma} \in \Gamma_{\mathrm{score},t}}\max_{\gamma^*}~ \mathbb{E}_{\gamma^*,t} |\gamma^* - \hat{\gamma}|^2 \leq 2\:A_t^2\: \epsilon_{t}^2\: . \label{eq:upperBound} \end{align}\qquad{(1)}\] Simultaneously, if the packing set cardinality \(N(\gamma^*, \epsilon_{t}) \geq e^{\sigma_t^2\:V(\epsilon_{t})},\) for all \(\gamma^*,\) it holds that the minimax lower bound achievable for target score-restricted estimators is the following: \[\begin{align} \label{eq:lowerBound} \min_{\hat{\gamma} \in \Gamma_{\mathrm{score},0}}\max_{\gamma^*}~ \mathbb{E}_{\gamma^*,0} |\gamma^* - \hat{\gamma}|^2 \geq 4\: A_t^2\: \epsilon_{t}^2. \end{align}\qquad{(2)}\]
3.0.0.5 Proof.
The upper bound in ?? is a direct application of Theorem 2 from [29]. Then, using a generalized Fano’s method (Lemma 3 of [30]) and an appropriate choice of \(\delta\) as a function of \(\epsilon\) in 2 that limits the packing entropy of the cover of \((0,1)\) with cardinality \(e^{V(\epsilon_{t})},\) we obtain the desired lower bound in ?? . We defer a more detailed proof to Appendix 9. 0◻
Theorem 1 elucidates the conditions under which a diffusion model learns the mixture weights accurately when target score estimation followed by Langevin sampling does not. The above analysis applies to any bimodal mixture density and to any scalar parameter in place of the mixture weight. The main mechanism that enables parameter recovery is that the score field at some time \(t\) changes more rapidly with parameter changes. In contrast, the target score at time 0 may not vary, creating an information-theoretic barrier that limits the accuracy of any estimator. That is, no estimator can distinguish two close values of \(\gamma\) using a finite number of target score samples.
Remark 1. A key contribution of our work is to explain the surprising phenomenon in which a diffusion model can accurately recover a parameter even when the target score remains insensitive to small perturbations of the parameter. In response, we defined the notion of DSSI, 1 , which measures the sensitivity of the scores along the entire noising process, as opposed to only the target score, to infinitesimal changes in the parameter. The most interesting scenario that is directly relevant to explaining “accurate parameter recovery despite target score insensitivity” is one where an intermediate score is “much more” sensitive to the parameter than the score of the target distribution. Our definition of the \(\epsilon\)-sensitivity is designed to capture this scenario and further allows us to fully explain the information-theoretic mechanism underlying accurate recovery (Theorem 1).
We remark that [29] constructs a Bayes density estimator that belongs to the mixture class and achieves the upper bound ?? . Thus, we can conclude from Theorem 1 that using score observations at time \(t\) leads to provably better parameter recovery.
Resolving data-processing inequality. Saying a family of distributions exhibits an \((\epsilon,\delta)\)-sensitivity essentially means that more information about the parameter is created in the noising process, which, at first glance, appears to violate the data processing inequality. The law of noised samples, \(p_{t}^{(\gamma)},\) cannot be more informative about \(\gamma\) at any time \(t\) than \(p_{}^{(\gamma)}\) (the target mixture), due to the data processing inequality. More precisely, since \(p_{t}^{(\gamma)}\) results from a \(\gamma\)-independent convolution operation on \(p_{0}^{(\gamma)},\) the data processing inequality gives: \(\mathrm{KL}(p_{t}^{(\gamma')}||p_{t}^{(\gamma)}) \leq \mathrm{KL}(p_{}^{(\gamma')}||p_{}^{(\gamma)}).\) But, our key idea here is that diffusion models use the score observation fields, \(\hat{\mathbf{s}}^{(\gamma)}_t,\) to recover \(\gamma\), even when \(\hat{\mathbf{s}}^{(\gamma)}\) can not separate \(\gamma\) values sufficiently. Specifically, under the assumptions of Theorem 1 on the mixture family, the score observations can be more informative about \(\gamma\) at time \(t\) than at time 0, i.e., \(\mathrm{KL}(q_{t}^{(\gamma')}||q_{t}^{(\gamma)}) \geq \mathrm{KL}(q_{}^{(\gamma')}||q_{}^{(\gamma)}),\) without violating data processing. In other words, 2 defines a structural property of mixture distributions: the family, \(\{q_{t}^{(\gamma)}\}\) at intermediate \(t\) can be more separated in \(\gamma\) than at \(t=0\) when mixture components do not overlap. In Lemma 1, we prove why this counterintuitive inequality does not violate data processing and why no estimator in \(\Gamma_0\) can reconstruct the score observation model that \(\Gamma_t\) has access to.
3.1 Effect of modified sensitivity to mixture weight↩︎
While Theorem 1 provides a mechanism for parameter recovery, we now discuss experimental evidence for the same mechanism but using the practically computable diffusion score sensitivity index 1 , rather than the information-theoretic quantities used in 2. Our numerical results show that the score sensitivity index quantifies how accurately a diffusion model captures the mixture weight of a Gaussian mixture model. We generate samples using the true analytical score function \(\mathbf{s}_{t}^{(\gamma^*)}(\mathbf{x})\) and a sequence of modified noise schedules with a coarse time discretization (\(T=100\)), designed to artificially reduce the diffusion score sensitivity index \(L(\gamma^*)\) (for details, see 11.1). Recall that \(L(\gamma^*)\) and \(\ell_\mathrm{DSM}\) depend on the choice of noise schedule \(\{\alpha_t\}_t\). We modify both linear and squared-cosine noise schedules and numerically estimate \(L(\gamma^*)\) for these schedules. 2 (c) shows the estimated mixture weight \(\hat{\gamma}\) (calculated using Expectation-Maximization) plotted against \(L(\gamma^*)\). We see that as the sensitivity index \(L(\gamma^*)\) decreases, both the bias and variance in \(\hat{\gamma}\) increase. This demonstrates that noise schedules with lower \(L(\gamma^*)\) yield mixture-weight estimates that are less robust to discretization errors, even with a perfect score estimator.
On the other hand, the means and covariances of the modes are uniformly well-approximated across the different noise schedules. This shows that errors in estimating the mixture weight can occur without noticeable artifacts in the geometry of individual modes. Modern accelerated diffusion sampling schemes (e.g., [31]) leverage coarse time discretizations to reduce the computational cost of sampling while maintaining comparable sample fidelity. Our results suggest that such sampling approaches may be vulnerable to misrepresenting parameters of interest in the data distribution while appearing to generate accurate samples.
Figure 2: Score sensitivity for Gaussian mixture models in dimension \(d=10\). (a) \(\ell_\mathrm{SM}(p_{t}^{(\gamma)}, p_{t}^{(\gamma^*)})\) as a function of \(t\) and \(\gamma\). (b) \(\ell_\mathrm{DSM}(p_{}^{(\gamma)}, p_{}^{(\gamma^*)})\) vs. \(\gamma-\gamma^*\) for multiple \(\gamma^*\), showing non-negligable sensitivity throughout. (c) Mixture weight estimates \(\hat{\gamma}\) from generated data for \(\gamma^*=0.6\) under modified linear and cosine noise schedules (see 11.1). The bias and variance of \(\hat{\gamma}\) increase as \(L(\gamma^*)\) decreases. (d) \(\ell_\mathrm{SM}\) and \(\ell_\mathrm{DSM}\) as a function of separation for \(\gamma^*=0.8\) and \(\gamma=0.9\). \(\ell_\mathrm{DSM}\) remains nonnegligible while \(\ell_\mathrm{SM}\) rapidly decays..
4 Diffusion sensitivity to GMM mixture weights↩︎
We show that for a bimodal isotropic Gaussian mixture in any dimension, the diffusion score sensitivity index (DSSI) \(L(\gamma^*)\) is non-negligible. As a result, a diffusion model with a well-trained score generates samples with the correct weights even for extreme values of \(\gamma\).
Theorem 2 (Diffusion sensitivity in bimodal Gaussian mixtures). Consider the family \(p_{}^{(\gamma)}(\mathbf{x}) =(1-\gamma) \; \mathcal{N}(\mathbf{x}; \mu_0, \mathbf{I}) + \gamma \;\mathcal{N}(\mathbf{x}; \mu_1, \mathbf{I})\) and a linear noise schedule \(\beta_t=\beta_1 t\) for \(\beta_1>0\). Let \(\alpha_t=\exp\left(\int_0^t\beta_s \mathrm{d}s\right)\). Then for any \(a \in (0,1/2)\) \[\begin{align} \min_{\hat{\gamma} \in [a,1-a]}\frac{\ell_\mathrm{DSM}(p_{}^{(\hat{\gamma})},p_{}^{(\gamma^*)})}{(\hat{\gamma}-\gamma^*)^2}>\frac{(1-\Phi(\norm{\mu_1-\mu_0}\sqrt{\alpha_1}))\left(\sqrt{k^2+16}-k\right)}{7\beta_1\gamma^*(1-\gamma^*)}\exp\left(-(k^2+k)/4\right) \label{eq:main95thm} \end{align}\qquad{(3)}\] with \(k:=\left|\log \frac{1-a}{a}\right|+\left|\log\frac{1-\gamma^*}{\gamma^*}\right|\) and \(\Phi(\cdot)\) the CDF of \(\mathcal{N}(0,1)\).
Proof Sketch: The proof proceeds in three steps. We first decompose \(\ell_\mathrm{DSM}(p_{}^{(\hat{\gamma})}, p_{}^{(\gamma^*)})\) by mixture component 24 to reduce the problem to lower-bounding the expectations over paths from \(p_{}^{(0)}\) and \(p_{}^{(1)}\) independently. Second, using the Probability Flow ODE (which is equivalent in probability to the forward process) [23], we find an analytical form for \(\norm{\mathbf{s}_{t}^{(\hat{\gamma})}-\mathbf{s}_{t}^{(\gamma^*)}}^2\) along paths \(\mathbf{z}_t\) starting from \(p_{}^{(0)}\) (4). Finally, we identify the interval of \(t\) values where this expression is bounded below (5), integrate over this region of \(t\) (6), and finally integrate over initial conditions \(\mathbf{y} _ {0}\). The bound is dimension-independent because for a given initial condition \(\mathbf{y}_{0}\) starting from e.g. mode \(p_{}^{(0)}\), the bound only depends on \(\left\langle\mathbf{y}_0-\mu_0, \mu_1-\mu_0\right\rangle/\norm{\mu_1-\mu_0}\sim \mathcal{N}(0,1)\) in any dimension. For full details, see 10.2. 0◻
Remark 2. Recall that if \(\ell_\mathrm{DSM}(p_{}^{(\gamma)},p_{}^{(\gamma^*)})<\delta\) then \((\gamma-\gamma^*)^2 \le \delta/L(\gamma^*)\) by the definition of \(L(\gamma^*)\) in 1 . 2 shows that \(L(\gamma^*)\) cannot be too small under mild conditions for two-component Gaussian mixture models. Further, recall that the neural network loss \(\ell_{\mathrm{NN}}\) 8 is an upper bound for \(\ell_\mathrm{DSM}\), and so if we train a neural network to an error \(\ell_{\mathrm{NN}}(p_{}^{({\boldsymbol{\theta}})},p_{}^{(\gamma)})<\delta\) and our neural network learns a score corresponding to the mixture distributions \(\{p_{t}^{(\gamma')}\}_{t\in (0,1)}\), then \(\ell_\mathrm{DSM}(p_{}^{(\gamma')},p_{}^{(\gamma^*)})<\delta\). Our theory more directly implies that any score model which closely approximates the true score \(\mathbf{s}_{}^{(\gamma^*)}\) in \(L^2\) will have small \(\ell_\mathrm{DSM}\), and thus its sample distribution will accurately reflect \(\gamma^*\).
Remark 3. If, when generating new samples at inference time, we alter the noise schedule \(\{\alpha_t\}\), this may alter the mixture proportions in the generated samples. Specifically, altering the noise schedule changes \(\ell_\mathrm{DSM}\) and \(L(\gamma^*)\), as noted below 7 . Hence, similar loss under \(\ell_\mathrm{DSM}\) for two different noise schedules could potentially correspond to very different guarantees for the accuracy of the recovered mixture weight \(\hat{\gamma}\). Discretizing the reverse process SDE accumulates time-discretization errors regardless of score accuracy, meaning noise schedules that lower \(L(\gamma^*)\) may worsen mixture-weight errors in the generative distribution, even for well-trained diffusion models. 3.1 showcases this effect for GMMs. In 5, we conduct further experiments on real-world data.
4.0.0.1 Effect of parameters on DSSI.
The DSSI \(L(\gamma^*)\) depends on the mode separation \(\norm{\mu_1-\mu_0}\), true mixture weight \(\gamma^*\), and choice of noise schedule \(\set{\alpha_t}_t\). First, consistent with the fact that \(\ell_\mathrm{DSM}(p_{}^{(\gamma)}, p_{}^{(\gamma^*)})\) is an increasing function of the separation (2 (d)), 3 (a) shows \(L(\gamma^*)\) is an increasing function of the mode separation. Second, although 2 is only stated for linear noise schedules, different continuous noise schedules are equivalent up to time-reparameterization, guaranteeing \(\ell_\mathrm{SM}(p_{t}^{(\hat{\gamma})}, p_{t}^{(\gamma^*)})\) is nonnegligible for some range of \(t\) under any noise schedule (3 (b)). Note however that \(L(\gamma^*)\) is defined relative to \(\ell_\mathrm{DSM}\), which averages \(\ell_\mathrm{SM}(p_{t}^{(\hat{\gamma})},p_{t}^{(\gamma^*)})\) over \(t\in[0,1]\), meaning a noise schedule which reduces this range of \(t\) to a small interval \([a,a+\epsilon]\) scales down \(L(\gamma^*)\) roughly proportional to \(\epsilon\), regardless of the peak value of \(\ell_\mathrm{SM}(p_{t}^{(\hat{\gamma})},p_{t}^{(\gamma^*)})\). Third, 3 (c) shows that for mode separation \(\norm{\mu_1-\mu_0}=9\), \(L(\gamma^*)\) is bounded below by 0.37 across all \(\gamma^*\), implying that the squared prediction error in \(\gamma\) is at most \(\approx 2.7\) times larger than the squared diffusion score error. Notably, this bound is dimension-independent, and hence this holds for bimodal isotropic Gaussian mixtures in any dimension. In 11.4, we showcase additional experiments on parameters of Gaussian mixtures other than the mixture weight.
Figure 3: Effect of Gaussian mixture model parameters on score sensitivity in dimension \(d=10\). (a) \(L(\gamma^*)\) vs. mode separation for multiple \(\gamma^*\leq 0.5\) (by symmetry \(L(\gamma^*)=L(1-\gamma^*)\)). \(L(\gamma^*)\) is an increasing function of separation. (b) \(\ell_\mathrm{SM}(p_{t}^{(\gamma)}, p_{t}^{(\gamma^*)})\) as a function of \(t\) and \(\gamma\) for a cosine noise schedule. Compared to a linear noise schedule, the noise schedule only affects the range of \(t\) in the forward process at which the score is sensitive to \(\gamma\). (c) \(\ell_\mathrm{DSM}(p_{}^{(\gamma)}, p_{}^{(\gamma^*)})/(\gamma -\gamma^*)^2\) vs. \(\gamma\); \(L(\gamma^*)\geq 0.27\) for all \(\gamma^*\), with the minimum at \(\gamma^*=0.5\) ..
4.1 Related Work↩︎
4.1.0.1 Multimodal capabilities of diffusion models.
The ability of diffusion models to sample multimodal targets varies substantially across common choices of neural architecture [16], training recipes [25], and sampling algorithms [14], [15]. Existing evaluation metrics [21], [32]–[35] are typically structure-agnostic and may fail to reliably detect mode coverage [36], [37]. In contrast, the diffusion score sensitivity index 1 provides a rigorous framework to reason about multimodal capability.
4.1.0.2 Learning GMMs with Diffusion Score Matching.
Diffusion score matching (DSM) is in principle expressive enough to learn Gaussian mixture: gradient descent recovers component means in the balanced isotropic regime4 [38], learning algorithms can be tailored for density estimation of GMMs with polynomial sample complexity [39], [40]. Here, [39] develop noise-sensitivity bounds similar in spirit to ours. Unlike the aforementioned articles, our work focuses on the ability of diffusion models to sample from target GMMs rather than on density estimation. When the learned score is \(L^2\)-accurate, the resulting diffusion sampler is provably close to the target in total variation [41], [42]. However, this may be insufficient to recover structural parameters that shape the geometry of distributions in low-probability regions. [11] suggest diffusion models may fail to sample multimodal targets due to the lack of expressivity of the Gaussian transition steps in the reverse diffusion process. Our work complements [11] by identifying a different mechanism for multimodal sampling, the diffusion score sensitivity index 1 , and by providing additional insights into the role of noise schedules in alleviating mode amplification.
4.1.0.3 Benefits of Annealing.
Classical (unannealed) score matching [26] is unsuitable for generative modeling of multimodal targets: the asymptotic efficiency of the score matching estimator relative to the maximum-likelihood estimator is governed by the log-Sobolev (LS) constant [9], which grows exponentially in mode separation for Gaussian mixtures [10]. [22] proposed annealing over noise scale to handle multimodality, inspired by [43]. [44] and [45] observe a phase transition at intermediate noise scales linked to the emergence of multimodal structure in practice. Amidst the empirical evidence, [46] establishes the first formal statistical benefit of annealing in the context of diffusion. [46] adapts techniques from the acceleration of Markov Chain mixing to derive a better score matching loss, with the resulting outcome closely resembling the diffusion score matching loss. While [46] establishes sample complexity bounds, our work addresses a complementary question of parameter recovery from generated samples.
5 Sensitivity in real datasets↩︎
We empirically show that mixture weight information is captured in the score functions of intermediate forward process distributions, even outside the Gaussian mixture setting. We additionally explore how altering the noise schedule affects the DSSI and empirically demonstrate that a small DSSI value quantitatively degrades mixture-weight recovery on non-Gaussian data. This is particularly salient to methods like [25], [47], [48] where noise schedules are selected to increase the speed of inference-time sampling; we show that such methods have the potential to amplify some modes over others, augmenting the recent observation by [14].
To study this, we train a simple autoencoder on the ones and eights from MNIST ([49]). Our target distribution \(p_{}^{(\gamma^*)}\) is then the mixture distribution of latent representations of eights and ones (\(p_{}^{(1)}\) is the distribution of ones). We set \(\gamma^*=0.4\) and use deep ResNets [50] with a linear noise schedule as our score matching models.
Figure 4: Influence of noise schedule on mixture weight recovery in MNIST. (a) Linear noise schedule; (b-c) two synthetic noise schedules designed to reduce \(L(\gamma^*)\) (see 11.1). Top row: \(\ell_\mathrm{SM}(p_{t}^{(\gamma)},p_{t}^{(\gamma^*)})\) as a function of \(t\) and \(\gamma\) estimated using trained diffusion models, with \(\gamma^*=0.4\). Bottom row: generated samples using each schedule, with estimated \(\hat{\gamma}\) below. Larger \(L(\gamma^*)\) yields more accurate mixture weight recovery..
5.0.0.1 Sensitivity at intermediate \(t\).
In 4, the top row displays the score sensitivity over time and \(\gamma\) using trained diffusion models (each with a different noise schedule) to approximate the true score function (see 11.3 for more details). The first column uses the same linear noise schedule as in 2 (a). Note that relative to the Gaussian mixture example, \(\ell_\mathrm{SM}(p_{t}^{(\gamma)}, p_{t}^{(\gamma^*)})\) is nonnegligable over a wider range of \(t\) values. We numerically estimate that the sensitivity index is \(\approx 5\times\) larger than in the Gaussian mixture model example, so our theory guarantees a five-fold smaller error in the predicted mixture weight for the same level of error under \(\ell_\mathrm{DSM}\) as in the Gaussian setting.
5.0.0.2 Sensitivity affects parameter recovery.
In 4 (d), the bottom row shows generated samples from a diffusion model with \(T=100\). We estimate the value of \(\gamma^*\) using generated samples and a trained classifier, shown below the generated samples. The generated samples accurately reflect the true value of the mixture weight, with only a \(1\%\) relative error in estimating \(\gamma^*\).
To demonstrate the effect of the score sensitivity index \(L(\gamma^*)\) on the error in estimating \(\gamma^*\), we perform an ablation study on the trained diffusion model by modifying the noise schedule to reduce \(L(\gamma^*)\) as in 3.1 (for details see 11.1). 4 (e) and 4 (f) show the score sensitivity and generated samples for the modified noise schedules, along with the corresponding estimates \(\hat{\gamma}\) from the samples. When the score sensitivity index reduces by a factor of \(10\) (4 (f)), we see that the estimate of \(\gamma^*\) has a relative error of \(31\%\). Surprisingly, the samples from the modified noise schedules have comparable fidelity to those from the linear noise schedule (4 (d)). This is in line with the theoretical and empirical observations that say that diffusion models under errors and accelerated sampling schemes can still accurately learn the geometry of the support of the target distribution [47], [51]–[53]. More details on our MNIST experiments can be found in 11.3.
6 Conclusion↩︎
We investigate the ability of diffusion models to recover the relative mode amplitudes of a multimodal target, despite the insensitivity of the score of the target distribution to mixture weights under high separation. Our analysis clarifies the impact of annealing: the score of the noisy target distribution across noise scales provides the necessary information for accurate sampling of multimodal distributions. Our experiments indicate that inference-time sampling algorithms can be tuned for diversity and fidelity, suggesting a mechanism for tailoring sampling algorithms to data distributions arising in scientific applications.
6.0.0.1 Limitations.
We prove a quantitative relationship between the sensitivity index and the diffusion score error only in the Gaussian mixture setting. Theorem 1 provides an information theoretic mechanism that cannot be computationally verified since it involves minimax bounds. Further, evaluation of the diffusion score-sensitivity index requires knowledge of the true parameter of interest \(\gamma^*\) and the score function of the target distribution \(\mathbf{s}_{t}^{(\gamma^*)}\). Thus, the application to benchmark datasets requires surrogate generative models in place of the unknown target score.
7 Notation↩︎
| Symbol | Definition | In words |
|---|---|---|
| \(\mathcal{P}(p_0, p_1)\) | \(\set{p_{}^{(\gamma)}: p_{}^{(\gamma)} = \gamma \; p_1+(1-\gamma) \; p_0,\;\gamma \in (0,1)}\) | Family of bimodal mixtures of \(p_0\) and \(p_1\) |
| \(p_t\) | \(\mathrm{d} \mathbf{x}_t = -\frac{\beta_t}{2} \mathbf{x}_t \mathrm{d}t + \sqrt{\beta_t}\, \mathrm{d}W_t \text{ and } \mathbf{x}_0 \sim p \Longrightarrow \mathbf{x}_t \sim p_t\) | Distribution of samples at time \(t\) in the forward process starting from \(p\) |
| \(\ell_{\mathrm{NN}}(\mathbf{s}_{\boldsymbol{\theta}}, p)\) | \(\expect[{ \mathbf{x}_0 \sim p,\;t,\;\mathbf{x}_t \mid \mathbf{x}_0}]{\lambda_t\, \norm{\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t, t) - \nabla \log p_{t}^{(\gamma^{*})}(\mathbf{x}_t \mid \mathbf{x}_0)}^2}\) | Diffusion model training objective |
| \(\ell_\mathrm{SM}(p,q)\) | \(\expect[{\mathbf{x}\sim q}]{\norm{\nabla \log p(\mathbf{x}) - \nabla \log q(\mathbf{x})}^2}\) | Classical score matching loss |
| \(\ell_\mathrm{DSM}(p,q)\) | \(\begin{aligned}[t] &\expect[{t \in (0,1),\;\mathbf{x}_t \sim q_t}]{\lambda_t\, \norm{\nabla \log p_t(\mathbf{x}_t) - \nabla \log q_t(\mathbf{x}_t)}^2} \\ &=\expect[{t \in (0,1)}]{\lambda_t\ell_\mathrm{SM}(p,q)} \end{aligned}\) | Diffusion score matching loss |
| \(L(\gamma^*)\) | \(\displaystyle \min_{\gamma'} \frac{\ell_\mathrm{DSM}(p_{}^{(\gamma')},p_{}^{(\gamma^*)})}{(\gamma^* - \gamma')^2}\) | Diffusion score sensitivity index |
| \(\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}, t)\) | \(\forall\, {\boldsymbol{\theta}}\in \boldsymbol{\Theta} = \mathbb{R}^{\#\,\text{Parameters}},\;\mathbf{s}_{\boldsymbol{\theta}}: \mathbb{R}^d \times [0,1] \rightarrow \mathbb{R}^d\) | Neural network with parameters \({\boldsymbol{\theta}}\) |
8 Comparison of sampling algorithms for multimodal distributions↩︎
Sampling from multimodal distributions using classical Langevin sampling is known to be challenging [54], [55]. For a two-component Gaussian mixture, 2 (d) shows that as mode separation increases, the sensitivity of the classical score matching loss to the mixture weight decays rapidly to 0 while the diffusion score matching loss does not. As a consequence, sampling algorithms based solely on the target score are unsuitable for sampling from multimodal distributions. 5 illustrates this effect, comparing classical Langevin sampling to diffusion sampling.
Figure 5: Comparison between different sampling approaches on a high-separation Gaussian mixture model. (a) shows a KDE of the true data distribution. (b)-(c) show, in order, KDEs of samples using (b) Langevin sampling (run to time \(T=10^5\)), (c) diffusion sampling using the analytic score function, and (d) diffusion sampling using a trained neural network. We can see that, unlike Langevin sampling, diffusion sampling accurately reflects the relative masses of the two modes..
9 Minimax upper and lower bounds↩︎
Here we provide a more detailed proof of Theorem 1 and provide additional remarks on the Theorem and the underlying assumptions. We begin with a remark on score-restricted estimators (Definition 1) and explain the motivation for their definition.
In minimax theory, and in particular, Fano’s method, the accuracy of an estimator for a parameter learned from a finite number of data samples is inherently limited by how well the data samples can distinguish between two parameter values. This distinguishability is often measured in terms of the KL divergence between the data distributions at two nearby parameter values. At first glance, the data samples given to a diffusion model appear to be the target samples, and as noted previously, the data distributions with well-separated modes do not distinguish between nearby values of weight parameters. Further, note \(\mathrm{KL}(p_{t}^{(\gamma)}||p_{t}^{(\gamma')}) \leq \mathrm{KL}(p_{}^{(\gamma)}||p_{}^{(\gamma')}).\) As a result, it might seem that the generated samples from the denoising process do not accurately recover the parameter (particularly, given the role that the KL divergence plays in Fano’s inequality [30]), but our empirical results disagree with this expectation. In our empirical results, we see that the noising/denoising process affects parameter recovery compared with score approximation followed by Langevin dynamics. A key idea that helps resolve this mismatch is to recognize that the data a diffusion model uses to accurately estimate parameters is the noisy score (approximated by a neural network), not the noisy sample data. In other words, the generated samples, from the reverse dynamics, are a function of the noisy scores (interpolated by neural networks). These neural network interpolations at some time \(t\) can have a distribution that distinguishes between parameter values more than at time \(0\). In that sense, the time \(t\) scores contain sufficient information to generate samples corresponding to a nearby parameter value, while the time \(0\) score (target score) may not.
Our score-restricted estimators define the notion of recovering \(\gamma^*\) from “score data” rather than sample data. Such a notion models parameter estimation from the generated samples in score-based generative models where the score field is first approximated and then used for sampling. To distinguish Langevin sampling from diffusion models, we define two separate classes of estimators: one class has access to only a target score oracle and the other class has access to a score oracle at an intermediate time along the noising process. Hence, the estimator classes \(\Gamma_{\mathrm{score},t}\) and \(\Gamma_{\mathrm{score},0}\) differ in the data distributions given to them, which are respectively, \((q_{t}^{(\gamma)})\) and \((q_{0}^{(\gamma)}).\) In Lemma 1, we show that the oracle model \(q_{t}^{(\gamma)}\) cannot be reconstructed in a \(\gamma\)-independent manner from the oracle manner \(q_{0}^{(\gamma)},\) and this induces a separation between the two estimator classes.
Remark 4. While in Theorem 1, we are only comparing \(\Gamma_{\mathrm{score},0}\) and a set of estimators that use score data at one other time \(t,\) the result continues to hold if instead under a specific redefinition of the two estimator classes. Let \(t\) still be a time at which the diffusion sensitivity definition holds for the values of \(\epsilon\) and \(\delta\) defined in Theorem 1. The estimator class \(\Gamma_{\mathrm{score},0}\) contains any sampling scheme that does not use samples from \(q_{t'}^{(\gamma)}\), for any time except \(t' = 0.\) On the other hand, the estimators in \(\Gamma_{\mathrm{score},0}\) must include sampling (denoising) steps that access score data from \(q_{t}^{(\gamma)}.\)
Remark 5. To compute the KL divergence, \(\mathrm{KL}(q_{t}^{(\gamma)}||q_{t}^{(\gamma')})\), we first recognize that both Gaussian measures are on \(L^2(p_{t}^{(\gamma)}).\) Hence, by Cameron-Martin formula, the Radon-Nikodym derivative, \(dq_{t}^{(\gamma)}/dq_{t}^{(\gamma')}.\) We can take expectations with respect to \(q_{t}^{(\gamma)}\) to arrive at the formula \(\mathrm{KL}(q_{t}^{(\gamma)}||q_{t}^{(\gamma')}) = \|\mathbf{s}_{t}^{(\gamma)} - \mathbf{s}_{t}^{(\gamma')}\|^2_{L^2(p_{t}^{(\gamma)}}/(2\sigma_t^2).\) This resembles the KL divergence between finite dimensional Gaussian distributions with the same covariance.
Lemma 1. There is no \(\gamma\)-independent post processing of score observations at time 0, \(\hat{\mathbf{s}}^{(\gamma)}_0,\) that can reproduce \(\hat{\mathbf{s}}^{(\gamma)}_t.\)
Proof. We can prove why by contradiction. If there exists a \(\gamma\)-independent processing of \(q_{}^{(\gamma)}\) that produces the distribution \(q_{t}^{(\gamma)},\) the data processing inequality gives \(\mathrm{KL}(q_{t}^{(\gamma)}||q_{t}^{(\gamma')}) \leq \mathrm{KL}(q_{}^{(\gamma)}||q_{}^{(\gamma')}).\) But, we have assumed that at time \(t,\) our mixture family is such that \(\mathrm{KL}(q_{}^{(\gamma)}||q_{}^{(\gamma')}) \leq \epsilon_t^2/4 \leq (1/128) (\gamma - \gamma')^2/A_t^2 \leq (1/128)\mathrm{KL}(q_{t}^{(\gamma)}||q_{t}^{(\gamma')}).\) Thus, the oracle observations and the statistical experiment that can be performed with the model \(q_{t}^{(\gamma)}\) cannot be simulated from having access to the oracle, \(q_{}^{(\gamma)}.\) ◻
Before we give the proof of Theorem 1, we recapitulate a corollary of Generalized Fano’s method from Lemma 3 of [30] and define covering entropy, which is used in Theorem 1.
Lemma 2. [Lack of identifiability of mixture weights] Let \(G:= \{\gamma_1, \cdots, \gamma_N\} \in (0,1)^N\) be a set of \(N\geq 16\) values of mixture weights such that \(|\gamma_i - \gamma_j| \geq \varepsilon,\: \: i \neq j, \; 1\leq i,j \leq N.\) Then, for the class of score-restricted estimators (1), \(\Gamma_{\mathrm{score},0},\) with oracle access to the model, \(q_{}^{(\gamma^*)}\), if \(\mathrm{KL}(q_{}^{(\gamma_i)}||q_{}^{(\gamma_j)}) \leq \log N/4,\) we have \[\begin{align} \inf_{\hat{\gamma} \in \Gamma_{\mathrm{score},0}} \max_{\gamma^* \in G} \mathbb{E}_{q_{}^{(\gamma^*)}} [|\hat{\gamma} - \gamma^*|^2] \geq \varepsilon^2/8. \label{eq:fano} \end{align}\qquad{(4)}\]
We note that the proof follows from [30], which is an application of Fano’s inequality. To get the error in the squared norm, \(\mathbb{E}[|\hat{\gamma} - \gamma^*|^2,\) as we have in ?? , we need to additionally use Markov’s inequality: \(\mathbb{E}[|\hat{\gamma} - \gamma|^2 \geq (\varepsilon/2)^2 \mathbb{P}(|\hat{\gamma} - \gamma^*| > \varepsilon/2) \geq (\varepsilon/2)^2 \: (1/2).\) The second inequality directly follows from the proof of Lemma 3 of [30]. The main difference here is that the data used by the estimator corresponds to the score (1), chosen to exemplify the diffusion model setting.
Definition 4. [\(\epsilon\)-covering entropy]Let \(O_t\) be a finite open cover of \([0,1],\) with centers \(\gamma_1, \cdots, \gamma_{N_t},\) such that for any \(\gamma \in (0,1),\) \(\min_j \mathrm{KL}(q_{t}^{(\gamma)}||q_{t}^{(\gamma_j)}) < \epsilon^2.\) Then, \(V_{\epsilon,t} := \log |O_t|\) is called the \(\epsilon\)-covering entropy.
The main novelty in the proof of Theorem 1 is the setup of score oracle-based estimators. Note that the above standard definitions of packing and covering entropies extend verbatim to our infinite-dimensional statistical models.
Secondly, while annealing has been known to be statistically beneficial [9], [46], our result proves better minimax bounds for parameter recovery, rather than density estimation. To complete the upper bound in ?? , we apply directly Theorem 2 [29]. To obtain the lower bound, we use Lemma 2 with a possibly different packing set of size \(N(\gamma^*)\) around each \(\gamma^*.\) Using the definition of \(\epsilon\)-sensitivity and assuming that \(\log N(\gamma^*) \geq \max\{\log 16, \sigma_t^2\: V(\epsilon_{t}) \}\) at all \(\gamma^*,\) we obtain that for any \(\gamma, \gamma' \in S(\gamma^*),\) \(\mathrm{KL}(q_{}^{(\gamma)}||q_{}^{(\gamma')}) \leq \delta = \epsilon_{t}^2/(4 \sigma_t^2) \leq (\log N(\gamma^*)/V(\epsilon_{t})) \epsilon_{t}^2/4 = \log N(\gamma^*)/4,\) since \(\epsilon_{t}^2 \: \sigma_t^2 = V(\epsilon_{t}).\) As a result, the conditions of Lemma 2 apply at every \(\gamma^*.\)
10 Diffusion process on Gaussian mixture models↩︎
10.1 Mixture Family Preliminaries↩︎
Let \(p_{}^{(0)}\) and \(p_{}^{(1)}\) be two fixed probability distributions on \(\mathbb{R}^d\). We recall the family of mixed probability densities \[\mathcal{P}(p_{}^{(0)},p_{}^{(1)}) \;\mathrel{\vcenter{:}}=\; \big\{\, p_{}^{(\gamma)} \mathrel{\vcenter{:}}= (1-\gamma)\,p_{}^{(0)}({\mathbf{x}}) + \gamma\,p_{}^{(1)}({\mathbf{x}}) \;\big|\; \gamma \in [0,1] \,\big\}.\] The set of noisy marginals at noise level \(t\) produced by the forward process 2 applied to samples \({\mathbf{x}}_0 \sim p_{}^{(\gamma)}\) is denoted \(\{p_{t}^{(\gamma)}\}_{t\in [0,1]}\). The forward process preserves the mixture structure: the noisy marginal \(p_{t}^{(\gamma)}\) is itself a mixture of the noisy component marginals \(p_{t}^{(0)}\) and \(p_{t}^{(1)}\), with the same mixture weight \(\gamma\). The family of noisy mixtures and the noisy family of mixtures therefore coincide.
Lemma 3 (Mixture closure under forward VP-SDE). For every \(\gamma \in (0,1)\) and every \(t \in [0,1]\), \[p_{t}^{(\gamma)} ({\mathbf{x}}) = (1-\gamma) p_{t}^{(0)}({\mathbf{x}}) + \gamma p_{t}^{(1)}({\mathbf{x}}) \label{eq:mixture-closure}\qquad{(5)}\]
Proof. Let \(k_t(\mathbf{y} \mid {\mathbf{x}}) := \mathcal{N}(\mathbf{y};\, \sqrt{\alpha_t}\,{\mathbf{x}},\, (1-\alpha_t)\,\mathbf{I})\) denote the forward conditional density 3 . Since \(k_t\) does not depend on the initial distribution, the marginal at noise level \(t\) is obtained by a linear convolution against the initial density: \[p_{t}^{(\gamma)}({\mathbf{x}}) \;=\; \int_{\mathbb{R}^d} k_t({\mathbf{x}} \mid {\mathbf{x}}_0)\,p_{}^{(\gamma)}({\mathbf{x}}_0)\,\mathrm{d}{\mathbf{x}}_0.\] Substituting the definition of the mixture and applying the linearity of integration, \[\begin{align} p_{t}^{(\gamma)}({\mathbf{x}}) \;&=\; (1-\gamma) \int_{\mathbb{R}^d} k_t({\mathbf{x}} \mid {\mathbf{x}}_0)\,p_{}^{(0)}({\mathbf{x}}_0)\,\mathrm{d}{\mathbf{x}}_0 \;+\; \gamma \int_{\mathbb{R}^d} k_t({\mathbf{x}} \mid {\mathbf{x}}_0)\,p_{}^{(1)}({\mathbf{x}}_0)\,\mathrm{d}{\mathbf{x}}_0 \\ \;&=\; (1-\gamma)\,p_{t}^{(0)}({\mathbf{x}}) + \gamma\,p_{t}^{(1)}({\mathbf{x}}). \end{align}\] ◻
A useful consequence of 3 is that samples from \(p_{t}^{(\gamma)}\) can be drawn by first sampling a component label \(c \sim \mathrm{Ber}(\gamma)\) and then drawing \(\mathbf{y}_t \sim p_{t}^{(c)}\), a component-wise decomposition that simplifies subsequent analysis.
10.1.1 The score of the noisy mixture↩︎
Differentiating \(\log p_{t}^{(\gamma)}\) and applying 3 gives the score of the noisy mixture as a pointwise convex combination of the component scores, \[\nabla_{\mathbf{x}} \log p_{t}^{(\gamma)}({\mathbf{x}}) \;=\; w_0^t({\mathbf{x}})\,\nabla_{\mathbf{x}} \log p_{t}^{(0)}({\mathbf{x}}) \;+\; w_1^t({\mathbf{x}})\,\nabla_{\mathbf{x}} \log p_{t}^{(1)}({\mathbf{x}}), \label{eq:mixture-score-app}\tag{9}\] where the weights are the EM responsibilities — the posterior probabilities under the generative model that \({\mathbf{x}}\) was produced by component \(0\) or component \(1\), respectively [56], \[w_0^t({\mathbf{x}}) \;\mathrel{\vcenter{:}}=\; \frac{(1-\gamma)\,p_{t}^{(0)}({\mathbf{x}})}{p_{t}^{(\gamma)}({\mathbf{x}})}, \qquad w_1^t({\mathbf{x}}) \;\mathrel{\vcenter{:}}=\; \frac{\gamma\,p_{t}^{(1)}({\mathbf{x}})}{p_{t}^{(\gamma)}({\mathbf{x}})}, \label{eq:mixture-weights-app}\tag{10}\] and satisfy \(w_0^t({\mathbf{x}}) + w_1^t({\mathbf{x}}) = 1\).
10.1.2 Specialization: Gaussian components↩︎
When the components are Gaussian, \(p_{}^{(c)} = \mathcal{N}(\boldsymbol{\mu}_c^*, \boldsymbol{\Sigma}_c^*)\) for \(c \in \{0,1\}\), the noisy marginals remain Gaussian, \(p_{t}^{(c)} = \mathcal{N}(\boldsymbol{\mu}_{c,t}^*, \boldsymbol{\Sigma}_{c,t}^*)\), and the score 9 takes the explicit form \[\nabla_{\mathbf{x}} \log p_{t}^{(\gamma)}({\mathbf{x}}) \;=\; \sum_{c \in \{0,1\}} w_c^t({\mathbf{x}})\, (\boldsymbol{\Sigma}_{c,t}^*)^{-1} (\boldsymbol{\mu}_{c,t}^* - {\mathbf{x}}), \label{eq:gaussian-mixture-score}\tag{11}\] a responsibility-weighted sum of precision-scaled displacements \((\boldsymbol{\Sigma}_{c,t}^*)^{-1}(\boldsymbol{\mu}_{c,t}^* - {\mathbf{x}})\). The responsibility weights 10 admit the explicit closed form \[w_c^t({\mathbf{x}}) \;=\; \frac{\pi_c\, \mathcal{N}\!\left({\mathbf{x}};\, \boldsymbol{\mu}_{c,t}^*,\, \boldsymbol{\Sigma}_{c,t}^*\right)}{\sum_{c' \in \{0,1\}} \pi_{c'}\, \mathcal{N}\!\left({\mathbf{x}};\, \boldsymbol{\mu}_{c',t}^*,\, \boldsymbol{\Sigma}_{c',t}^*\right)}, \label{eq:responsibility-weights}\tag{12}\] with mixture weights \(\pi_0 = 1-\gamma\) and \(\pi_1 = \gamma\). The numerator and denominator together carry information about \(\gamma\), the Mahalanobis distance from \({\mathbf{x}}\) to each component mean, and the Gaussian normalizing constants \(\det(\boldsymbol{\Sigma}_{c,t}^*)^{-1/2}\).
10.1.3 DDPM-style forward process↩︎
Lemma 4. For \(p_{}^{(0)}:=\mathcal{N}(\mathbf{x};\mu_0,\mathbf{I})\) and \(p_{}^{(1)}:=\mathcal{N}(\mathbf{x};\mu_1,\mathbf{I})\) denote two individual component distributions. For a variance-preserving forward process with noise schedule \((\bar{\alpha}_t)_{t \in [0,T]}\), the parameters of the mixture modes at time \(t\) are \[\boldsymbol{\mu}_{c,t}^* \;=\; \sqrt{\bar{\alpha}_t}\, \boldsymbol{\mu}_c^*, \qquad \boldsymbol{\Sigma}_{c,t}^* \;=\; \bar{\alpha}_t\, \boldsymbol{\Sigma}_c^* + (1 - \bar{\alpha}_t)\, \mathbf{I}, \label{eq:ddpm-perturbation}\qquad{(6)}\] Hence, the noisy marginal distributions are, \[\begin{align} p_{t}^{(0)}(\mathbf{x}) = \mathcal{N}(\mathbf{x};\sqrt{\alpha_t} \mu_0, \mathbf{I}), \qquad p_{t}^{(1)}(\mathbf{x}) = \mathcal{N}(\mathbf{x};\sqrt{\alpha_t} \mu_1, \mathbf{I}) \end{align}\]
Proof. Let \(\mathbf{x}_0\sim \mathcal{N}(\mu,\mathbf{\Sigma})\). We have that \[\begin{align} \mathbf{x}_t | \mathbf{x}_0 &\sim \mathcal{N}(\sqrt{\alpha_t} \mathbf{x}_0, (1-\alpha_t) \mathbf{I}) \\ \Longrightarrow \mathbf{x}_t &= \sqrt{\alpha_t}\mathbf{x}_0 + \sqrt{1-\alpha_t} \mathbf{z} \end{align}\] where \(\mathbf{z}\sim \mathcal{N}(0,\mathbf{I})\) is independent of \(\mathbf{x}_0\). Therefore, \[\begin{align} \mathbf{x}_t &\sim \mathcal{N}\left(\sqrt{\alpha_t} \mathbf{\mu}, \alpha_t \mathbf{\Sigma }+(1-\alpha_t)\mathbf{I}\right) \end{align}\] ◻
As per ?? \(\alpha_t\) decreases from \(1\) to \(0\), the component means contract toward the origin and the covariances inflate toward \(\mathbf{I}\). Substituting ?? into 12 reveals two limits of interest. At small \(t\) (\(\alpha_t \to 1\)), the perturbed components remain well separated, and the responsibilities concentrate sharply on the closest mode, recovering the M-step regime of the EM algorithm. At large \(t\) (\(\alpha_t \to 0\)), the perturbed components all approach \(\mathcal{N}(\mathbf{0}, \mathbf{I})\) and the responsibilities flatten toward the prior weights \((1-\gamma, \gamma)\), becoming nearly independent of \({\mathbf{x}}\). Sensitivity of the score to the mixture parameter \(\gamma\) is therefore concentrated at the noise scales for which the responsibility weights are neither saturated nor uninformative — a window whose location and width are determined by the geometry of the components.
We now demonstrate an equivalence between the neural network loss 8 and the diffusion score matching loss 7 .
Theorem 3. Let \(p_{}^{(\gamma)}\) be a data distribution on \(\mathbb{R}^d\) with \(\mathbb{E}_{p_{}^{(\gamma)}}\|\mathbf{x}_0\|^2 < \infty\), and let \(\{p_{t}^{(\gamma)}\}_{t\in[0,1]}\) be its marginals under the VP-SDE 2 –3 . Let \(\mathbf{s}_{\boldsymbol{\theta}}:\mathbb{R}^d \times [0,1] \to \mathbb{R}^d\) be a score model with \(\mathbb{E}_{p_{t}^{(\gamma)}}\|\mathbf{s}_{\boldsymbol{\theta}}(\cdot,t)\|^2 < \infty\) for almost every \(t \in [0,1]\). Then \[\ell_{\mathrm{NN}}(p_{}^{({\boldsymbol{\theta}})}, p_{}^{(\gamma)}) \;=\; \ell_\mathrm{DSM}(p_{}^{({\boldsymbol{\theta}})}, p_{}^{(\gamma)}) \;+\; C(p_{}^{(\gamma)}), \label{eq:nn-dsm-equivalence}\qquad{(7)}\] where the constant \(C(p_{}^{(\gamma)})\) does not depend on \({\boldsymbol{\theta}}\) and admits the closed form \[\begin{align} \label{eq:nn-dsm-constant} C(p_{}^{(\gamma)}) &\;=\; \int_0^1 \lambda_t \cdot \frac{\alpha_t}{(1-\alpha_t)^2}\,\mathrm{MMSE}_t(p_{}^{(\gamma)})\, dt,\\ \mathrm{MMSE}_t(p_{}^{(\gamma)}) &\;:=\; \mathbb{E}_{\substack{\mathbf{x}_0 \,\sim\, p_{}^{(\gamma)} \\ \mathbf{x}_t \,\sim\, p_{t}^{(\gamma)}(\cdot\mid\mathbf{x}_0)}}\!\left[\,\big\|\mathbf{x}_0 - \widehat{\mathbf{x}}_0(\mathbf{x}_t)\big\|^2\,\right],\\ \widehat{\mathbf{x}}_0(\mathbf{x}_t) &\;:=\; \mathbb{E}_{\mathbf{x}_0' \,\sim\, p_{}^{(\gamma)}(\cdot\mid\mathbf{x}_t)}[\mathbf{x}_0']. \end{align}\qquad{(8)}\] Moreover, \(C(p_{}^{(\gamma)}) \ge 0\), with equality if and only if \(p_{}^{(\gamma)}\) is a Dirac mass. In particular, \[\ell_{\mathrm{NN}}(p_{}^{({\boldsymbol{\theta}})}, p_{}^{(\gamma)}) \;\ge\; \ell_\mathrm{DSM}(p_{}^{({\boldsymbol{\theta}})}, p_{}^{(\gamma)}) \label{eq:nn-dsm-ordering}\qquad{(9)}\] for every \({\boldsymbol{\theta}}\in \boldsymbol{\Theta}\).
Proof. Our arguments follow the reasoning in [57] with an additional logic for guaranteeing the required relation between the neural network loss and the diffusion score matching loss. Throughout, all expectations are taken under \(p_{}^{(\gamma)}(\mathbf{x}_0)\), \(p_{t}^{(\gamma)}(\mathbf{x}_t\mid \mathbf{x}_0)\) for the joint, and under \(p_{t}^{(\gamma)}(\mathbf{x}_t)\) for the marginal. We suppress the dependence on \(t\) in \(\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t, t)\) when convenient. Fix \(t \in [0,1]\). Expanding the squared norm in the inner expectation of \(\ell_\mathrm{DSM}\), \[\mathbb{E}_{\mathbf{x}_t}\!\big\|\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t) - \mathbf{s}_{t}^{(\gamma)}(\mathbf{x}_t)\big\|^2 = \mathbb{E}_{\mathbf{x}_t}\|\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t)\|^2 \,-\, 2\,b_t({\boldsymbol{\theta}}) \,+\, \mathbb{E}_{\mathbf{x}_t}\!\big\|\mathbf{s}_{t}^{(\gamma)}(\mathbf{x}_t)\big\|^2, \label{eq:proof-esm-expand}\tag{13}\] with \(b_t({\boldsymbol{\theta}}) := \mathbb{E}_{\mathbf{x}_t}\!\big\langle \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t),\, \mathbf{s}_{t}^{(\gamma)}(\mathbf{x}_t)\big\rangle\). Similarly, \[\begin{align} \label{eq:proof-dsm-expand} &\mathbb{E}_{\mathbf{x}_0,\mathbf{x}_t}\!\big\|\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t) - \nabla_{\mathbf{x}_t}\log p_{t}^{(\gamma)}(\mathbf{x}_t \mid \mathbf{x}_0)\big\|^2\\\nonumber &\qquad = \mathbb{E}_{\mathbf{x}_0,\mathbf{x}_t}\|\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t, t)\|^2 \,-\, 2\,a_t({\boldsymbol{\theta}}) \,+\, \mathbb{E}_{\mathbf{x}_0,\mathbf{x}_t}\!\big\|\nabla_{\mathbf{x}_t}\log p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\big\|^2, \end{align}\tag{14}\] with \(a_t({\boldsymbol{\theta}}) := \mathbb{E}_{\mathbf{x}_0,\mathbf{x}_t}\!\big\langle \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t),\, \nabla_{\mathbf{x}_t}\log p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\big\rangle\). The first term in each expansion coincides because \(\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t)\) depends only on \(\mathbf{x}_t\). We now claim that for all \(\theta\), \(b_t({\boldsymbol{\theta}}) = a_t({\boldsymbol{\theta}})\). Using \(\mathbf{s}_{t}^{(\gamma)} = \nabla p_{t}^{(\gamma)}/p_{t}^{(\gamma)}\) and the marginalization \(p_{t}^{(\gamma)}(\mathbf{x}_t) = \int p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\,p_{}^{(\gamma)}(\mathbf{x}_0)\, d\mathbf{x}_0\), \[\begin{align} b_t({\boldsymbol{\theta}}) &= \int p_{t}^{(\gamma)}(\mathbf{x}_t)\,\big\langle \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t),\, \mathbf{s}_{t}^{(\gamma)}(\mathbf{x}_t) \big\rangle\, d\mathbf{x}_t \\ &= \int \big\langle \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t),\, \nabla_{\mathbf{x}_t}p_{t}^{(\gamma)}(\mathbf{x}_t)\big\rangle\, d\mathbf{x}_t \\ &= \int \Big\langle \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t),\, \nabla_{\mathbf{x}_t}\!\!\int p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\,p_{}^{(\gamma)}(\mathbf{x}_0)\, d\mathbf{x}_0\Big\rangle\, d\mathbf{x}_t \\ &= \iint \big\langle \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t),\, \nabla_{\mathbf{x}_t}p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\big\rangle\,p_{}^{(\gamma)}(\mathbf{x}_0)\, d\mathbf{x}_0\, d\mathbf{x}_t \\ &= \iint p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\,p_{}^{(\gamma)}(\mathbf{x}_0)\,\big\langle \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t,t),\, \nabla_{\mathbf{x}_t}\log p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\big\rangle\, d\mathbf{x}_0\, d\mathbf{x}_t \\ &= a_t({\boldsymbol{\theta}}), \end{align}\] where Leibniz’s rule applies by smoothness and Gaussian decay of the forward kernel, and Fubini’s theorem applies due to the regularity assumption on \(\mathbf{s}_{{\boldsymbol{\theta}}}\). Combining the above, \[\mathbb{E}_{\mathbf{x}_0,\mathbf{x}_t}\!\big\|\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t, t) - \nabla_{\mathbf{x}_t}\log p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\big\|^2 - \mathbb{E}_{\mathbf{x}_t}\!\big\|\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}_t, t) - \mathbf{s}_{t}^{(\gamma)}(\mathbf{x}_t) \big\|^2 = \mathrm{gap}(t), \label{eq:proof-residual}\tag{15}\] where \[\mathrm{gap}(t) := \mathbb{E}_{\mathbf{x}_0,\mathbf{x}_t}\!\big\|\nabla_{\mathbf{x}_t}\log p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\big\|^2 - \mathbb{E}_{\mathbf{x}_t}\!\big\|\mathbf{s}_{t}^{(\gamma)}(\mathbf{x}_t)\big\|^2, \label{eq:proof-Dt-def}\tag{16}\] which is independent of \({\boldsymbol{\theta}}\). The Gaussian forward kernel 3 gives \[\nabla_{\mathbf{x}_t}\log p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0) = -\frac{\mathbf{x}_t - \sqrt{\alpha_t}\mathbf{x}_0}{1-\alpha_t} = -\frac{\mathbf{z}}{\sqrt{1-\alpha_t}},\quad \mathbf{z}\sim\mathcal{N}(\mathbf{0},\mathbf{I}), \label{eq:proof-cond-score}\tag{17}\] hence the first term of 16 equals \[\mathbb{E}~\!\big\|\nabla_{\mathbf{x}_t}\log p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\big\|^2 = \frac{d}{1-\alpha_t}. \label{eq:proof-cond-norm}\tag{18}\] For the marginal score, take expectations of 17 conditional on \(\mathbf{x}_t\) we obtain Tweedie’s formula: \[\mathbf{s}_{t}^{(\gamma)}(\mathbf{x}_t) = \mathbb{E}~\!\big[\nabla_{\mathbf{x}_t}\log p_{t}^{(\gamma)}(\mathbf{x}_t\mid\mathbf{x}_0)\,\big|\,\mathbf{x}_t\big] = \frac{\sqrt{\alpha_t}\,\mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t] - \mathbf{x}_t}{1-\alpha_t}. \label{eq:proof-tweedie}\tag{19}\] Setting \(\hat{\boldsymbol{\eta}}(\mathbf{x}_t) := \sqrt{\alpha_t}\,\mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t] = \mathbb{E}[\sqrt{\alpha_t}\mathbf{x}_0 \mid \mathbf{x}_t]\), we see that, \[\mathbb{E}~\!\big\|\mathbf{s}_{t}^{(\gamma)}(\mathbf{x}_t)\big\|^2 = \frac{\mathbb{E}\|\mathbf{x}_t - \hat{\boldsymbol{\eta}}(\mathbf{x}_t)\|^2}{(1-\alpha_t)^2}. \label{eq:proof-score-norm}\tag{20}\] Decomposing \(\mathbf{x}_t - \hat{\boldsymbol{\eta}} = (\mathbf{x}_t - \sqrt{\alpha_t}\mathbf{x}_0) + (\sqrt{\alpha_t}\mathbf{x}_0 - \hat{\boldsymbol{\eta}}) = \sqrt{1-\alpha_t}\,\mathbf{z} + \sqrt{\alpha_t}(\mathbf{x}_0 - \mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t])\) and expanding the expectation we get, \[\mathbb{E}\|\mathbf{x}_t - \hat{\boldsymbol{\eta}}\|^2 = (1-\alpha_t)\,d + 2\sqrt{\alpha_t(1-\alpha_t)}\,\mathbb{E}~\!\big\langle\mathbf{z},\,\mathbf{x}_0 - \mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t]\big\rangle + \alpha_t\,\mathrm{MMSE}_t(p_{}^{(\gamma)}). \label{eq:proof-decomp}\tag{21}\] For the cross term, use \(\mathbf{z} = (\mathbf{x}_t - \sqrt{\alpha_t}\mathbf{x}_0)/\sqrt{1-\alpha_t}\) and orthogonality of conditional expectation, \(\mathbb{E}~\!\big[(\mathbf{x}_0 - \mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t])\,\widehat{\mathbf{x}}_0(\mathbf{x}_t)\big] = 0\) for any \(g\): \[\begin{align} \mathbb{E}~\!\big\langle\mathbf{z},\,\mathbf{x}_0 - \mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t]\big\rangle &= \tfrac{1}{\sqrt{1-\alpha_t}}\,\mathbb{E}~\!\big\langle \mathbf{x}_t - \sqrt{\alpha_t}\mathbf{x}_0,\,\mathbf{x}_0 - \mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t]\big\rangle \\ &= \tfrac{1}{\sqrt{1-\alpha_t}}\, \widehat{\mathbf{x}}_0(0 - \sqrt{\alpha_t}\,\mathrm{MMSE}_t(p_{}^{(\gamma)})\big) \\ &= -\frac{\sqrt{\alpha_t}\,\mathrm{MMSE}_t(p_{}^{(\gamma)})}{\sqrt{1-\alpha_t}}. \end{align}\] Substituting into 21 , \[\mathbb{E}\|\mathbf{x}_t - \hat{\boldsymbol{\eta}}\|^2 = (1-\alpha_t)\,d - \alpha_t\,\mathrm{MMSE}_t(p_{}^{(\gamma)}),\] and therefore via 20 , \[\mathbb{E}~\!\big\|\mathbf{s}_{t}^{(\gamma)}(\mathbf{x}_t)\big\|^2 = \frac{d}{1-\alpha_t} - \frac{\alpha_t\,\mathrm{MMSE}_t(p_{}^{(\gamma)})}{(1-\alpha_t)^2}. \label{eq:proof-score-final}\tag{22}\] Combining 18 and 22 in 16 , \[\mathrm{gap}(t) = \frac{\alpha_t\,\mathrm{MMSE}_t(p_{}^{(\gamma)})}{(1-\alpha_t)^2} \;\ge\; 0. \label{eq:proof-Dt-final}\tag{23}\] Hence, multiplying 15 by \(\lambda_t\) and integrating over \(t\sim \mathrm{Unif}(0,1)\), the difference in the neural network loss and the diffusion score matching loss can be written as, \[\ell_{\mathrm{NN}}(p_{}^{({\boldsymbol{\theta}})}, p_{}^{(\gamma)}) - \ell_\mathrm{DSM}(p_{}^{({\boldsymbol{\theta}})}, p_{}^{(\gamma)}) = \int_0^1 \lambda_t\,\mathrm{gap}(t)\,dt = \int_0^1 \lambda_t\,\frac{\alpha_t\,\mathrm{MMSE}_t(p_{}^{(\gamma)})}{(1-\alpha_t)^2}\,dt =: C(p_{}^{(\gamma)}),\] proving ?? –?? . For the equality case, note that \(C(p_{}^{(\gamma)}) = 0\) iff \(\mathrm{MMSE}_t(p_{}^{(\gamma)}) = 0\) for almost every \(t \in (0,1)\). The latter requires \(\mathbf{x}_0 = \mathbb{E}[\mathbf{x}_0\mid\mathbf{x}_t]\) almost surely, i.e., \(\mathbf{x}_0\) is \(\sigma(\mathbf{x}_t)\)-measurable. Since \(\mathbf{x}_t = \sqrt{\alpha_t}\mathbf{x}_0 + \sqrt{1-\alpha_t}\mathbf{z}\) with \(\mathbf{z}\) independent of \(\mathbf{x}_0\) and \(1-\alpha_t > 0\), this forces \(\mathbf{x}_0\) to be deterministic. ◻
10.2 Proof of 2↩︎
Remark 6. In this theorem, we use \(\lambda_t=1\) rather than \(\lambda_t= 1-\alpha_t\). The value of \(t\) where \(\ell_\mathrm{SM}(p_{t}^{(\gamma)},p_{t}^{(\gamma^*)})\) is large correspond to when \(\alpha_t\) is close to \(0\) for when the modes are reasonably well-separated, so that \(1-\alpha_t\) is close to \(1\). The same analysis can be performed in the presence of any weighting function \(\lambda_t\) if one is willing to accept a more complicated expression for the constant \(C\); we elaborate on this in 6. For all of our numerical results, we use the convention that \(\lambda_t=1-\alpha_t\) to show that our results still hold for the weighting factor used in practice.
Throughout the proof, let \(\Delta:=\norm{\mu_1-\mu_0}\). From 3 and 4, we can generate samples from \(p_{t}^{(\gamma)}\) by sampling \(c\sim \mathrm{Ber}(\gamma)\) and sampling \(\mathbf{y}_t\sim p_{t}^{(c)}\). Furthermore, the next lemma will show that we can generate these samples using the Probability Flow ODE ([23]) instead of the stochastic forward process.
First, note that \[\begin{align} \label{eq:32mode32decomposition} \ell_\mathrm{DSM}(p_{}^{(\hat{\gamma})},p_{}^{(\gamma^*)})&= \expect[t,\mathbf{x}_t\sim p_{t}^{(\gamma^*)}]{\norm{\mathbf{s}_{t}^{(\gamma^*)}(\mathbf{x}_t)-\mathbf{s}_{t}^{(\hat{\gamma})}}(\mathbf{x}_t)} \\ &= (1-\gamma^*)\expect[t,\mathbf{x}_t\sim p_{t}^{(\gamma^*)}|c=0]{\norm{\mathbf{s}_{t}^{(\gamma^*)}(\mathbf{x}_t)-\mathbf{s}_{t}^{(\hat{\gamma})}}(\mathbf{x}_t)} \\ &\quad\quad + \gamma^*\expect[t,\mathbf{x}_t\sim p_{t}^{(\gamma^*)}|c=1]{\norm{\mathbf{s}_{t}^{(\gamma^*)}(\mathbf{x}_t)-\mathbf{s}_{t}^{(\hat{\gamma})}}(\mathbf{x}_t)} \\ &= (1-\gamma^*)\expect[t,\mathbf{x}_t\sim p_{t}^{(0)}]{\norm{\mathbf{s}_{t}^{(\gamma^*)}(\mathbf{x}_t)-\mathbf{s}_{t}^{(\hat{\gamma})}}(\mathbf{x}_t)} \\ &\quad\quad + \gamma^*\expect[t,\mathbf{x}_t\sim p_{t}^{(1)}]{\norm{\mathbf{s}_{t}^{(\gamma^*)}(\mathbf{x}_t)-\mathbf{s}_{t}^{(\hat{\gamma})}}(\mathbf{x}_t)} \end{align}\tag{24}\]
In particular, we can lower-bound \(\ell_\mathrm{DSM}\) by independently lower-bounding each of these expectations.
Lemma 5. Consider any distribution \(\pi\) and let \(\pi_t\) be the noisy marginal at time \(t\) of the forward process. Then, if \(\mathbf{y}_t\) satisfies the Probability Flow ODE \[\label{eqn:32Probability32Flow32ODE} \mathrm{d}\mathbf{y}_t = -\frac{1}{2}\beta_t\left(\mathbf{y}_t + \nabla \log p_{\mathbf{y}_t}(\mathbf{y}_t)\right)\mathrm{d}t\qquad{(10)}\] and \(\mathbf{y}_0\sim \pi\), then \(\mathbf{y}_t\sim \pi_t\).
Proof. Let \(\mathbf{x}_t\) satisfy the Variance Preserving SDE and \(\mathbf{y}_t\) be as in the lemma statement. Then, we have the Fokker-Planck equations \[\begin{align} \frac{\partial p_{\mathbf{x}_t} }{\partial t}&= -\nabla \cdot \left(-\frac{1}{2}\beta_t \mathbf{x}p_{\mathbf{x}_t} \right) + \frac{1}{2}\beta_t\Delta p_{\mathbf{x}_t} \\ &=\frac{1}{2}\beta_t \nabla \cdot \left(\mathbf{x}p_{\mathbf{x}_t} \right) + \frac{1}{2}\beta_t\Delta p_{\mathbf{x}_t} \\ \\ \frac{\partial p_{\mathbf{y}_t} }{\partial t} &= \frac{1}{2}\beta_t \nabla \cdot \left(\mathbf{y} p_{\mathbf{y}_t} \right) +\frac{1}{2}\beta_t \nabla \cdot \left(p_{\mathbf{y}_t}\nabla \log p_{\mathbf{y}_t} \right) \\ &= \frac{1}{2}\beta_t \nabla \cdot \left(\mathbf{y} p_{\mathbf{y}_t} \right) +\frac{1}{2}\beta_t \Delta p_{\mathbf{y}_t} \end{align}\] ◻
The Probability Flow ODE (?? ) substantially simplifies the analysis because of the following fact:
Lemma 6. For \(\mathbf{y}_0\), and \(\mathbf{y}_t\) following the Probability Flow ODE associated with \(c\) and starting from \(\mathbf{y}_0\sim p_{}^{(c)}\), we have that \[\mathbf{y}_t - \sqrt{\alpha_t}\mu_c = \mathbf{y}_0 - \mu_c\]
Proof. Note that \(\frac{\mathrm d}{\mathrm d t}\sqrt{\alpha_t}=\frac{1}{2}\beta_t \sqrt{\alpha_t}\). Therefore, \[\begin{align} \frac{\mathrm d}{\mathrm dt}(\mathbf{y}_t -\sqrt{\alpha_t}\mu_c) &=-\frac{1}{2}\beta_t \left(\mathbf{y}_t + \nabla \log p_{t}^{(c)}(\mathbf{y}_t)\right) - \frac{1}{2}\beta_t \sqrt{\alpha_t} \mu_c \\ &=-\frac{1}{2}\beta_t \left(\mathbf{y}_t + \sqrt{\alpha_t} \mu_c - \mathbf{y}_t - \sqrt{\alpha_t} \mu_c\right) \\ &=0 \end{align}\] ◻
We now focus on the case where \(c=0\) (the case of \(c=1\) will follow via symmetry). Let \(\mathbf{z}_t = \mathbf{y}_t|c=0\). Our first goal is to lower bound \(\norm{\mathbf{s}_{t}^{(\gamma^*)}(\mathbf{z}_t) - \mathbf{s}_{t}^{(\hat{\gamma})}(\mathbf{z}_t)}^2\) pointwise in \(t\) and \(\mathbf{z}_t\). We first show an analytic expression for this score difference, then identify a range of \(t\)’s along each path where this expression is sufficiently large. Finally, by integrating over initial conditions \(\mathbf{z}_0\), we get a lower bound for \(\expect[t,\mathbf{z}_t\sim p_{t}^{(0)}]{\norm{\mathbf{s}_{t}^{(\gamma^*)}(\mathbf{z}_t)-\mathbf{s}_{t}^{(\hat{\gamma})}(\mathbf{z}_t)}^2}\).
Proposition 4. Let \(\hat{\gamma}, \gamma^*\in [0,1]\). Then, \[\norm{\mathbf{s}_{t}^{(\hat{\gamma})}(\mathbf{z}_t) - \mathbf{s}_{t}^{(\gamma^{*})}(\mathbf{z}_t)}^2 = ({\hat{\gamma}}-\gamma^*)^2 \frac{\norm{\mu_0-\mu_1}^2 \alpha_te^{2f_0(\mathbf{z}_0, t)}}{(1-{\hat{\gamma}} + {\hat{\gamma}}e^{f_0(\mathbf{z}_0, t)})^2 (1-\gamma^* + \gamma^* e^{f_0(\mathbf{z}_0, t)})^2}\] where \[f_0(\mathbf{z}_0, t):=-\frac{1}{2}\norm{\mu_1-\mu_0}^2\alpha_t + \left\langle\mathbf{z}_0 - \mu_0, \mu_1 - \mu_0\right\rangle\sqrt{\alpha_t}\]
Proof. We \[p_{t}^{(0)}(\mathbf{x})=\frac{1}{Z_t}\exp\left(-\frac{1}{2}\norm{\mathbf{x}-\sqrt{\alpha_t}\mu_0}^2\right)\] \[p_{t}^{(1)}(\mathbf{x})=\frac{1}{Z_t}\exp\left(-\frac{1}{2}\norm{\mathbf{x}-\sqrt{\alpha_t}\mu_1}^2\right)\]
Expanding the log likelihoods, we have \[\begin{align} \log \frac{p_1(\mathbf{z}_t)}{p_0(\mathbf{z}_t)}&= -\frac{1}{2} \norm{\mu_0-\mu_1}^2 \alpha_t + \left\langle\mathbf{z}_0 - \mu_0, \mu_1 - \mu_0\right\rangle\sqrt{\alpha_t} \end{align}\]
Let \(f(\mathbf{z}_0,t) :=\log \frac{p_1(\mathbf{z}_t)}{p_0(\mathbf{z}_t)}\). Then, \[\begin{align} \mathbf{s}_{t}^{(\gamma)}(\mathbf{z}_t)&= \frac{(1-\gamma)\;p_{t}^{(0)}(\mathbf{z}_t)\mathbf{s}_{t}^{(0)}(\mathbf{z}_t) + \gamma\;p_{t}^{(1)}(\mathbf{z}_t)\mathbf{s}_{t}^{(1)}(\mathbf{z}_t)}{p_{t}^{(\gamma)}(\mathbf{z}_t)} \\ &=\mathbf{s}_{t}^{(0)}(\mathbf{z}_t) + \frac{ \gamma\;p_{t}^{(1)}(\mathbf{z}_t)\left(\mathbf{s}_{t}^{(1)}(\mathbf{z}_t)-\mathbf{s}_{t}^{(0)}(\mathbf{z}_t)\right)}{p_{t}^{(\gamma)}(\mathbf{z}_t)} \\ &=\mathbf{s}_{t}^{(0)}(\mathbf{z}_t) + \frac{ \gamma\;p_{t}^{(1)}(\mathbf{z}_t)}{p_{t}^{(\gamma)}(\mathbf{z}_t)}\sqrt{\alpha_t}\left( \mu_1- \mu_0\right) \\ &=\mathbf{s}_{t}^{(0)}(\mathbf{z}_t) + \frac{ \gamma\;e^{f(\mathbf{z}_0,t)}}{1-\gamma+\gamma \; e^{f(\mathbf{z}_0,t)}}\sqrt{\alpha_t}\left( \mu_1- \mu_0\right) \end{align}\] where we used \(\nabla \log p_{t}^{(c)}(\mathbf{x}) = \sqrt{\alpha_t}\mu_c-\mathbf{x}\). Therefore, \[\begin{align} &\mathbf{s}_{t}^{(\hat{\gamma})}(\mathbf{z}_t) - \mathbf{s}_{t}^{(\gamma^{*})}(\mathbf{z}_t)= \frac{ ({\hat{\gamma}}-\gamma^*)\;e^{f(\mathbf{z}_0,t)}}{\left(1-{\hat{\gamma}}+{\hat{\gamma}} \; e^{f(\mathbf{z}_0,t)}\right)\left(1-\gamma^*+\gamma^* \; e^{f(\mathbf{z}_0,t)}\right)}\sqrt{\alpha_t}\left( \mu_1- \mu_0\right) \end{align}\]
This gives \[\norm{\mathbf{s}_{t}^{(\hat{\gamma})}(\mathbf{z}_t) -\mathbf{s}_{t}^{(\gamma^{*})}(\mathbf{z}_t)}^2=({\hat{\gamma}}-\gamma^*)^2\frac{ \norm{ \mu_1- \mu_0}^2\alpha_t\;e^{2f(\mathbf{z}_0,t)}}{\left(1-{\hat{\gamma}}+{\hat{\gamma}} \; e^{f(\mathbf{z}_0,t)}\right)^2\left(1-\gamma^*+\gamma^* \; e^{f(\mathbf{z}_0,t)}\right)^2}\] ◻
Lemma 7. For all \(x\in \mathbb{R}\) and \(\gamma\in [0,1]\), we have \[\frac{e^x}{(1-\gamma + \gamma e^x)^2}\geq \frac{1}{4\gamma (1-\gamma)}e^{-\frac{1}{4}\left(x-\log \frac{1-\gamma}{\gamma}\right)^2}\] with equality only when \(x=\log \frac{1-\gamma}{\gamma}\). In particular, \[\begin{align} &\frac{e^{2f_0(\mathbf{z}_0, t)}}{\left(1-{\hat{\gamma}} + {\hat{\gamma}} e^{f_0(\mathbf{z}_0, t)}\right)^2 \left(1-\gamma^* + \gamma^* e^{f_0(\mathbf{z}_0, t)}\right)^2}\geq \frac{e^{-\frac{1}{4}\left(f_0(\mathbf{z}_0,t)-\log \frac{1-{\hat{\gamma}}}{{\hat{\gamma}}}\right)^2-\frac{1}{4}\left(f_0(\mathbf{z}_0,t)-\log \frac{1-\gamma^*}{\gamma^*}\right)^2}}{16 {\hat{\gamma}} (1-{\hat{\gamma}}) \gamma^* (1-\gamma^*)} \end{align}\]
Proof. Let \(g(x)= \frac{e^x}{(1-\gamma + \gamma e^x)^2}\) and \(h(x)=\frac{1}{4\gamma (1-\gamma)}e^{-\frac{1}{4}\left(x-\log \frac{1-\gamma}{\gamma}\right)^2}\). Then, \[\begin{align} \frac{\mathrm d}{\mathrm d x}\log g(x) &=\frac{1-\gamma -\gamma e^x}{1-\gamma + \gamma e^x} \end{align}\]
For all \(\gamma\in [0,1]\) \(1-\gamma + \gamma e^x>0\). Therefore, \(x^*:=\log \frac{1-\gamma}{\gamma}\) is the only critical point of \(g(x)\), and \(g(x^*)=\frac{1}{4\gamma\left(1-\gamma\right)}\). Similarly \(h(x^*)=\frac{1}{4\gamma\left(1-\gamma\right)}\) and \(\frac{\mathrm d}{\mathrm d x}\log h(x^*)=0\). We have \[\begin{align} \frac{\mathrm{d}}{\mathrm{d}x}[\log h(x)-\log g(x)] &=-\frac{1}{2}(x-x^*) - \frac{1-e^{x-x^*}}{1+e^{x-x^*}} \\ \frac{\mathrm{d}^2}{\mathrm{d}x^2}[\log h(x)-\log g(x)] &= -\frac{(1-e^{x-x^*})^2}{2(1+e^{x-x^*})^2} \end{align}\]
So, \(\frac{\mathrm{d}^2}{\mathrm{d}x^2}[\log h(x)-\log g(x)]\leq 0\) with equality only when \(x=x^*\). Therefore, \(\log h-\log g\) is strictly concave away from \(x^*\) and maximized at \(x^*\), where \(h(x^*)=g(x^*)\). Therefore, \(h< g\). ◻
Lemma 8. Let \(\mathbf{z}_t\) and \(\mathbf{z}_0\) be as above. Let \(M(\mathbf{z}_0):=\left\langle\mathbf{z}_0 - \mu_0, \mu_1 - \mu_0\right\rangle\), and define the log-odds ratios \[\begin{align} \hat{r} &:= \log \frac{1-{\hat{\gamma}}}{{\hat{\gamma}}} \\ r^* &:= \log \frac{1-\gamma^*}{\gamma^*} \end{align}\]
Let \(\epsilon>0\) be such that \[\begin{align} \epsilon &\leq \exp\left(-\frac{1}{8}\left[(\hat{r}-r^*)^2 + \left(\hat{r}+r^*-\frac{M(\mathbf{z}_0)^2}{\Delta^2}\right)_+^2\right]\right) \end{align}\] where \((\cdot)_+:= \max\set{\cdot, 0}\).
Let \[\begin{align} a&=\max\left\{\frac{M(\mathbf{z}_0)}{\Delta^2} + \frac{1}{\Delta}\sqrt{\left(\frac{M(\mathbf{z}_0)^2}{\Delta^2} - \left( \hat{r} + r^*\right)-\sqrt{8\log \frac{1}{\epsilon}-\left(\hat{r}-r^*\right)^2}\right)_+},\sqrt{\alpha_1}\right\} \\ b&= \min\left\{\frac{M(\mathbf{z}_0)}{\Delta^2} + \frac{1}{\Delta}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} - \left( \hat{r} + r^*\right)+\sqrt{8\log \frac{1}{\epsilon}-\left(\hat{r}-r^*\right)^2}},1\right\} \end{align}\]
Then for all \(t\) such that \(\sqrt{\alpha_t}\in [a,b]\), \[e^{-\frac{1}{4}\left(f_0(\mathbf{z}_0,t)-\hat{r} \right)^s2-\frac{1}{4}\left(f_0(\mathbf{z}_0,t)-r^*\right)^2} \geq \epsilon\]
Proof. First, the inequality is equivalent to \[\begin{align} \left(f_0(\mathbf{z}_0,t)-r^*\right)^2+\left(f_0(z_0,t)-r^*\right)^2\leq 4\log \frac{1}{\epsilon} \end{align}\]
Plugging in the definition of \(f\), the inequality holds whenever \[\left[\frac{\Delta^2}{2} \left(\sqrt{\alpha_t}-\frac{M(\mathbf{z}_0)}{\Delta^2}\right)^2 - \frac{M(\mathbf{z}_0)^2}{2\Delta^2}+ \frac{\hat{r}+r^*}{2}\right]^2 \leq 2\log \frac{1}{\epsilon}-\frac{(\hat{r}-r^*)^2}{4}\]
We wish to identify the values of \(\sqrt{\alpha_t}\) which make this inequality hold. This function is continuous, so its sign can only change at the roots of the corresponding equality constraint: \[\left[\frac{\Delta^2}{2} \left(\sqrt{\alpha_t}-\frac{M(\mathbf{z}_0)}{\Delta^2}\right)^2 - \frac{M(\mathbf{z}_0)^2}{2\Delta^2}+ \frac{\hat{r}+r^*}{2}\right]^2 = 2\log \frac{1}{\epsilon}-\frac{(\hat{r} -r^*)^2}{4}\]
Solving this quartic gives the roots \[\sqrt{\alpha_t}=\frac{M(\mathbf{z}_0)}{\Delta^2} \pm\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^4}-\frac{\hat{r}+r^*}{\Delta^2}\pm \frac{1}{\Delta^2}\sqrt{8\log \frac{1}{\epsilon}-(\hat{r}-r^*)^2}}\]
The condition on \(\epsilon\) is such that the square roots are all well-defined. Let \(b'\) be the largest of these roots (both \(\pm\)’s are set to \(+\)). As \(\sqrt{\alpha_t}\rightarrow \infty\), we would have \(f_0(\mathbf{z}_0,t)\rightarrow -\infty\) (because \(-\Delta^2/2\) is strictly negative), and hence \[e^{-\frac{1}{4}\left(f_0(\mathbf{z}_0,t)-\hat{r} \right)^2-\frac{1}{4}\left(f_0(\mathbf{z}_0,t)-r^*\right)^2} \rightarrow 0\]
Therefore, for \(\sqrt{\alpha_t}>b'\) the inequality does not hold. Note that the four roots above are distinct (with two being complex for \(\epsilon\) sufficiently small), meaning that each root has multiplicity 1, and hence the polynomial changes signs at \(b'\). Therefore, for \(\sqrt{\alpha_t}\) between \(b'\) and the next largest real root, the inequality holds. The next largest real root is \[a'=\begin{cases} \frac{M(\mathbf{z}_0)}{\Delta^2} +\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^4}-\frac{\hat{r}+r^*}{\Delta^2}-\frac{1}{\Delta^2}\sqrt{8\log \frac{1}{\epsilon}-(\hat{r}-r^*)^2}} & \text{if}\quad \frac{M(\mathbf{z}_0)^2}{\Delta^2}-(\hat{r}+r^*)-\sqrt{8\log \frac{1}{\epsilon}-(\hat{r}-r^*)^2} \geq0 \\ \frac{M(\mathbf{z}_0)}{\Delta^2} -\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^4}-\frac{\hat{r}+r^*}{\Delta^2}+\frac{1}{\Delta^2}\sqrt{8\log \frac{1}{\epsilon}-(\hat{r}-r^*)^2}} & \text{otherwise} \end{cases}\]
In the second case we have \(a'\leq \frac{M(\mathbf{z}_0)}{\Delta^2}\), so \[a'\leq \frac{M(\mathbf{z}_0)}{\Delta^2} +\sqrt{\left(\frac{M(\mathbf{z}_0)^2}{\Delta^4}-\frac{\hat{r}+r^*}{\Delta^2}-\frac{1}{\Delta^2}\sqrt{8\log \frac{1}{\epsilon}-(\hat{r}-r^*)^2}\right)_+}=:a''\]
Therefore, the inequality holds on \([a',b']\) and \([a'',b']\subseteq [a',b']\), so the inequality holds on \([a'',b']\). Note that for \(t\in [0,1]\), \(\alpha_t\in [\alpha_1, 1]\) where \(\alpha_1 >0\) is small, so we let \(b:= \min\set{b', 1}\) and \(a:=\max\set{a,\sqrt{\alpha_1}}\). ◻
Proposition 5. For all \(\epsilon\) satisfying the condition in 8, for all \(t\) s.t. \(\sqrt{\alpha_t}\in[a,b]\) (with \(a\) and \(b\) as defined in the previous lemma) we have \[\begin{align} \norm{\mathbf{s}_{t}^{(\hat{\gamma})}(\mathbf{z}_t) - \mathbf{s}_{t}^{(\gamma^{*})}(\mathbf{z}_t)}^2\geq\frac{(\hat{\gamma}-\gamma^*)^2\norm{ \mu_1- \mu_0}^2 \alpha_t}{16 \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)}\epsilon \end{align}\]
Proof. Combining the above two lemmas, \[\begin{align} \norm{\mathbf{s}_{t}^{(\hat{\gamma})}(\mathbf{z}_t) - \mathbf{s}_{t}^{(\gamma^{*})}(\mathbf{z}_t)}^2 &\geq\frac{(\hat{\gamma}-\gamma^*)^2\norm{ \mu_1- \mu_0}^2 \alpha_t}{16 \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)}e^{-\frac{1}{4}\left(f_0(\mathbf{z}_0,t)-\hat{r}\right)^2-\frac{1}{4}\left(f_0(\mathbf{z}_0,t)-r^*\right)^2} \\ &\geq\frac{(\hat{\gamma}-\gamma^*)^2\norm{ \mu_1- \mu_0}^2 \alpha_t}{16 \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)}\epsilon \end{align}\] ◻
We now turn our attention to lower bounding \[\int_0^1\norm{\mathbf{s}_{t}^{(\hat{\gamma})}(\mathbf{z}_t) - \mathbf{s}_{t}^{(\gamma^{*})}(\mathbf{z}_t)}^2\mathrm{d}t\] along a given path starting from \(\mathbf{z _ 0}\). From the lemmas above, we have \[\int_0^1 \norm{\nabla \log p_{t}^{(\hat{\gamma})}(\mathbf{z}_t) - \nabla \log p_{t}^{(\gamma^{*})}(\mathbf{z}_t)}^2\mathrm{d}t > \frac{({\hat{\gamma}}-\gamma^*)^2\norm{ \mu_1- \mu_0}^2}{16{\hat{\gamma}}(1-{\hat{\gamma}}) \gamma^* (1-\gamma^*)}\epsilon\int_{t_{b^2}}^{t_{a^2}}\alpha_t\mathrm{d}t\] where \(t_{a^2}\) and \(t_{b^2}\) are such that \(\alpha_{t_{a^2}}={a^2}\) and \(\alpha_{t_{b^2}}={b^2}\). Note \(\alpha_t\) is a decreasing function of \(t\), hence why the bounds of the integral are reversed. If we let \(\alpha^{-1}(\tau)\) be the inverse function of \(\alpha_t\), such that \(\alpha_{\alpha^{-1}(\tau)} = \tau\), then \[\int_{t_{b^2}}^{t_{a^2}}\alpha_t\mathrm{d}t=2\int_{a}^b \frac{\tau}{\beta_{\alpha^{-1}(\tau^2)}}\mathrm{d}\tau\]
In particular, if we consider the linear noise schedule \[\beta_t:=\beta_1 t\] then \[\int_{t_{b^2}}^{t_{a^2}}\alpha_t\mathrm{d}t=\frac{1}{\sqrt{\beta_1}}\int_{a}^b \frac{\tau}{\sqrt{-\ln \tau}}\mathrm{d}\tau\]
Proposition 6. Let \(a\), \(b\), \(\Delta\), \(z_0\), \(\hat{r}\), \(r^*\), and \(M\), be as above, and assume \(\epsilon\) satisfies the bound from 8. Define \[R_0:=\sqrt{8\ln \frac{1}{\epsilon}-(\hat{r}-r^*)^2}-(\hat{r}+r^*)\]
Then, when \(\frac{M(\mathbf{z}_0)}{\Delta^2}\geq \sqrt{\alpha_1}\) and \[\epsilon \geq \exp\left(-\frac{1}{8}\left[(\Delta^2 - 2M(\mathbf{z}_0)+(\hat{r}+r^*))^2 + (\hat{r}-r^*)^2\right]\right)\] then \[\begin{align} \int_{a}^b \frac{\tau}{\sqrt{-\ln \tau}}\mathrm{d}\tau&\geq \frac{\frac{M(\mathbf{z}_0)}{\Delta}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} + R_0}+\frac{M(\mathbf{z}_0)^2}{2\Delta^2} + \frac{1}{2}R_0}{\Delta^2\sqrt{\ln(\Delta)- \ln\left(\frac{M(\mathbf{z}_0)}{\Delta}+ \frac{1}{2}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} + R_0}\right)}} \end{align}\]
Remark 7. This proposition requires \(\lambda_t=1\), as the analysis is far more challenging when the integrand \(g\) is nonconvex. If one wishes to analyze alternative weighting factors or, indeed, an alternative noise schedule, all the analysis up to this point applies. The difference is that one now needs to lower bound the more challenging expression \(\int_{t_{b^2}}^{t_{a^2}}\alpha_t\lambda_t\mathrm{d}t\). For the linear noise schedule and \(\lambda_t=1-\alpha_t\), this take the form \(\frac{1}{\sqrt{\beta_1}}\int_{a}^b \frac{\tau(1-\tau^2)}{\sqrt{-\ln \tau}}\mathrm{d}\tau\).
Proof. Let \[g(t):= \frac{\tau}{\sqrt{-\ln \tau}}\]
Note that \[\begin{align} g''(t)&= \frac{1}{2t(-\ln t)^{3/2}} + \frac{3}{4t(-\ln t)^{5/2}} >0 \end{align}\]
So, \(g(t)\) is convex, meaning \[\int_a^b g(\tau)\mathrm{d}\tau \geq (b-a)g\left(\frac{a+b}{2}\right)\]
The condition that \(\frac{M(\mathbf{z}_0)}{\Delta^2}\geq \sqrt{\alpha_1}\) guarantees \(a\geq\sqrt{\alpha_1}\). Similarly, the condition on \(\epsilon\) guarantees \(b\leq 1\), so \[\begin{align} b-a&=\frac{1}{\Delta}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} - \left( \hat{r} + r^*\right)+\sqrt{8\log \frac{1}{\epsilon}-\left(\hat{r}-r^*\right)^2}} \\ &= \frac{1}{\Delta}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} + R_0} \end{align}\]
We have \[\begin{align} \frac{a+b}{2}&= \frac{M(\mathbf{z}_0)}{\Delta^2}+\frac{1}{2\Delta}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} + R_0} \end{align}\] and so \[\begin{align} \Longrightarrow g\left(\frac{a+b}{2}\right)&= \frac{\frac{M(\mathbf{z}_0)}{\Delta}+\frac{1}{2}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} + R_0}}{\Delta\sqrt{\ln(\Delta)- \ln\left(\frac{M(\mathbf{z}_0)}{\Delta}+ \frac{1}{2}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} + R_0}\right)}} \end{align}\]
Combining these, \[\begin{align} \int_a^b g(\tau)\mathrm{d}\tau &\geq \frac{\frac{M(\mathbf{z}_0)}{\Delta}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} + R_0}+\frac{M(\mathbf{z}_0)^2}{2\Delta^2} + \frac{1}{2}R_0}{\Delta^2\sqrt{\ln(\Delta)- \ln\left(\frac{M(\mathbf{z}_0)}{\Delta}+ \frac{1}{2}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} + R_0}\right)}} \end{align}\] ◻
Combining this lower bound with what we showed before, we have \[\begin{align} &\int_0^1\norm{\nabla \log p_{t}^{(\hat{\gamma})}(\mathbf{z}_t) - \nabla \log p_{t}^{(\gamma^{*})}(\mathbf{z}_t)}^2 \mathrm{d}t \\ &\quad\quad\geq\frac{(\hat{\gamma}-\gamma^*)^2 \epsilon}{16\sqrt{\beta_1}\;\hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)} \frac{\frac{M(\mathbf{z}_0)}{\Delta}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} + R_0}+\frac{M(\mathbf{z}_0)^2}{2\Delta^2} + \frac{1}{2}R_0}{\sqrt{\ln(\Delta)- \ln\left(\frac{M(\mathbf{z}_0)}{\Delta}+ \frac{1}{2}\sqrt{\frac{M(\mathbf{z}_0)^2}{\Delta^2} + R_0}\right)}} \end{align}\]
Recall that \(\mathbf{z}_0\sim \mathcal{N}(\mu_0, \mathbf{I})\), so \[\begin{align} M(\mathbf{z}_0) &= \left\langle\mathbf{z}_0 -\mu_0, \mu_1-\mu_0\right\rangle \sim \mathcal{N}(0, \norm{\mu_1-\mu_0}^2) \\ \Longrightarrow \frac{M(\mathbf{z}_0)}{\Delta}&\sim \mathcal{N}(0,1) \end{align}\]
Proposition 7. Let \(\Phi\) be the CDF of \(\mathcal{N}(0,1)\). Let \(\tau\geq \Delta\sqrt{\alpha_1}\), and let \(\epsilon>0\) satisfy \(\epsilon \in (\epsilon_{\mathrm{min}}^0, \epsilon_{\mathrm{max}}^0)\) with \[\begin{align} \epsilon_{\mathrm{max}} ^0& =\exp\left(-\frac{1}{8}\left[(\hat{r}-r^*)^2 + \left(\hat{r}+r^*-\tau^2\right)_+^2\right]\right) \\ \epsilon_{\mathrm{min}}^0 &= \exp\left(-\frac{1}{8}\left[(\Delta(\Delta- 2\tau)+(\hat{r}+r^*))^2 + (\hat{r}-r^*)^2\right]\right) \end{align}\]
Then, with probability \(1-\Phi(\tau)\) over samples of \(\mathbf{z}_0\), \[\begin{align} &\int_0^1\norm{\nabla \log p_{t}^{(\hat{\gamma})}(\mathbf{z}_t) - \nabla \log p_{t}^{(\gamma^{*})}(\mathbf{z}_t)}^2 \mathrm{d}t \\ &\quad\quad\geq\frac{(\hat{\gamma}-\gamma^*)^2 \epsilon}{16\sqrt{\beta_1}\; \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)} \frac{\tau\sqrt{\tau^2 + R_0}+\frac{1}{2}\tau^2 + \frac{1}{2}R_0}{\sqrt{\ln(\Delta)- \ln\left(\tau+ \frac{1}{2}\sqrt{\tau^2 + R_0}\right)}} \end{align}\]
Proof. If we let \(K:= \frac{M(\mathbf{z}_0)}{\Delta}\) our bound becomes \[\begin{align} &\int_0^1\norm{\nabla \log p_{t}^{(\hat{\gamma})}(\mathbf{z}_t) - \nabla \log p_{t}^{(\gamma^{*})}(\mathbf{z}_t)}^2 \mathrm{d}t \\ &\quad\quad\geq\frac{(\hat{\gamma}-\gamma^*)^2 \epsilon}{16 \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)} \frac{K\sqrt{K^2 + R_0}+\frac{1}{2}K^2 + \frac{1}{2}R_0}{\sqrt{\ln(\Delta)- \ln\left(K+ \frac{1}{2}\sqrt{K^2 + R_0}\right)}} \end{align}\]
Assume that \(K\geq \tau\), which occurs with probability \(1-\Phi(\tau)\). Our bound is an increasing function in \(K\) and is thus minimized at \(K=\tau\). Similarly, our upper bound from before on \(\epsilon\) is a decreasing function of \(K\) and our lower bound from before is an increasing function of \(\epsilon\), so our bounds on \(\epsilon\) are tightest when we plug in \(\tau\), giving \(\epsilon_{\mathrm{max}}^0\) and \(\epsilon_{\mathrm{min}}^0\) from the statement of the proposition. ◻
We now state the corresponding proposition for the case where \(c=1\). Let \(\mathbf{w}_t := \mathbf{y}_t |c=1\).
Proposition 8. Let \(\tau\geq \Delta\sqrt{\alpha_1}\), and let \(\epsilon>0\) satisfy \(\epsilon \in (\epsilon_{\mathrm{min}}^1, \epsilon_{\mathrm{max}}^1)\) with \[\begin{align} \epsilon_{\mathrm{max}} ^1& =\exp\left(-\frac{1}{8}\left[(\hat{r}-r^*)^2 + \left(-(\hat{r}+r^*)-\tau^2\right)_+^2\right]\right) \\ \epsilon_{\mathrm{min}}^1 &= \exp\left(-\frac{1}{8}\left[(\Delta(\Delta- 2\tau)+(\hat{r}+r^*))^2 - (\hat{r}-r^*)^2\right]\right) \end{align}\]
Then, with probability \(1-\Phi(\tau)\) over samples of \(\mathbf{w}_0\), \[\begin{align} &\int_0^1\norm{\nabla \log p_{t}^{(\hat{\gamma})}(\mathbf{w}_t) - \nabla \log p_{t}^{(\gamma^{*})}(\mathbf{w}_t)}^2 \mathrm{d}t \\ &\quad\quad\geq\frac{(\hat{\gamma}-\gamma^*)^2 \epsilon}{16\sqrt{\beta_1}\; \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)} \frac{\tau\sqrt{\tau^2 + R_1}+\frac{1}{2}\tau^2 + \frac{1}{2}R_1}{\sqrt{\ln(\Delta)- \ln\left(\tau+ \frac{1}{2}\sqrt{\tau^2 + R_1}\right)}} \end{align}\] where \[R_1:=\sqrt{8\ln \frac{1}{\epsilon}-(\hat{r}-r^*)^2}+(\hat{r}+r^*)\]
Proof. This follows from repeating all of our analysis above for the \(c=0\) case with \(\gamma^*\) and \(\hat{\gamma}\) replaced with \(1-\gamma^*\) and \(1-\hat{\gamma}\), respectively. ◻
Combining 7 and 8, we have by our decomposition of \(\ell_\mathrm{DSM}\) from the start of this section that \[\ell_\mathrm{DSM}(p_{}^{(\hat{\gamma})},p_{}^{(\gamma^*)})\geq C (\hat{\gamma} - \gamma^*)^2\] where \[C=\frac{ \epsilon ( 1-\Phi(\tau))}{16\sqrt{\beta_1}\; \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)}\left[(1-\gamma^*)\frac{\tau\sqrt{\tau^2 + R_0}+\frac{1}{2}\tau^2 + \frac{1}{2}R_0}{\sqrt{\ln(\Delta)- \ln\left(\tau+ \frac{1}{2}\sqrt{\tau^2 + R_0}\right)}}+\gamma^* \frac{\tau\sqrt{\tau^2 + R_1}+\frac{1}{2}\tau^2 + \frac{1}{2}R_1}{\sqrt{\ln(\Delta)- \ln\left(\tau+ \frac{1}{2}\sqrt{\tau^2 + R_1}\right)}}\right]\] for all \(\tau\geq \Delta\sqrt{\alpha_1}\) and \(\epsilon\in [ \epsilon_{\mathrm{min}}, \epsilon_{\mathrm{\max}}]\), where \[\begin{align} \epsilon_{\mathrm{max}}&=\min\set{\epsilon_{\mathrm{max}}^0,\epsilon_{\mathrm{max}}^1}=\exp\left(-\frac{1}{8}\left[(\hat{r}-r^*)^2 + \max\set{\left(\hat{r}+r^*-\tau^2\right)_+^2,\left(-(\hat{r}+r^*)-\tau^2\right)_+^2}\right]\right) \\ \epsilon_{\mathrm{min}}&=\max\set{\epsilon_{\mathrm{min}}^0,\epsilon_{\mathrm{min}}^1}=\exp\left(-\frac{1}{8}\left[\min\set{(\Delta(\Delta- 2\tau)+(\hat{r}+r^*))^2,(\Delta(\Delta- 2\tau)-(\hat{r}+r^*))^2} + (\hat{r}-r^*)^2\right]\right) \end{align}\]
To simplify this expression, note that if \[g(x):=\frac{x(\tau+\frac{1}{2}x)}{\sqrt{\ln(\Delta) - \ln(\tau + \frac{1}{2}x)}}\] then direct calculation shows that if \(h(x):=\ln(\Delta) - \ln(\tau + \frac{1}{2}x)\), \[g'(x)=\frac{\frac{x}{4}+(\tau+x)h(x)}{h(x)^{3/2}}>0\]
Therefore, \(g\) is increasing, meaning \[\begin{align} C&=\frac{ \epsilon ( 1-\Phi(\tau))}{16\sqrt{\beta_1}\; \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)}\left[(1-\gamma^*)g\left(\sqrt{\tau^2+R_0}\right)+\gamma^* g\left(\sqrt{\tau^2+R_1}\right)\right] \\ &\geq \frac{ \epsilon ( 1-\Phi(\tau))}{16\sqrt{\beta_1}\; \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)}g\left(\min\set{\sqrt{\tau^2+R_0}, \sqrt{\tau^2+R_1}}\right) \\ &= \frac{ \epsilon ( 1-\Phi(\tau))}{16\sqrt{\beta_1}\; \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)}g\left(\sqrt{\tau^2+\sqrt{8\ln \frac{1}{\epsilon}-(\hat{r} - r^*)^2}- \abs{\hat{r}-r^*}}\right) \end{align}\]
Let \(R_{\mathrm{min}}:=\sqrt{8\ln \frac{1}{\epsilon}-(\hat{r} - r^*)^2}- \abs{\hat{r}-r^*}\). To obtain the bound from the body of the paper, we now specify values for \(\tau\) and \(\epsilon\) which simplify the expression. First, we will set \(\tau\) to be its minimum value: \(\tau=\Delta\sqrt{\alpha_1}\). We choose this because numerically, it seems to be the value of \(\tau\) that maximizes the expression. This gives us \[C\geq \frac{ \epsilon ( 1-\Phi(\Delta\sqrt{\alpha_1}))}{16\sqrt{\beta_1}\; \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)} \frac{\sqrt{\Delta^2\alpha_1 + R_{\mathrm{min}} }\left(\Delta\sqrt{\alpha_1} +\frac{1}{2}\sqrt{\Delta^2\alpha_1+ R_{\mathrm{min}}}\right)}{\sqrt{\log(\Delta)-\log\left(\sqrt{\alpha_1}\right)-\log\left(\Delta+\frac{1}{2}\sqrt{\Delta^2+\frac{R_0}{\alpha_1}}\right)}}\]
Noting that \(\log(\Delta)<\log\left(\Delta+\frac{1}{2}\sqrt{\Delta^2+\frac{R_0}{\alpha_1}}\right)\), the denominator of the second term is less than \(\sqrt{-\log(\sqrt{\alpha_1})}=\frac{1}{4}\beta_1\) (recall \(\alpha_t=\exp\left(-\int_0^t \beta_s\mathrm{d}s\right)=\exp\left(-\frac{1}{2}\beta_1 t^2\right)\)). Additionally, note that the numerator of the second term is an increasing function of \(\Delta\), meaning we can lower-bound it by setting \(\Delta=0\). Combining these gives us \[C\geq \frac{1-\Phi(\Delta\sqrt{\alpha_1})}{8\beta_1\; \hat{\gamma} (1-\hat{\gamma}) \gamma^* (1-\gamma^*)} \epsilon \,R_{\mathrm{min}}\]
Finally, direct calculation shows that \(\epsilon\, R_{\mathrm{min}}\) is maximized when \(\epsilon=\exp\left(-\frac{1}{8}\left((\hat{r}-r^*)^2 + \delta^2\right)\right)\) where \(\delta = \frac{|\hat{r}-r^*|+ \sqrt{(\hat{r}-r^*)^2+16}}{2}\), giving \[\epsilon \, R_{\mathrm{min}} =\frac{\sqrt{(\hat{r}-r^*)^2+16}-|\hat{r}-r^*|}{2} \exp\left(-\frac{1}{8}\left((\hat{r}-r^*)^2 + \delta^2\right)\right)\]
Note \(\delta\leq \abs{\hat{r}-r^*}+2\) (because \(\sqrt{a^2+b^2}\leq |a|+|b|\)), and \(\delta^2\leq \abs{\hat{r}-r^*}^2+2\abs{\hat{r}-r^*}+4\). Applying this, and the fact that \(\hat{\gamma}(1-\hat{\gamma})\leq\frac{1}{4}\) \[\begin{align} C&\geq \frac{\left(1-\Phi(\Delta\sqrt{\alpha_1})\right)\left(\sqrt{(\hat{r}-r^*)^2+16}-|\hat{r}-r^*|\right)}{4\beta_1\; \gamma^* (1-\gamma^*)}\exp\left(-\frac{1}{4}\left((\hat{r}-r^*)^2 +\abs{\hat{r}-r^*}\right) - \frac{1}{2}\right) \\ &=\frac{\left(1-\Phi(\Delta\sqrt{\alpha_1})\right)\left(\sqrt{(\hat{r}-r^*)^2+16}-|\hat{r}-r^*|\right)}{4\sqrt{e}\beta_1\; \gamma^* (1-\gamma^*)}\exp\left(-\frac{1}{4}\left((\hat{r}-r^*)^2 +\abs{\hat{r}-r^*}\right)\right) \end{align}\]
Finally, for \(\hat{\gamma}\in [a,1-a]\), \(\abs{\hat{r} - r^*} \leq \log\frac{1-a}{a}+\abs{\log \frac{1-\gamma^*}{\gamma^*}}\). Our expression is a decreasing function of \(\abs{\hat{r}-r^*}\), and thus for all \(\hat{\gamma}\in [a,1-a]\), if we let \(k:=\log\frac{1-a}{a}+\abs{\log \frac{1-\gamma^*}{\gamma^*}}\) \[C\geq \frac{\left(1-\Phi(\Delta\sqrt{\alpha_1})\right)\left(\sqrt{k^2+16}-k\right)}{7\beta_1\; \gamma^* (1-\gamma^*)}\exp\left(-\left(k^2 +k\right)/4\right)\] where we additionally used \(4\sqrt{e}<7\). Bringing this all together, \[\begin{align} \min_{\hat{\gamma} \in [a,1-a]}\frac{\ell_\mathrm{DSM}(p_{}^{(\hat{\gamma})},p_{}^{(\gamma^*)})}{(\hat{\gamma}-\gamma^*)^2}>\frac{(1-\Phi(\norm{\mu_1-\mu_0}\sqrt{\alpha_1}))\left(\sqrt{k^2+16}-k\right)}{7\beta_1\gamma^*(1-\gamma^*)}\exp\left(-(k^2+k)/4\right) \end{align}\] 0◻
11 Additional numerical experiments↩︎
11.1 Noise schedules and time reparameterization↩︎
In this work, we reference two noise schedules used in practice. The linear noise schedule is defined by \(\beta_t:=\beta_1 t\) for a constant \(\beta_1>0\), such that \(\alpha_t=\exp\left(-\frac{\beta_1}{2}t^2\right)\). The squared-cosine noise schedule [58]
\(\bar\alpha_t=\cos\left(\frac{t+s}{1+s}\cdot \frac{\pi}{2}\right)^2\) and \(\alpha_t =\frac{\bar\alpha_t}{\bar\alpha_0}\) for \(s=0.008\).
We want to analyze the effect of artificially reducing the sensitivity index on the estimation of \(\gamma^*\). To do this, we can reparameterize the time in the forward process as \(t=t_{a,b,v}(\tau)\), where \[t_{a,b,v}(\tau)=\begin{cases} \tau & 0\leq \tau<a \\ a + v(\tau-a)& a \leq \tau <a + \frac{b-a}{v} \\ b + \frac{1-b}{1-\left(a + \frac{b-a}{v}\right)}\left(t - a - \frac{b-a}{v}\right) & a + \frac{b-a}{v} \leq t\leq 1 \end{cases}\]
Under this parameterization, if we let \(\mathbf{y}_\tau:= \mathbf{x}_{t(\tau)}\) where \(\mathbf{x}_t\) follows the Variance Preserving SDE with noise schedule \(\set{\alpha_t}\), then \(\mathbf{y}_t\) satisfies \[\begin{align} \mathrm{d}\mathbf{y}_\tau &= -\frac{\beta_{t(\tau)}}{2}\mathbf{y}_\tau \frac{\mathrm{d}t}{\mathrm{d} \tau}\mathrm{d}\tau + \sqrt{\beta_{t(\tau)}}\cdot \sqrt{\frac{\mathrm{d}t}{\mathrm{d}\tau}}\mathrm dW_t \\ &= -\frac{\tilde{\beta}_{\tau}}{2}\mathbf{y}_\tau \mathrm{d}\tau + \sqrt{\tilde{\beta}_{\tau}}\mathrm{d}W_t \end{align}\] where \(\tilde{\beta}_\tau:= \beta_{t(\tau)} \cdot \frac{\mathrm{d}t}{\mathrm{d}\tau}\). In particular, reparameterizing time is equivalent to choosing a different noise schedule \(\set{\tilde{\alpha}_\tau}\). For \(v\geq 1\), \(t_{a,b,v}\) maps the interval \([a,b]\) to \(\left[a, a+\frac{b-a}{v}\right]\), scaling the length of the interval by \(1/v\). To reduce the value of \(L(\gamma^*)\), we let \([a,b]\) contain the region of time where \(\frac{\ell_\mathrm{SM}(p_{t}^{(\gamma)}, p_{t}^{(\gamma^*)})}{(\gamma-\gamma^*)^2}>\epsilon\) for fixed \(\epsilon>0\) and increase the value of \(v\).
11.2 Neural network experiment details↩︎
For the score models \(\mathbf{s}_{\boldsymbol{\theta}}\) in our experiments, we use deep ResNets [50] with width equal to the ambient dimension and depth \(8\). Unless otherwise specified, we use a linear noise schedule with \(\beta_1=20\) and a time discretization of \(T=4000\) steps on the interval \([10^{-5},1]\) for training and \([10^{-3},1]\) for sampling, following [23].
We train our models using the loss described in this paper (and [23]) using AdaBelief [59] with a linear warmup and cosine decay learning rate. We use JAX [60] and Equinox [61] for our neural network code, SciPy [62] for additional machine learning algorithms, Matplotlib [63] for figure generation, and Weights and Biases [64] for experiment tracking. 5
11.3 MNIST ablation study details↩︎
The architecture of our autoencoder is simple: two convolutional layers followed by a dense layer for the encoder, and the reverse for the decoder. For our MNIST experiments, we use a latent dimension of \(12\). We train the autoencoder on all examples of \(1\)s and \(8\)s in MNIST. We additionally train an MNIST classifier on only \(1\)s and \(8\)s using the architecture in the Equinox documentation [61].
Unlike in the case of Gaussian mixture models, we do not know the analytical score function for the latents of our autoencoder, and thus estimation of \(\ell_\mathrm{DSM}(p_{t}^{(\gamma)}, p_{t}^{(\gamma^*)})\) and \(L(\gamma^*)\) analytically is intractable. Instead, for \(\gamma^*\in \set{0.2,0.22,...,0.8}\), we train diffusion models \(\mathbf{s}_{{\boldsymbol{\theta}}_\gamma}\) on data from \(p_{}^{(\gamma)}\) as approximations for the analytical scores using a linear noise schedule. We keep the number of training examples constant across all experiments (5800, the smaller of the counts of 1s and 8s in MNIST) and train the models identically. For each \(t\in [0,1]\), we then have the estimate \[\ell_\mathrm{DSM}(p_{t}^{(\gamma)}, p_{t}^{(\gamma^*)})\approx \expect[\mathbf{x}_t\sim p_{t}^{(\gamma^{*})}]{\lambda_t \norm{\mathbf{s}_{{\boldsymbol{\theta}}_{\gamma^*}}(\mathbf{x}_t, t) - \mathbf{s}_{{\boldsymbol{\theta}}_{\gamma'}}(\mathbf{x}_t, t)}^2}\]
We generate all samples used in 4 using \(\mathbf{s}_{{\boldsymbol{\theta}}_{0.6}}\), varying the noise schedule used for sampling.
11.4 DSSI for other parameters↩︎
Here, we apply our DSSI framework to parameters other than the mixture weight in a Gaussian model. In what follows, we will let \(\Sigma_0\) and \(\Sigma_1\) be fixed random positive definite matrices, such that \[\Sigma_i = A^\top A+10^{-8}I, \;\;A_{ij}\sim \mathcal{N}(0,1)\] for \(i\in \set{0,1}\) and use a linear noise schedule. We focus here on one-dimensional families of distributions for consistency with the body of the paper, although we note that our framework can be extended to multi-parameter recovery directly. For this section, we examine GMMs in dimension \(d=10\).
For our analysis of the sensitivity of \(\ell_\mathrm{DSM}\) to the modes of the mixture components, we consider the one-dimensional family of distributions \[p_{}^{(\delta)}:=0.6 \; \mathcal{N}(\mathbf{0}, \mathbf{\Sigma}_0) + 0.4 \; \mathcal{N}\left(\frac{12\;\delta}{\sqrt{d}}\mathbf{1},\mathbf{\Sigma}_1\right)\] in dimension \(d=10\) where we range \(\delta\) between \(0\) and \(1\). 6 (a) shows \(\ell_\mathrm{SM}(p_{t}^{(\delta)}, p_{t}^{(\delta^*)})\) as a function of \(\delta\) and \(t\), and 6 (b) shows \(\frac{\ell_\mathrm{DSM}(p_{t}^{(\delta)}, p_{t}^{(\delta^*)})}{(\delta - \delta^*)^2}\) as a function of \(\delta\) for multiple values of \(\delta^*\). We see that, as with the mixture weight, \(\ell_\mathrm{SM}(p_{t}^{(\delta)}, p_{t}^{(\delta^*)})\) is nonnegligible only for intermediate times in the forward process. However, note that the minimum value of \(L(\delta^*)\) we see here is greater than \(2\) (compared to \(0.37\) in 2 (b)), meaning that our theory implies a \(5.4\times\) stronger control on the error in estimation \(\mu\) using a diffusion model trained to the same level of error under \(\ell_\mathrm{DSM}\).
Figure 6: Sensitivity of \(\ell_\mathrm{DSM}\) to means of mixture components. (a) \(\ell_\mathrm{SM}(p_{t}^{(\delta)}, p_{t}^{(\delta^*)})\) as a function of \(t\) and \(\delta\). (b) \(\frac{\ell_\mathrm{DSM}(p_{}^{(\delta)}, p_{}^{(\delta^*)})}{(\delta - \delta^*)^2}\) as a function of \(\delta\)..
We now apply our DSSI framework to covariance matrix recovery. Consider the one-dimensional family of distributions \[p_{}^{(\kappa)}:=0.6 \; \mathcal{N}(\mathbf{0}, \mathbf{I}) + 0.4 \; \mathcal{N}\left(\frac{9.0}{\sqrt{d}}\mathbf{1},\mathbf{\Sigma}_\kappa\right)\] in dimension \(d=10\) where we range \(\delta\) between \(0\) and \(1\). Here, we define \[\mathbf{\Sigma}_\kappa:= \mathbf{\Sigma}_0^{1/2}\left(\mathbf{\Sigma}_0^{-1/2}\mathbf{\Sigma}_1\mathbf{\Sigma}_0^{-1/2}\right)^{\kappa}\mathbf{\Sigma}_0^{1/2}\] to be the geodesic interpolating between \(\mathbf{\Sigma}_0\) and \(\mathbf{\Sigma}_1\). 7 (a) shows \(\ell_\mathrm{SM}(p_{t}^{(\kappa)}, p_{t}^{(\kappa^*)})\) as a function of \(\kappa\) and \(t\), and 7 (b) shows \(\frac{\ell_\mathrm{DSM}(p_{}^{(\kappa)}, p_{}^{(\kappa^*)})}{(\kappa - \kappa^*)^2}\) as a function of \(\kappa\).
Again, \(\ell_\mathrm{SM}(p_{t}^{(\delta)}, p_{t}^{(\delta^*)})\) is nonnegligable for intermediate times in the forward process. We additionally see that \(L(\kappa^*)\) is greater than \(0.6\) uniformly across \(\kappa\).
Figure 7: Sensitivity of \(\ell_\mathrm{DSM}\) to covariance matrix of mixture components. (a) \(\ell_\mathrm{SM}(p_{t}^{(\kappa)}, p_{t}^{(\kappa^*)})\) as a function of \(t\) and \(\kappa\). (b) \(\frac{\ell_\mathrm{DSM}(p_{}^{(\kappa)}, p_{}^{(\kappa^*)})}{(\kappa - \kappa^*)^2}\) as a function of \(\kappa\)..
References↩︎
As a slight abuse of terminology, the sequences \(\set{\alpha_t}\) and \(\set{\beta_t}\) will interchangably be referred to as the “noise schedule”.↩︎
The superscript here indicates the trainable parameters \({\boldsymbol{\theta}}\) of the learned network.↩︎
The mixture model \(q\mathrel{\vcenter{:}}= \sum_{i=1}^k \frac{1}{k}\, \mathcal{N}(\boldsymbol{\mu}^{(i)},\, \mathbf{I})\) where each component has equal mixture weight and identity covariance↩︎
We will make our code available upon publication.↩︎