1 Introduction↩︎

Score-based diffusion models [1], [2] have achieved state-of-the-art performance across a wide range of generative modeling applications, including image and video generation ([3]; [4]), signal processing [5], and language modeling [6], [7]. At a high level, diffusion models generate samples by starting from Gaussian noise and iteratively denoising via a learned reverse-time diffusion dynamics. This procedure relies on accurate estimation of the score function (the gradient of the log-density) along a forward noising process.

Diffusion models can be viewed as unsupervised distribution learners—given finite training samples from an unknown data distribution, they aim to generate new samples that faithfully follow the same law. This perspective raises a fundamental statistical question: how many training samples are needed for diffusion-based sampling to accurately learn the underlying data distribution, and can this sample complexity match the information-theoretic limit? At a conceptual level, diffusion sampling is inherently two-stage: it first uses training data to estimate the time-indexed scores along a forward diffusion process, and then plugs these learned scores into an iterative sampling procedure to generate an output. Therefore, addressing the above question calls for sample complexity guarantees that jointly control both the score estimation error and the error accumulated during sampling.

1.0.0.1 Leveraging intrinsic structures.

Motivated by this, a growing body of statistical theory has been developed to understand the sample complexity of diffusion models [8]–[13].

For a broad class of \(d\)-dimensional distributions with \(\beta\)-Hölder smooth densities (without assuming smooth scores or log-concave/uniformly-bounded densities), state-of-the-art theory shows that both DDPM [14] and DDIM [15] require on the order of (up to logarithmic factors) \[\begin{align} \label{eq:sample32complexity32general} \varepsilon^{-\frac{d+2\beta}{\beta}} \end{align}\tag{1}\] training samples to generate an output within \(\varepsilon\) total variation (TV) distance to the target distribution. While this sample complexity is (nearly-)minimax optimal for general smooth-density classes, it suffers from the curse of dimensionality as the ambient dimension \(d\) grows. Consequently, such guarantees fail to fully explain the empirical effectiveness of diffusion models in modern high-dimensional applications, suggesting that the worst-case bounds in 1 may be overly pessimistic for structured distributions arising in practice.

To narrow this gap, recent work has explored whether, and in what sense, diffusion models can exploit intrinsic structures underlying the data distributions. However, statistical theory of diffusion models for structured data distributions remains far from complete. Existing results in this direction typically focus on distributions supported on a single low-dimensional structure, such as a linear subspace or manifold [16]–[19], a factor model [20], or certain dependence structures [21]. Although these works establish improved sample complexities governed by intrinsic rather than ambient dimension, they require strong assumptions on the data distribution. A prominent example is the requirement for the density to be uniformly bounded away from zero on its support. While this condition is standard in the nonparametric statistics literature and technically convenient, it excludes important multi-modal structures with well-separated components, where the density necessarily becomes small, or even vanishes, between modes.

More fundamentally, the prevailing “single manifold/subspace” paradigm limits our theoretical understanding of diffusion models’ capabilities to learn heterogeneous distributions whose different modes concentrate near different low-dimensional structures. Such geometry is common in modern high-dimensional data, where distinct classes or clusters may occupy separate subspaces or manifolds [22], [23]. As a result, existing theory still falls short of explaining the empirical effectiveness of diffusion models when learning low-dimensional, multi-modal distributions.

1.1 Main contributions↩︎

In this paper, we develop a statistical theory for diffusion models, aimed at understanding how many samples are required to learn low-dimensional, multi-modal distributions.

Concretely, we consider a target data distribution \(p^\star\) supported on a union of subspaces (UoS), i.e., \[\begin{align} \mathsf{supp}(p^\star) \subseteq \cup _{i=1}^M V_i, \end{align}\] where each \(V_i\) is a linear subspace with dimension \(k_i\). We denote by \(k \mathrel{\vcenter{:}}=\max_{i\in[M]} k_i\) the maximum intrinsic dimension. In addition, we assume that the restriction of the target distribution to each subspace is \(\sigma\)-subgaussian.

Under these two assumptions, we construct a kernel-based regularized score estimator \(\widehat s_t\) for the score function of the Gaussian-smoothed distribution \(p_t \mathrel{\vcenter{:}}= p^\star \ast \mathcal{N}(0,tI_d)\) for any \(t>0\). Given \(n\) samples drawn from the target distribution \(p^\star\), we establish a finite-sample \(L^2\) score estimation error bound: \[\begin{align} \mathop{\mathbb{E}}\big[\|\widehat{s}_t(X)-s_t^\star(X)\|_2^2\big] = \widetilde{O}\bigg(\frac{1}{n} \Big( \frac{1}{t}+\frac{1}{t^{{(k\vee2)}/2+1}}\Big) \bigg). \end{align}\] Here the expectation is taken over both the training data and \(X\sim p_t\), where we write \(a\vee b \mathrel{\vcenter{:}}=\max\{a,b\}\).

Building on this score estimation guarantee, we prove that diffusion samplers require at most the order of (up to logarithmic factors) \[\begin{align} \varepsilon^{- (k \vee 2)} \end{align}\] training samples to generate a sample that is \(\varepsilon\)-close in 1-Wasserstein distance to the target distribution \(p^\star\). Importantly, this convergence rate depends only on the intrinsic dimension \(k\), rather than the ambient dimension \(d\), and matches the minimax optimal rate for learning a \(k\)-dimensional distribution [24]. Moreover, our theory requires only subgaussian tails on each subspace, without imposing any restrictive assumptions on scores or densities. As a result, our framework naturally accommodates multi-modal distributions with well-separated components.

Finally, we emphasize that the kernel-based score estimator developed in this paper is primarily a theoretical proof device, rather than a practical alternative to neural network (NN)-based score estimation. Nevertheless, our results provide an important step toward statistical guarantees for NN score-based diffusion models. In particular, they establish the achievability of the fundamental statistical limit and identify the structural properties that analysis of NN score estimators should capture, while also providing an explicit low-dimensional target for NN approximation. More discussion on extensions to NN-based scores can be found in Section 6.

1.2 Related works↩︎

Statistical theory for diffusion models.

Recent work has begun to provide finite-sample guarantees for diffusion-based sampling by studying statistical bounds for score estimation error. Under strong density regularity assumptions (e.g., boundedness on compact domains), [9] showed that neural-network-based ERM score estimators lead to minimax-optimal rates in both TV and \(W_1\) distances when used with reverse SDE samplers. Using nonparametric constructions under density lower-bound conditions, [13] derived minimax-optimal score estimation rates and corresponding sampling guarantees. More recently, for subgaussian targets with \(\beta\)-Hölder smooth densities, [14] established minimax-optimality for DDPM using truncated kernel score estimators, and [15] obtained an end-to-end minimax-optimal convergence analysis for ODE-based diffusion (DDIM/probability flow) by combining smoothed score estimation with convergence analysis of the sampling dynamics.

Adaptation to low-dimensional structures.

For distributions supported on a linear subspace, [16] established convergence rates governed by the subspace dimension under smooth score assumptions. For manifold-supported targets, [18] proved analogous intrinsic-dimension rates, but the analysis requires controlling geometric approximation error (e.g., via Hausdorff distance) and typically relies on density lower-bound conditions on the support. Beyond geometric support constraints, [21] obtained minimax-optimal rates for diffusion learning under structured dependence (exponential-interaction) models. In addition, [25] analyzed mixtures of low-rank Gaussians, focusing on the special case of orthogonal subspaces. [26] showed that shallow NN-based diffusion models can provably adapt to hidden low-dimensional subspace structure under independent component data models and smoothness assumptions on the latent scores.

Complementary to the statistical perspective, a parallel line of work studies whether the sampling stage of diffusion models can automatically exploit low-dimensional data structure [27]–[30]. These works show that the iteration complexity, the number of sampling iterations required to achieve a desired accuracy, also depends only on the intrinsic dimension rather than the ambient dimension. In addition, low-dimensional adaptation has also been investigated for discrete diffusion models when learning discrete distributions [31]–[35].

1.3 Notation↩︎

For \(a,b\in\mathbb{R}\), we denote \(a \vee b \mathrel{\vcenter{:}}=\max\{a,b\}\) and \(a\wedge b \mathrel{\vcenter{:}}=\min\{a,b\}\). For positive integer \(M\), let \([M]\mathrel{\vcenter{:}}=\{1,\cdots,M\}\). For random vector \(X\), we use \(p_X\) to denote its distribution or probability density function, depending on the context. For any vector \(x\in \mathbb{R}^d\), we denote \(\| \cdot \|_p\) as its \(p\)-norm, i.e., \(\|x\|_p \mathrel{\vcenter{:}}=(\sum_{i=1}^d |x_i|^p)^{1/p}\), and write \(\|x\|_{\infty} \mathrel{\vcenter{:}}=\max_{i}|x_i|\). We use \(\|\cdot\|\) to denote the 2-norm for simplicity. For any vector \(x\in \mathbb{R}^d\) and any \(i,j\in[d]\) with \(i<j\), we denote by \(x_{i:j} \in \mathbb{R}^{j-i+1}\) the subvector consisting of the \(i\)-th through \(j\)-th entries of \(x\). For vectors \(\{\alpha_i\}_{i=1}^k\), we denote by \(\mathsf{span}(\{ \alpha_i \}_{i=1}^k)\) the linear space spanned by these vectors.

For a probability distribution \(p\) and a random vector \(X\), we write \(X \sim p\) to mean that \(X\) follows the distribution \(p\). We denote by \(\mathsf{supp}(p)\) the support of probability measure \(p\), i.e., the smallest closed set \(S\) such that \(p(S)=1\). In addition, let \(\mathbb{1}\{ \cdot \}\) denote the indicator function. For random vectors \(X, Y\), we define the 1-Wasserstein distance between their distributions \(p_X\) and \(p_Y\) by \[\begin{align} W_1(p_X,p_Y) \mathrel{\vcenter{:}}=\inf_{\gamma \in \Gamma(p_X,p_Y)}\iint \|x-y\| \, \gamma(\!\,\mathrm{d}x,\!\,\mathrm{d}y), \end{align}\] where \(\Gamma(p_{X},p_{Y})\) denotes the set of couplings of \(p_X\) and \(p_Y\). For probability distributions \(P,Q\), we denote their convolution by \(P\ast Q\). Finally, we use \(\mathsf{poly}(n)\) to denote a polynomial function of \(n\) where the specific degree may vary across different contexts.

2 Problem formulation↩︎

2.1 Preliminaries↩︎

In this section, we briefly introduce the score-based diffusion models.

Forward process.

The forward process starts from the target distribution \(p^\star\) and gradually adds Gaussian noise. A popular choice is the Ornstein-Uhlenbeck (OU) process [36]: \[\begin{align} \,\mathrm{d}X_t = -X_t \,\mathrm{d}t + \sqrt{2} \,\mathrm{d}B_t, \quad \text{with}\quad X_0 \sim p^\star. \label{eq:forward32SDE} \end{align}\tag{2}\] Here \((B_t)_{t\in [0,T]}\) is a standard Brownian motion in \(\mathbb{R}^d\). A key property of this OU process is that the conditional distribution of \(X_t\) given \(X_0\) remains Gaussian for all \(t\). More precisely, one can verify that \[\begin{align} X_t \mid X_0 \overset{\mathrm d}=c_t X_0 + \sigma_t W_t \label{eq:forward32conditional32dist} \end{align}\tag{3}\] where \(c_t \mathrel{\vcenter{:}}= e^{-t}\), \(\sigma_t \mathrel{\vcenter{:}}=\sqrt{1-e^{-2t}}\), and \(W_t \sim \mathcal{N}(0,I_d)\) is independent of \(X_0\). In particular, the parameter \(t\) fully determines the noise level of the forward process. As \(t\) becomes sufficiently large, \(c_t\) approaches zero and the distribution of \(X_t\) becomes close to the standard Gaussian distribution \(\mathcal{N}(0,I_d)\).

Reverse process.

Running the forward dynamics backward in time transforms Gaussian noise into samples from \(p^\star\), forming the basis of diffusion-based sampling. For the OU process in 2 , its time-reversal SDE is given by \[\begin{align} Y_0 &\sim p_{X_T}, \notag\\ \,\mathrm{d}Y_t &= \big(Y_t + 2\nabla \log p_{X_{T-t}}(Y_t)\big) \,\mathrm{d}t + \sqrt{2} \,\mathrm{d}\overline{B}_t. \label{eq:reverse32SDE} \end{align}\tag{4}\] Here \(p_{X_{T-t}}\) denotes the density of the forward process (2 ) at time \(T-t\) and \(\{ \overline{B}_t\}_{t\in [0,T]}\) is a standard Brownian motion in \(\mathbb{R}^d\). By classical time-reversal results for SDEs [37], this process satisfies \(Y_{T-t}\overset{\mathrm d}=X_t\) for all \(t\in[0,T]\).

The crucial ingredient in the reverse dynamics is the score function of the marginal distributions of the forward process. For a random vector \(X \in \mathbb{R}^d\) with density \(p_X\), its score function is given by \[\begin{align} s^\star_{X}(x) \mathrel{\vcenter{:}}=\nabla \log p_{X}(x) = \frac{\nabla p_{X}(x)}{p_{X}(x)}. \label{eq:score32function} \end{align}\tag{5}\] Since these scores are unknown in practice, they must be estimated from training samples \(\{X^{(i)}\}_{i=1}^n\) drawn from \(p^\star\).

Sampling procedure.

Since \(X_T\overset{\mathrm d}{\to}\mathcal{N}(0,I_d)\) as \(T\to\infty\), diffusion-based sampling can be implemented by initializing the reverse dynamics from \(\mathcal{N}(0,I_d)\) and replacing the true score \(s_{X_t}^\star\) with a learned estimator \(\widehat{s}_{X_t}\). The resulting procedure is summarized in Algorithm 1. We introduce an early stopping time \(\tau>0\) to avoid the small-time regime, where score estimation is most challenging. In practice, the reverse SDE can be implemented using numerical methods such as Euler-Maruyama.

Score estimation reduction.

To estimate the score function \(s^\star_{X_t}\), it is often more convenient to construct score estimator for the following variance-exploding (VE) process \[\begin{align} \,\mathrm{d}Z_t = \,\mathrm{d}B_t, \quad \text{with } Z_0 \sim p^\star. \label{eq:VE32process} \end{align}\tag{6}\] Here \(\{B_t\}_{t\geq 0}\) also denotes the Brownian motion and thus \(Z_t\) follows the distribution \(p^\star \ast \mathcal{N}(0,tI_d)\). It is straightforward to verify that the score functions of \(X_t\) and \(Z_t\) satisfy \[\begin{align} s^\star_{X_t}(x) = \frac{1}{c_t} s^\star_{Z_{h(t)}} (\frac{x}{c_t}) \quad \text{with}\quad h(t)\mathrel{\vcenter{:}}=\frac{\sigma_t^2}{c_t^2}. \label{eq:score32between32VE32process32and32VP} \end{align}\tag{7}\] As a result, it suffices to estimate the score function of \(Z_t\) for any \(t > 0\) and then define \(\widehat{s}_{X_t}(x)\mathrel{\vcenter{:}}=\frac{1}{c_t} \widehat{s}_{Z_{h(t)}}(\frac{x}{c_t})\) as the estimator of \(s^\star_{X_t}(x)\). For notational simplicity, we denote by \(s_t^\star \mathrel{\vcenter{:}}=\nabla \log p_{Z_t}\) the score function of \(Z_t\), and let \(\widehat{s}_t\) denote its estimator. The derivation of (7 ) is provided in Appendix 7.2.

2.2 Assumptions↩︎

In this section, we introduce the assumptions imposed on the target distribution \(p^\star\).

First, to capture low-dimensional, multi-modal structure, we assume that the support of \(p^\star\) is contained in a finite union of low-dimensional linear subspaces.

This assumption provides a tractable abstraction for low-dimensional, multi-modal distributions, where different modes may concentrate on different subspaces. Such union-of-subspaces structure has been widely used in the modeling of heterogeneous high-dimensional data [25] and has also been observed empirically in real-world datasets [23], [38].

For each \(i\in [M]\), let \(p_i^\star \mathrel{\vcenter{:}}= p^\star \mid_{V_i}\) denote the restriction of the target distribution \(p^\star\) to subspace \(V_i\). By Assumption 1, we can decompose the target distribution as \(p^\star = \sum_{i=1}^M p_i^\star\).

For each subspace \(V_i\), let \(A_i \in \mathbb{R}^{d\times k_i}\) be a matrix whose columns form an orthogonal basis of \(V_i\): \[\begin{align} V_i = \mathsf{span}(\mathsf{col}(A_i)), \quad A_i^{\top}A_i = I_{k_i}. \end{align}\] Denote by \(\mathsf{proj}_i:\mathbb{R}^d \rightarrow V_i\) the projection onto \(V_i\), given by \[\begin{align} \mathsf{proj}_i(x) = A_iA_i^{\top}x. \end{align}\]

Next, we impose a mild subgaussian assumption on the target distribution within each subspace.

The subgaussian assumption is fairly mild in the sense that it subsumes any distribution with bounded support, which covers a wide range of practical data such as image data.

3 Main results↩︎

This section introduces our score estimator and presents theoretical guarantees for both score estimation and sampling.

3.1 Algorithm↩︎

Given \(n\) i.i.d.samples \(\{ X^{(i)}\}_{i=1}^n\) drawn from the target distribution \(p^\star\), our goal is to build a score estimator \(\widehat{s}_t\) that learns the score function \(s_t^\star\) of the perturbed data distribution \(p_t = p^\star \ast \mathcal{N}(0,tI_d)\) for any time \(t>0\).

3.1.0.1 Motivation.

Observe that the density \(p_t\) can be written as \[\begin{align} p_t(x) &= \big(p^\star \ast \mathcal{N}(0, t I_d)\big)(x) = \int \varphi_t(x-y;d) p^\star(\!\,\mathrm{d}y), \end{align}\] where \(\varphi_t(x;d) \mathrel{\vcenter{:}}=(2\pi t)^{-d/2}\exp\bigl(-\|x\|_2^2/(2t)\bigr)\) is the density of \(\mathcal{N}(0,tI_d)\) in \(\mathbb{R}^d\). Recall that under the UoS assumption, we can decompose \(p^\star = \sum_{i=1}^M p_i^\star\), yielding \[\begin{align} p_t(x) &= \sum_{i=1}^M \int_{V_i} \varphi_t(x-y;d) p_i^\star( \!\,\mathrm{d}y). \end{align}\] The gradient of \(p_t\) can then be computed as \[\nabla p_t(x) = \sum_{i=1}^M \int_{V_i} -\frac{x-y}{t} \varphi_t(x-y;d) p_i^\star( \!\,\mathrm{d}y) .\] Therefore, the score function \(s_t^\star = {\nabla p_t}/{p_t}\) of \(p_t\) admits the following mixture-type decomposition: \[\tag{10} \begin{align} s_t^\star(x) &= \sum_{i=1}^M \frac{1}{p_t(x)}\int_{V_i} -\frac{x-y}{t} \varphi_t(x-y;d) p_i^\star(\!\,\mathrm{d}y) \notag\\ & \eqqcolon\sum_{i=1}^M w_t(i,x) \cdot s_t(i,x), \end{align} where we define \begin{align} w_t(i,x) &\mathrel{\vcenter{:}}=\frac{\int_{V_i} \varphi_t(x-y;d) p_i^\star(\!\,\mathrm{d}y)}{p_t(x)}\eqqcolon\frac{q_t(i,x)}{p_t(x)}, \tag{11} \\ s_t(i,x)&\mathrel{\vcenter{:}}=\frac{\int_{V_i} -\frac{x-y}{t} \varphi_t(x-y;d) p_i^\star(\!\,\mathrm{d}y)}{\int_{V_i} \varphi_t(x-y;d) p_i^\star(\!\,\mathrm{d}y)} \tag{12}. \end{align}\] Intuitively, \(w_t(i,x)\) can be interpreted as the posterior probability (or effective mixture weight) that \(x\) originates from the \(i\)-th subspace after Gaussian smoothing, while \(s_t(i,x)\) is the score of the corresponding smoothed component.

Our key observation is that each score component \(s_t(i,x) \in \mathbb{R}^d\) admits a favorable normal-tangent decomposition. The normal part is essentially the score of a time-dependent Gaussian distribution and has a closed-form expression, while the tangent component is determined entirely by a \(k_i\)-dimensional score function [16]: \[\begin{align} s_t(i,x) = -\frac{1}{t}\big(x-\mathsf{proj}_i(x)\big) + A_i s_t^\mathsf{low}(i,A_i^{\top}x), \label{eq:s95t40i44x4132decomp} \end{align}\tag{13}\] where \(\mathsf{proj}_i(x)\mathrel{\vcenter{:}}= A_i A_i^{\top} x\) is the projection of \(x\) onto subspace \(V_i\), and \(s_t^\mathsf{low}(i,\cdot):\mathbb{R}^{k_i}\rightarrow \mathbb{R}^{k_i}\) is the score function of the \(k_i\)-dimensional smoothed distribution \(p^\mathsf{low}_i \ast \mathcal{N}(0,tI_{k_i})\) (see 9 ) on the subspace \(V_i\). This decomposition reduces score estimation to a low-dimensional problem: once \(V_i\) is identified, estimating \(s_t(i,\cdot)\) is governed by the difficulty of estimating the low-dimensional score \(s_t^\mathsf{low}(i,\cdot)\) in dimension \(k_i\), rather than the ambient dimension \(d\).

Motivated by this observation, we propose a two-step score estimation procedure. We first use the data to estimate the subspaces \(\{A_i\}_{i=1}^M\) and construct a classification function \(c:\mathbb{R}^d\to[M]\) that assigns points to subspaces (such that \(c(x)=i\) if and only if \(x\in V_i\)). Given these estimates, we then estimate the component scores and mixture weights, and combine them to form the full score estimator. For theoretical clarity, we employ sample splitting, where \(n_0\) samples are used for subspace recovery and the remaining \(N = n-n_0\) samples are used for score estimation.

In what follows, we describe the proposed score estimator in reverse order.

3.1.0.2 Score estimator.

We begin by presenting the score estimator assuming access to subspace estimates \(\{A_i\}_{i=1}^M\) and a classification function \(c(\cdot)\).

Inspired by the score decomposition in (10 ), we construct the score estimator as a weighted combination of score components associated with each subspace: \[\begin{align} \widehat{s}_t(x) \mathrel{\vcenter{:}}=\sum_{i=1}^M \widehat{w}_t(i,x) \widehat{s}_t(i,x) \label{eq:score32estimator}, \end{align}\tag{14}\] where \(\widehat{w}_t(i,x)\) and \(\widehat{s}_t(i,x)\) estimate \(w_t(i,x)\) in 11 and \(s_t(i,x)\) in 12 , respectively.

• Score component estimator \(\widehat{s}_t(i,x)\). In light of the low-dimensional structure in (13 ), it suffices to learn an estimator \(\widehat{s}_t^\mathsf{low}\) for each \(k_i\)-dimensional score \(s_t^\mathsf{low}(i,\cdot)\), and then construct the estimator for the \(d\)-dimensional score \(s_t(i,x)\) associated with \(V_i\) as \[\begin{align} \widehat{s}_t(i,x) \mathrel{\vcenter{:}}=-\frac{x-\mathsf{proj}_i(x)}{t} + A_i \widehat{s}_t^\mathsf{low}(i,A_i^{\top}x). \label{eq:s95hat95t40i44x4132decomp} \end{align}\tag{15}\]

  To build \(\widehat{s}_t^\mathsf{low}\), let \(\mathcal{C}_i \mathrel{\vcenter{:}}=\{j\in [N]: X^{(j)} \in V_i \}\) denote the index set of samples belonging to subspace \(V_i\). Since \(s_t^\mathsf{low}(i,\cdot)\) is the score of the smoothed low-dimensional distribution \(p_i^\mathsf{low}\ast\mathcal{N}(0,tI_{k_i})\), we first estimate its density using the Gaussian kernel estimator \[\begin{align} \widehat{g}_t(i,x) \mathrel{\vcenter{:}}=\frac{1}{|\mathcal{C}_i|} \sum_{j\in \mathcal{C}_i} \varphi_{t}(x - A_i^{\top} X^{(j)};k_i) \label{eq:estimator32for32p95t94l40i44x41}, \end{align}\tag{16}\] where \(\varphi_{t}(x;k_i)\) denotes the density of \(\mathcal{N}(0,tI_{k_i})\). We then define the \(k_i\)-dimensional score estimator as \[\begin{align} \widehat{s}_t^\mathsf{low}(i,x) \mathrel{\vcenter{:}}=\mathsf{clip}_{R} \bigg( \frac{\nabla \widehat{g}_t(i,x)}{\widehat{g}_t(i,x)} \psi \Big(\widehat{g}_t(i,x);\frac{\log N}{N(2\pi t)^{k_i/2}}\Big)\bigg). \label{eq:low-dimensional32score32estimator} \end{align}\tag{17}\] Here, \(\psi(x;\eta) \mathrel{\vcenter{:}}=\mathbb{1}\{x \geq \eta\}\) is a thresholding function, and the clip operator is defined by \[\begin{align} \mathsf{clip}_r(z)\mathrel{\vcenter{:}}= \begin{cases} z, &\|z\|_2 \leq r; \\ \frac{z}{\|z\|}\cdot r,&\text{otherwise}. \end{cases} \end{align}\] We set the clipping radius to be \(R = \sqrt{2\log N/t}\).

  In words, we first form the plug-in estimator \(\nabla \widehat{g}_t/\widehat{g}_t\) using the kernel density estimator 16 . We then apply a thresholding rule \(\psi(\widehat{g}_t;\eta_t)\), which regularizes this ratio according to the estimated density level \(\widehat{g}_t\) and the threshold \(\eta_t = N^{-1}(2\pi t)^{-k_i/2}\log N\) that depends on the sample size \(N\) and time \(t\). Specifically, in low-density regions where \(\nabla \widehat{g}_t/\widehat{g}_t\) is unstable due to small denominators and limited data, the resulting score estimator \(\widehat{s}_t\) is set to zero. This regularization is important not only for controlling the subsequent estimation error, but also for improving generalization by preventing the estimator from closely fitting empirical artifacts.

• Weight estimator \(\widehat{w}_t(i,x)\). As for the mixture weight \(w_t(i,x)\) defined in (11 ), we first construct Gaussian kernel density estimators for \(p_t(x)\) and \(q_t(i,x)\): \[\begin{align} \widehat{p}_t(x) \mathrel{\vcenter{:}}=\frac{1}{N} \sum_{j=1}^N \varphi_t(x-X^{(j)};d), \tag{18} \end{align} and for any i\in [M], \begin{align} \widehat{q}_t(i,x) \mathrel{\vcenter{:}}=\frac{1}{N} \sum_{j=1}^N \varphi_t(x-X^{(j)};d) \mathbb{1}_{\{ c(X^{(j)})=i\}}. \tag{19} \end{align}\] We then define the weight estimator as \[\begin{align} \widehat{w}_t(i,x):= \frac{\widehat{q}_t(i,x)}{\widehat{p}_t(x)} \mathbb{1}_{\{ x \in \mathcal{G}_t(i)\}}, \label{eq:estimate32w95t40i44x41} \end{align}\tag{20}\] where \(\mathcal{G}_t(i)\) is a set given by \[\begin{align} \mathcal{G}_t(i) \mathrel{\vcenter{:}}=\Big\{ x: \| x-\mathsf{proj}_i(x)\|_2\leq R_t(i)\Big\}, \label{eq:regularization32set} \end{align}\tag{21}\] with \(R_t(i) = C_R\sqrt{ td \log (Ndt^{k_i/2})}\) for some universal constant \(C_R >0\).

  In a word, the weight estimator \(\widehat{w}_t(i,x)\) is a plug-in estimator for the true weight 11 , up to some low-probability set under \(p_t\). The indicator function \(\mathbb{1}_{\{ x \in \mathcal{G}_t(i)\}}\) is introduced for technical convenience in the analysis and could be removed with a shaper argument.

3.1.0.3 Subspace recovery.

Finally, we briefly discuss how the subspace estimates \(\{A_i\}_{i=1}^M\) and classification function \(c(\cdot)\) can be obtained from training data. This is a classical subspace clustering problem. Under standard identifiability and separation conditions, and assuming known bounds on the number of subspaces \(M\) and the maximal intrinsic dimension \(k\), several polynomial-time methods can recover the underlying subspaces and cluster assignments, such as sparse subspace clustering [39], thresholding-based subspace clustering [40] and greedy subspace clustering [41]. From a statistical perspective, this geometric recovery step is typically less demanding than learning the full target distribution.

3.2 Theoretical guarantees↩︎

We now state our theoretical guarantees for the proposed score estimator and the resulting sampler.

We first present the \(L^2\) error for the proposed score estimator \(\widehat s_t\) in (14 ). The proof can be found in Appendix 7.1.

In words, Theorem 1 shows that the convergence rate of the \(L^2\) error (with respect to diffusion time \(t\)) of the proposed score estimator depends on the intrinsic dimension \(k\), rather than the ambient dimension \(d\). This yields a substantial improvement over existing rate-optimal score estimation guarantees for general distributions [13]–[15], [42],which do not exploit intrinsic low-dimensional data structures and therefore suffer from the curse of dimensionality.

Moreover, we note that exact subspace recovery can be achieved with high probability using \(n_0=C_{\mathsf{sc}}M^2k\log n\) samples for subspace clustering, for a sufficiently large constant \(C_{\mathsf{sc}}>0\). This sample size is negligible compared with the remaining \(N=n-n_0\) samples used for score estimation, provided that \(n\) is sufficiently large.

We next translate the resulting score estimation guarantee into a sampling guarantee for the diffusion sampler. The proof is deferred to Appendix 7.2.

Theorem 2 provides the convergence rate of the \(W_1\)-sampling error for diffusion sampling equipped with the proposed score estimator. By exploiting the intrinsic low-dimensional data structure through kernel-based score estimation, the resulting convergence rate (with respect to sample size \(n\)) is governed by the intrinsic dimension \(k\), rather than the ambient dimension \(d\), with \(d\) appearing only linearly through the prefactor.

Several remarks are in order:

• (Near-)minimax optimality. The minimax risk of estimating a \(k\)-dimensional density scales as \(n^{-\frac{1}{k \vee 2}}\) [24]. Since sampling is always harder than density estimation, Theorem 2 shows that our sampling algorithm is minimax optimal (up to logarithmic factors).

• Sample complexity. The error bound in 22 demonstrates that in order to achieve an \(\varepsilon\)-accurate sampling in 1-Wasserstein distance, it suffices to have \(\varepsilon^{- (k \vee 2)}\) samples, up to logarithmic factors, thereby breaking the curse of dimensionality that plagues prior results [13]–[15], [42].

• Weak assumptions on the target distribution. Our results do not rely on stringent structural conditions commonly imposed in earlier work, such as smooth densities/scores, log-concavity, or exactly Gaussian components. As a consequence, our framework applies to a broader class of multi-modal distributions of practical interest. Moreover, to the best of our knowledge, even in the single low-dimensional setting, our result is the first to achieve a (near-)optimal rate under only a subgaussian assumption on the target distribution, without extra assumptions on the score or density.

4 Analysis↩︎

This section provides the proof sketches for Theorems 1–2.

4.0.0.1 Proof sketch of Theorem 1.

In light of the expressions of the true score (10 ) and the score estimator (14 ), the \(L^2\)-error decomposes as: \[\begin{align} &\mathbb{E}\big[\|\widehat{s}_t(Z_t)-s_t^\star(Z_t)\|_2^2\big] \notag \\ &\,\,\, \lesssim \sum_{i=1}^M \int \mathbb{E}\big[\big(w_t(i,x)-\widehat{w}_t(i,x)\big)^2 \|\widehat{s}_t(i,x) \|_2^2 \big] p_t(x) \!\,\mathrm{d}x \notag\\ &\quad + \sum_{i=1}^M \int w_t^2(i,x) \, \mathbb{E}\big[\| s_t(i,x) - \widehat{s}_t(i,x) \|_2^2 \big] p_t(x) \! \,\mathrm{d}x, \label{eq:22} \end{align}\tag{23}\] where the expectation in the first line is taken over both \(Z_t\sim p_t\) and the i.i.d.samples \(\{ X^{(i)}\}_{i=1}^N\), while the second and the third line only take expectation over the samples.

This decomposition suggests that we need to control the mean squared error of both the weight estimator \(\widehat{w}_t(i,x)\) and score estimator \(\widehat{s}_t(i,x)\).

For the weight estimator, we recall the true weight \(w_{t}(i,x) = q_t(i,x)/p_t(x)\) in (11 ). The mean squared error is controlled in the following lemma, with proof deferred to Appendix 8.2.

With Lemma 1 in hand, we can now bound the first term in 23 associated with weight estimation. It is easy to see \(p_t\) is \(\sqrt{\sigma^2+t}\)-subgaussian. Define \[\mathcal{B}_t \mathrel{\vcenter{:}}=\big\{ x \in \mathbb{R}^d: \| A_i^{\top} x\|_{2} \leq B_t, \forall i \in [M] \big\},\] with \(B_t \mathrel{\vcenter{:}}= C_B \sqrt{k(\sigma^2+t)\log N}\) for some universal constant \(C_B>0\). One can show that the estimation error within \(\mathcal{B}_t\) dominates since \(\mathcal{B}_t^c\) is a low probability region w.r.t. \(p_t\). This allows us to apply Lemma 1 to derive the following bound that only depends on \(k_i\): \[\begin{align} & \int \mathbb{E} \big[ (w_t(i,x)-\widehat{w}_t(i,x))^2\big]p_t(x)\,\mathrm{d}x \notag \\ &\quad \lesssim \sum_{j=1}^M\int_{\mathcal{B}_t} \frac{1}{N t^{d/2}} e^{-\frac{1}{2t} \| x-A_jA_j^{\top}x\|_2^2} \,\mathrm{d}x \notag \\ & \quad = \widetilde{O}\bigg( \sum_{j=1}^M \frac{(\sigma^2+t)^{k_j/2}}{Nt^{k_j/2}} \bigg). \label{temp1} \end{align}\tag{24}\]

Regarding the second term in 23 associated with score estimation, notice that \(w_t(i,x) = q_t(i,x)/p_t(x) \leq 1\). Thus, it suffices to bound \[\begin{align} &\int w_t^2(i,x) \mathbb{E}\big[\| s_t(i,x) - \widehat{s}_t(i,x) \|_2^2 \big] p_t(x) \,\mathrm{d}x \\ & \quad \leq \int \mathbb{E}\big[\| s_t(i,x) - \widehat{s}_t(i,x) \|_2^2 \big] q_t(i,x) \,\mathrm{d}x. \end{align}\] To this end, let \(N_i\mathrel{\vcenter{:}}=\sum_{j=1}^N \mathbb{1}_{\{c(X^{(j)})=i\}}\) denote the sample size that we use to estimate the score on subspace \(V_i\). The following lemma provides a mean squared error bound for the score estimator \(\widehat{s}_t(i,x)\) in (15 ). The proof is deferred to Appendix 8.4.

Combining 24 –26 completes the proof of Theorem 1.

4.0.0.2 Proof sketch of Theorem 2.

As we will show in Lemma 3 in Appendix 8, \(n_0=C_{\mathsf{sc}}M^2k\log n\) samples suffice to recover the subspaces exactly with probability at least \(1-Mn^{-10}\), for a sufficiently large constant \(C_{\mathsf{sc}}>0\). Thus, for large \(n\), the remaining sample size for score estimation satisfies \(N=n-n_0 \geq n/2\). Conditioned on the exact subspace recovery event, we apply the score estimation error bound in Theorem 1. We then relate the 1-Wasserstein error between the target distribution \(p^\star\) and the generated distribution \(p_{\widehat{Y}_{T-\tau}}\) to the integral of the score estimation error over time via the following stability bound [9], [17], [18]: \[\begin{align} & \mathbb{E}\big[W_1(p^\star,p_{\widehat{Y}_{T-\tau}})\big] \notag\\ & \,\,\,\lesssim \sqrt{d} \, \Big( \sqrt{\tau} + \delta + e^{-T}+ \notag\\ & \quad \sum_{j=0}^{L-1} \sigma_{T_{j+1}} \sqrt{\log \frac{1}{\delta} \int_{T_j}^{T_{j+1}} \mathop{\mathbb{E}}\big[ \| \widehat{s}_{X_t}(X)-s_{X_t}^\star(X)\|_2^2\big] \,\mathrm{d}t}\Big) \label{eq:W195convergence95rate} \end{align}\tag{27}\] for \(0<T_0=\tau < T_1< \cdots <T_L = T\) and any \(\delta>0\). Here the expectation is taken over the randomness of samples and \(X \sim p_{X_t}\).

By Theorem 1, together with \(N\geq n/2\) and the score relation in 7 , one can show that \[\begin{align} & \int_{\mathbb{R}^d}\mathbb{E}\big[ \| \widehat{s}_{X_t}(x)-s^\star_{X_t}(x)\|_2^2\big] p_{X_t}(x) \,\mathrm{d}x \\ &\quad =\widetilde{O}\bigg( \frac{dM^3}{n} \Big( \frac{1}{h(t)}+\frac{\sigma^{k\vee2}}{h(t)^{{(k\vee2)}/2+1}}\Big)\bigg). \end{align}\] Observe that \(h(t) = 2e^{2t} = 2/c_t^2,\) and thus \[\begin{align} & \int_{T_j}^{T_{j+1}} \int_{\mathbb{R}^d}\mathbb{E}\big[ \| \widehat{s}_{X_t}(x)-s^\star_{X_t}(x)\|_2^2\big] p_{X_t}(x) \,\mathrm{d}x \,\mathrm{d}t \\ & = \widetilde{O}\bigg( \frac{dM^3}{n} \Big( \log \frac{h(T_{j+1})}{h(T_j)} + \frac{2\sigma^{k\vee2}}{k\vee2}\frac{1}{h(T_j)^{(k\vee2)/2}} \Big)\bigg). \end{align}\]

We then choose a dyadic partition of the time interval by setting \[\begin{align} T_{j+1} = 2T_j, \quad T \asymp \log n, \quad \tau\asymp n^{-\gamma} \end{align}\] for some \(\gamma>0\) that will be specified below. Plugging these into 27 yields \[\begin{align} &\sum_{j=0}^{L-1} \sqrt{(1-e^{-4T_j})\int_{T_j}^{T_{j+1}} \mathbb{E}\big[ \| \widehat{s}_{X_t}(X)-s^\star_{X_t}(X)\|_2^2\big] \,\mathrm{d}t} \\ &\lesssim \sqrt{d} \, M^{3/2} \,\mathsf{poly}\log n \begin{cases} \frac{1}{\sqrt{n}}\tau^{-\frac{k}{4}+\frac{1}{2}}, \quad &k\geq 2 \\ \frac{1}{\sqrt{n}}, & k=1 \end{cases} \end{align}\]

Finally, taking \(\delta = n^{-1}\) and \(\tau = n^{-2/k}\), we conclude \[\begin{align} \mathbb{E}\big[W_1(p^\star,p_{\widehat{Y}_{T-\tau}})\big] \lesssim \frac{dM^{3/2}}{n^{1/(k\vee 2)}} \mathsf{poly}\log n. \end{align}\]

5 Numerical results↩︎

In this section, we provide numerical experiments to validate the theoretical findings of our paper. Since evaluating the Wasserstein distance is computationally prohibitive in high dimensions, we focus on the \(L^2\)-score estimation guarantee in Theorem 1.

We consider a synthetic target distribution in \(\mathbb{R}^d\) with \(d=48\). The support is a union of \(M=128\) randomly generated linear subspaces, each with intrinsic dimension \(k=3\). The restriction of the distribution to each subspace is chosen to be a two-component Gaussian mixture with randomized parameters. We construct the kernel-based score estimator using \(N=50,000\) i.i.d samples from the target distribution. For each time value \(t\), we approximate the \(L^2\)-score estimation error by Monte Carlo using \(10{,}000\) independent samples from \(p_t\) and average the result over \(20\) independent training datasets generated from the same target distribution.

Figure 2 plots the empirical \(L^2\) score estimation error versus the diffusion time \(t\). The observed scaling is consistent with the prediction of Theorem 1, where the score estimation error is governed by the intrinsic dimension of the data, rather than the ambient dimension. In particular, despite the relatively large ambient dimension \(d=48\), the empirical error exhibits a substantially milder dependence on \(t\) than would be suggested by ambient-dimensional worst-case bounds.

6 Discussion↩︎

This paper has studied the sample complexity of diffusion models for learning distributions supported on a union of low-dimensional subspaces, a tractable model for low-dimensional, multi-modal data commonly observed in practice. We construct a kernel-based score estimator and prove that diffusion-based sampling can learn the target distribution using at most \(\widetilde{O}(\varepsilon^{-(k\vee2)})\) samples, where \(k\) represents the intrinsic dimension. Our result shows that diffusion models can achieve statistical optimality by exploiting intrinsic low-dimensional structure while naturally accommodating multi-modal data.

Building on the results of our paper, several directions remain open for future work. First, it would be valuable to extend our theory to practical NN-based score estimators. A natural approach is to analyze an ERM estimator over a NN class, which requires controlling both approximation and generalization errors. The main challenge lies in the approximation step: constructing a NN approximation of the target score whose complexity depends on the intrinsic dimension. Our score decomposition, together with the kernel-based construction, makes the relevant low-dimensional approximation targets explicit. Thus, our construction provides a concrete roadmap for future analysis of NN score estimators. Second, it would be important to extend our framework to real data with more complex geometric structures, such as classes residing on manifolds of varying dimensions. A natural first step is to consider distributions concentrated near a union of low-dimensional subspaces. In this setting, once the underlying subspaces are learned from noisy observations, the score still admits an analogous normal-tangent decomposition, where the tangent component is governed by a low-dimensional score and the normal component remains Gaussian with an enlarged variance. This suggests that the framework developed here could be extended to noisy low-dimensional models. Third, it would be interesting to develop a fully end-to-end convergence analysis that explicitly accounts for both score estimation error and discretization of the reverse-time dynamics (SDE/ODE), ideally yielding non-asymptotic bounds in Wasserstein distance under mild distributional conditions.

Acknowledgements↩︎

C.Cai is supported in part by the NSF grant DMS-2515333.

Impact Statement↩︎

This paper presents work whose goal is to advance the field of diffusion model theory. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.

7 Proof of theorems↩︎

7.1 Proof of Theorem 1↩︎

In this section, we consider the score estimation error under exact subspace recovery. We first define the following event set \(\mathcal{A}\) as \[\begin{align} \mathcal{A} \mathrel{\vcenter{:}}=\Big\{ N_i \geq \frac{N}{2c_{p}M}, \quad \forall i \in [M] \Big\}. \end{align}\] Here \(N_i\mathrel{\vcenter{:}}=\sum_{j=1}^N \mathbb{1}_{\{ c(X^{(j)}) = i\}}\) is a random number of samples on \(V_i\). The following claim tells us that, event \(\mathcal{A}\) happens with high probability.

The proof follows from concentration inequality and can be found in Appendix 9.1.

We first consider the error on event \(\mathcal{A}\), where we have enough samples on each subspace. Notice the score decomposition (10 ) and the score estimator in (14 ), the \(L^2\) estimation error can be written as \[\begin{align} \int \mathbb{E}\Big[\big\| s_t^\star(x)-\widehat{s}_t(x)\big\|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x &= \int \mathbb{E}[\| \sum_{i=1}^M \big(w_t(i,x)s_t(i,x)- \widehat{w}_t(i,x)\widehat{s}_t(i,x)\big)\|_2^2 \mathbb{1}_{\mathcal{A}}] p_t(x) \!\,\mathrm{d}x \\ & \leq M \sum_{i=1}^M \int \mathbb{E}\Big[ \big\| w_t(i,x)s_t(i,x)-\widehat{w}_t(i,x)\widehat{s}_t(i,x) \big\|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x \quad (\text{C-S Ineq})\\ &\eqqcolon M \sum_{i=1}^M L_i . \end{align}\] We further decompose \(L_i\) using the error of weight estimator and the error of score estimator respectively, \[\begin{align} L_i &\mathrel{\vcenter{:}}=\int \mathbb{E}\Big[ \big\| w_t(i,x)s_t(i,x)-\widehat{w}_t(i,x)\widehat{s}_t(i,x) \big\|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x \\ & \lesssim \int \mathbb{E} \Big[ \big(w_t(i,x)-\widehat{w}_t(i,x)\big)^2 \cdot \|\widehat{s}_t(i,x) \|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x + \int \mathbb{E} \Big[ \big(w_t(i,x))^2 \cdot \| s_t(i,x) - \widehat{s}_t(i,x) \|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x \\ & \eqqcolon L_{i,1} + L_{i,2}. \end{align}\]

Bound of \(L_{i,1}\).

We utilize a set \(\mathcal{B}_t\) here and further decompose \(L_{i,1}\) as \[\begin{align} L_{i,1} &\mathrel{\vcenter{:}}=\int \mathbb{E} \Big[ \big(w_t(i,x)-\widehat{w}_t(i,x)\big)^2 \cdot \|\widehat{s}_t(i,x) \|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \! \,\mathrm{d}x \\ & \leq \int_{\mathcal{G}_{t}(i) \cap \mathcal{B}_t} \mathbb{E} \Big[ \big(w_t(i,x)-\widehat{w}_t(i,x)\big)^2 \cdot \|\widehat{s}_t(i,x) \|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x \\ &+ \int_{\mathcal{G}_t(i)^c \cap \mathcal{B}_t} \mathbb{E} \Big[ \big(w_t(i,x)-\widehat{w}_t(i,x)\big)^2 \cdot \|\widehat{s}_t(i,x) \|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x \\ & + \int_{\mathcal{B}_t^c} \mathbb{E} \Big[ \big(w_t(i,x)-\widehat{w}_t(i,x)\big)^2 \cdot \|\widehat{s}_t(i,x) \|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x \\ &\eqqcolon\kappa_1 + \kappa_2 + \kappa_3, \end{align}\] here we recall that \(\mathcal{G}_t(i)\) is the regularization set defined in (21 ) and \(\mathcal{B}_t\) is defined as, \[\begin{align} \mathcal{B}_t \mathrel{\vcenter{:}}=\Big\{ x \in \mathbb{R}^d: \| A_i^{\top} x\|_{2} \leq B_t, \forall i \in [M]\Big\}, \quad \text{with }B_t \mathrel{\vcenter{:}}= C_B \sqrt{k(\sigma^2+t)\log (N)}. \label{eq:def32of32B95t} \end{align}\tag{28}\] for certain absolute constant \(C_B>0\).

The following claim tells us that the probability outside \(\mathcal{B}_t\) is negligible and see Appendix 9.2 for its proof.

• For \(\kappa_1\), notice that for \(x\in \mathcal{G}_t(i)\), \[\begin{align} \| \widehat{s}_t(i,x)\| _2 \leq \frac{R_t(i)}{t} + \sqrt{\frac{2}{t}\log N} \end{align}\] using Lemma 2. Hence, \[\begin{align} \kappa_1 &\mathrel{\vcenter{:}}=\int_{\mathcal{G}_{t}(i) \cap \mathcal{B}_t} \mathbb{E} \Big[ \big(w_t(i,x)-\widehat{w}_t(i,x)\big)^2 \cdot \|\widehat{s}_t(i,x) \|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x \\ & \lesssim \Big( \frac{R_t(i)^2}{t^2} + \frac{1}{t}\Big) \cdot \int_{\mathcal{G}_t(i) \cap \mathcal{B}_t} \frac{1}{(2\pi t)^{d/2}N\cdot p_t(x)} \Big(\sum_{j=1}^M e^{-\frac{1}{2t}\|x-\mathsf{proj}_j(x)\|_2^2} \cdot q_t(j,x) \Big) \!\,\mathrm{d}x \quad (\text{Lemma \ref{lem:MSE32for32weight}}) \\ &\leq \Big( \frac{R_t(i)^2}{t^2} + \frac{1}{t}\Big) \cdot \sum_{j=1}^M \int_{\mathcal{B}_t} \frac{1}{N(2\pi t)^{d/2}} e^{-\frac{1}{2t} \| x-A_jA_j^{\top}x\|_2^2} \!\,\mathrm{d}x \\ & \leq \Big( \frac{R_t(i)^2}{t^2} + \frac{1}{t}\Big) \cdot \sum_{j=1}^M \frac{2^{k_j} B_t^{k_j}}{N(2\pi t)^{k_j/2}} \\ &\lesssim \frac{d}{t}\cdot \frac{M}{N}\Big( 1+\frac{\sigma^k}{t^{k/2}}\Big) \cdot \big(\mathsf{poly}\log N +\log t\big) \end{align}\] In the third inequality, we apply the following changing variable technique in integration and then use Tonelli’s Theorem, \[\begin{align} z = T_jx \mathrel{\vcenter{:}}= \begin{pmatrix} A_j^{\top}x\\ P_{V_j^{\perp}}^{\top} x \end{pmatrix} \label{eq:changing32var} \end{align}\tag{29}\] here columns of \(P_{V_j^{\perp}}\) denotes an orthogonal basis of \(V_j^{\perp}\) and this is an orthogonal transform with \[\begin{align} \|x-A_jA_j^{\top}x\|_2^2 = \|P_{V_j^{\perp}}^{\top} x\|_2^2. \end{align}\]

• For \(\kappa_2\), we have, \[\begin{align} \kappa_2 &\mathrel{\vcenter{:}}=\int_{\mathcal{G}_t(i)^c \cap \mathcal{B}_t} \mathbb{E} \Big[ \big(w_t(i,x)-\widehat{w}_t(i,x)\big)^2 \cdot \|\widehat{s}_t(i,x) \|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x \\ & = \int_{\mathcal{G}_t(i)^c \cap \mathcal{B}_t} \frac{q_t^2(i,x)}{p_t^2(x)} \mathbb{E} [\|\widehat{s}_t(i,x) \|_2^2] \cdot p_t(x)\!\,\mathrm{d}x \\ &\lesssim \int_{\mathcal{G}_t(i)^c \cap \mathcal{B}_t} \Big( \frac{1}{t} + \frac{\|x-A_iA_i^{\top}x\|_2^2}{t^2}\Big) \cdot q_t(i,x) \!\,\mathrm{d}x \quad (\text{Lemma \ref{lem:bound32of32s95t40i44x41} + } q_t(i,x) \leq p_t(x)) \end{align}\] Since it holds that, \[\begin{align} q_t(i,x) \mathrel{\vcenter{:}}=\int_{V_i} \varphi_t(x-y;d) p_i^\star(\!\,\mathrm{d}y) \leq p_i^\star(V_i)\cdot (2\pi t)^{-d/2} e^{-\frac{1}{2t}\|x-A_iA_i^{\top}x\|_2^2}, \end{align}\] we apply the similar linear transform as (29 ), \[\begin{align} \kappa_2 &\lesssim p_i^\star(V_i)\int_{\mathbb{R}^d} \mathbb{1}_{ \|z\|\leq B_t} \cdot \mathbb{1}_{\| z_{k_i+1:d}\|_2\geq R_t(i)} \Big( \frac{1}{t} + \frac{\| z_{k_i+1:d}\|_2^2}{t^2}\Big) \cdot (2\pi t)^{-d/2}e^{-\frac{1}{2t} \|z_{k_i+1:d}\|_2^2} \!\,\mathrm{d}z \\ &\lesssim p_i^\star(V_i) (2\pi t)^{-k_i/2} \int_{\mathbb{R}^{k_i}} \mathbb{1}_{\| z_{1:k_i}\|\leq B_t} \Big( \frac{1}{t} + \frac{d\cdot R_t(i)^2 + d^2tC_1}{t^2}\Big) \cdot \exp \Big( -\frac{R_t(i)^2}{d\cdot tC_1}\Big) \!\,\mathrm{d}z_{1:k_i} \quad (\text{Lemma \ref{lem:tail32bdd32for32subgaussian}})\\ &\lesssim \frac{2^{k_i} B_t^{k_i}}{(2\pi t)^{k_i/2}} p_i^\star(V_i)\cdot \Big( \frac{d^2}{t} + \frac{d\cdot R_t(i)^2}{t^2}\Big) \exp \Big( -\frac{R_t(i)^2}{d\cdot tC_1}\Big) \\ &\lesssim \frac{d}{Nt}\Big( 1+\frac{\sigma^{k_i}}{t^{k_i/2}}\Big)\cdot \mathsf{poly}\log N \lesssim \frac{d}{Nt}\Big( 1+\frac{\sigma^{k}}{t^{k/2}}\Big)\cdot \mathsf{poly}\log N. \end{align}\] Here \(C_1>0\) is a universal constant which is related with the sub-gaussian norm of standard Gaussian distribution.

• For \(\kappa_3\), \[\begin{align} \kappa_3 &\mathrel{\vcenter{:}}=\int_{\mathcal{B}_t^c} \mathbb{E} \Big[ \big(w_t(i,x)-\widehat{w}_t(i,x)\big)^2 \cdot \|\widehat{s}_t(i,x) \|_2^2 \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x \\ &\leq \int_{\mathcal{B}_t^c} \mathbb{E} \Big[ \|\widehat{s}_t(i,x) \|_2^2\Big] p_t(x) \!\,\mathrm{d}x \\ & \lesssim \int_{\mathcal{B}_t^c} \Big( \frac{\| x-A_iA_i^{\top}x\|_2^2}{t^2} + \frac{1}{t}\Big) \cdot p_t(x)\!\,\mathrm{d}x \quad (\text{Lemma \ref{lem:bound32of32s95t40i44x41}}) \\ &\leq \int_{\mathcal{B}_t^c} \Big( \frac{\|x\|_2^2}{t^2} + \frac{1}{t} \Big) p_t(x)\!\,\mathrm{d}x \\ &\lesssim \frac{dM(\sigma^2+t)}{t^2N^2} \mathsf{poly}\log N \quad (\text{Claim \ref{claim:bdd32of32set32B95t}}) \\ &\lesssim \frac{d}{tN} \Big( 1+\frac{\sigma^{k\vee2}}{t^{{(k\vee2)}/2}}\Big)\mathsf{poly}\log N . \end{align}\]

• For \(L_{i,1}\), we could sum them up, \[\begin{align} L_{i,1} &\leq \kappa_1 + \kappa_2 + \kappa_3\\ &\lesssim \frac{dM}{Nt} \Big( 1+\frac{\sigma^{k\vee2}}{t^{{(k\vee2)}/2}}\Big) \cdot \big(\mathsf{poly}\log N+\log t\big) . \end{align}\]

Bound of \(L_{i,2}\).

\[\begin{align} L_{i,2} &\mathrel{\vcenter{:}}=\int \mathbb{E} \Big[ \big(w_t(i,x))^2 \cdot \| s_t(i,x) - \widehat{s}_t(i,x) \|_2^2 \cdot \mathbb{1}_{\mathcal{A}}\Big] p_t(x) \!\,\mathrm{d}x \\ &= \int \frac{q_t^2(i,x)}{p_t^2(x)} \mathbb{E} \Big[ \big\| s_t(i,x) - \widehat{s}_t(i,x) \big\|_2^2\cdot \mathbb{1}_{\mathcal{A}} \Big] p_t(x) \!\,\mathrm{d}x \\ & \leq \int \mathbb{E}\Big[ \big\| s_t(i,x) - \widehat{s}_t(i,x)\big \|_2^2 \cdot \mathbb{1}_{\{ N_i\geq N/2c_{p}M\}} \Big] q_t(i,x) \! \,\mathrm{d}x \quad (\text{since }q_t(i,x) \leq p_t(x)) \\ &\lesssim p^\star(V_i) \frac{c_{p}M(4/\sqrt{\pi})^{k_i} }{N} \Big( \frac{1}{t} + \frac{\sigma^{k_i}}{t^{k_i/2+1}}\Big) \mathsf{poly}\log N \quad (\text{Lemma \ref{lem:bound32of32s95t40i44x41}}). \end{align}\]

Bound of error on \(\mathcal{A}\).

Therefore, \[\begin{align} \int \mathbb{E}[\| s_t^\star(x)-\widehat{s}_t(x)\|_2^2\mathbb{1}_{\mathcal{A}}] p_t(x)\!\,\mathrm{d}x &\leq M\sum_{i=1}^M (L_{i,1}+L_{i,2}) \\ &\lesssim \frac{dM^3}{Nt} \Big( 1+\frac{\sigma^{k\vee2}}{t^{{(k\vee2)}/2}}\Big) \cdot \big(\mathsf{poly}\log N +\log t\big). \end{align}\] Here for simplicity, we omit constant terms \(c_{p}\) and terms that are only related with intrinsic dimension \(k\).

Bound of error outside \(\mathcal{A}^c\).

\[\begin{align} \int \mathbb{E}[\| s_t^\star(x)-\widehat{s}_t(x)\|_2^2\mathbb{1}_{\mathcal{A}^c}] p_t(x)\!\,\mathrm{d}x &\lesssim \Big(\int \|s_t(x)\|_2^2p_t(x) \!\,\mathrm{d}x\Big) \cdot \mathbb{P}[\mathcal{A}^c] \\ &+ \int \Big( \frac{1}{t} + \frac{\|x\|_2^2}{t^2} \Big) p_t(x)\! \,\mathrm{d}x \cdot \mathbb{P}[\mathcal{A}^c] \quad (\text{Lemma \ref{lem:bound32of32s95t40i44x41}})\\ &\lesssim \Big(\frac{d}{t} + \frac{d(\sigma^2+t)}{t^2} \Big) \cdot 2M e^{-\frac{N}{2c_{p}^2M^2}} \end{align}\] Here we apply Lemma 11 in [15] and that \(p^\star \ast \mathcal{N}(0,tI_d)\) is \(c\sqrt{\sigma^2+t}\) subgaussian r.v as we have proven in Appendix 9.2. Notice that, for large enough \(N\), \(e^{-\frac{N}{2c_{p}^2M^2}} \leq \frac{M^2c_{p}^2}{N}\) and thus, \[\begin{align} \int \mathbb{E}[\| s_t^\star(x)-\widehat{s}_t(x)\|_2^2 \mathbb{1}_{\mathcal{A}^c}] p_t(x)\!\,\mathrm{d}x \lesssim \frac{dM^3}{Nt}\Big( 1+\frac{\sigma^2}{t}\Big). \end{align}\]

Summary.

In summary, our analysis above shows that, \[\begin{align} \int \mathbb{E}\big[\|\widehat{s}_t(x)-s_t^\star(x)\|_2^2\big] p_t(x) \!\,\mathrm{d}x &= \int \mathbb{E}[\| s_t^\star(x)-\widehat{s}_t(x)\|_2^2 \mathbb{1}_{\mathcal{A}}] p_t(x) \!\,\mathrm{d}x + \int \mathbb{E}[\| s_t^\star(x)-\widehat{s}_t(x)\|_2^2 \mathbb{1}_{\mathcal{A}^c}] p_t(x) \!\,\mathrm{d}x \\ &\lesssim \frac{dM^3}{Nt}\Big( 1+\frac{\sigma^{k\vee2}}{t^{{(k\vee2)}/2}}\Big) \cdot \big(\mathsf{poly}\log N+\log t\big) . \end{align}\] This proves Theorem 1.

7.2 Proof of Theorem 2↩︎

Proof of (7 ).

Recall the forward process (3 ), it holds that, \[\begin{align} p_{X_t}(x) &= \int p_{X_t|X_0}(x|y) p^\star(\! \,\mathrm{d}y) \\ &= \int \big(2\pi\sigma_t^2\big) ^{-d/2} e^{-\frac{1}{2\sigma_t^2}\|x-c_ty\|_2^2}p^\star( \! \,\mathrm{d}y) \end{align}\] Hence, \[\begin{align} \nabla_x p_{X_t}(x)&= (2\pi\sigma_t^2\big) ^{-d/2}\cdot \int \Big( -\frac{x-c_ty}{\sigma_t^2}\Big)e^{-\frac{1}{2\sigma_t^2}\|x-c_ty\|_2^2}p^\star(\! \,\mathrm{d}y) \end{align}\] Therefore, its score function, \[\begin{align} s^\star_{X_t}(x)&= \frac{\nabla_xp_{X_t}(x)}{p_{X_t}(x)} \\ &= \frac{\int \Big( -\frac{x-c_ty}{\sigma_t^2}\Big)e^{-\frac{1}{2\sigma_t^2}\|x-c_ty\|_2^2}p^\star(\! \,\mathrm{d}y) }{\int e^{-\frac{1}{2\sigma_t^2}\|x-c_ty\|_2^2}p^\star(\! \,\mathrm{d}y)} \end{align}\] Similarly, for any \(t>0\) and the VE process (6 ), it holds that, \[\begin{align} s^\star_{Z_t}(x) = \frac{\int \Big( -\frac{x-y}{t}\Big)e^{-\frac{1}{2t}\|x-y\|_2^2}p^\star(\! \,\mathrm{d}y) }{\int e^{-\frac{1}{2t}\|x-y\|_2^2}p^\star(\! \,\mathrm{d}y)} \end{align}\]

Therefore, \[\begin{align} s^\star_{X_t}(c_tx) &= \frac{\int \Big( -\frac{c_t(x-y)}{\sigma_t^2}\Big)e^{-\frac{c_t^2}{2\sigma_t^2}\|x-y\|_2^2}p^\star(\! \,\mathrm{d}y) }{\int e^{-\frac{c_t^2}{2\sigma_t^2}\|x-y\|_2^2}p^\star(\! \,\mathrm{d}y)}\\ &= \frac{1}{c_t}s^\star_{Z_{h(t)}}(x) \\ &=: \frac{1}{c_t}s^\star_{h(t)}(x) \end{align}\] Here \(h(t)\mathrel{\vcenter{:}}=\frac{\sigma_t^2}{c_t^2} = \frac{1-e^{-2t}}{e^{-2t}} = e^{2t}-1\).

Proof of Theorem 2.

We denote \(\mathcal{E}\) as the event set of exact subspace recovery using \(n_0\) samples. From Lemma 3, we could take \(n_0 = O(M^2k\log n)\) and thus \(\mathbb{P}[\mathcal{E}^c] \lesssim Mn^{-10}\). Hence, for large \(n\) such that \(n_0 \leq 0.5 n\), we have the sample size for score estimation: \(N \geq 0.5n\). Therefore, for the score estimation error of VE process, we have, \[\begin{align} \mathbb{E}\big[ \| \widehat{s}_{t}(X)-s_{t}^\star(X)\|_2^2\big] &= \mathbb{E}\big[ \| \widehat{s}_{t}(X)-s_{t}^\star(X)\|_2^2 \cdot \mathbb{1}_{\mathcal{E}} \big] + \mathbb{E}\big[ \| \widehat{s}_{t}(X)-s_{t}^\star(X)\|_2^2 \cdot \mathbb{1}_{\mathcal{E}^c} \big] \notag \\ & \lesssim \frac{dM^3}{nt} \Big( 1+\frac{\sigma^{k\vee2}}{t^{{(k\vee2)}/2}}\Big) \cdot \big(\mathsf{poly}\log n +\log t\big) +\frac{d\log n}{t} \cdot Mn^{-10} + \sqrt{\mathbb{E}[\| s_t^\star(X)\|_2^4] \cdot Mn^{-10}} \notag\\ &\lesssim \frac{dM^3}{nt} \Big( 1+\frac{\sigma^{k\vee2}}{t^{{(k\vee2)}/2}}\Big) \cdot \big(\mathsf{poly}\log n +\log t\big) \label{eq:score32estimation32error32for32VE} \end{align}\tag{30}\] In the first inequality, we apply Theorem 1 and that \(\|\widehat{s}_t(x)\|^2 \lesssim \frac{R_t^2(i)}{t^2} + \frac{\log n}{t} \lesssim \frac{d\log n}{t}\); in the last inequality, we apply Lemma 11 in [15] for the moment bound of true scores.

Now we consider the sampling error results from score estimation error. Notice that 1-Wasserstein distance between target distribution and generated distribution \(\widehat{Y}_{T-\tau}\) using Algorithm 1 has the following control, \[\begin{align} \mathbb{E}\big[W_1(p^\star,p_{\widehat{Y}_{T-\tau}})\big]\lesssim \sqrt{d} \Big( \sqrt{\tau} + \sum_{j=0}^{L-1} \sqrt{\log \delta^{-1}\cdot \sigma_{T_{j+1}}^2 \int_{T_j}^{T_{j+1}} \int_{\mathbb{R}^d}\mathbb{E}\big[ \| \widehat{s}_{X_t}(x)-s_{X_t}^\star(x)\|_2^2\big] p_{X_t}(x) \,\mathrm{d}x \,\mathrm{d}t} + \delta + e^{-T} \Big) \label{eq:W95132convergence32rate} \end{align}\tag{31}\] for certain \(0<T_0=\tau < T_1< \cdots <T_L = T\) and any \(\delta>0\). This bound can be found as (8) in [18], which cites Lemma D.7 in [9] as the proof, and also can be found as Lemma B.2 in [17]. Furthermore, the score estimator we define in (14 ) satisfies, \[\begin{align} \widehat{w}_t(i,x) \neq 0 \Longrightarrow x \in \mathcal{G}_t(i) \Longrightarrow \|\widehat{s}_t(i,x)\| \lesssim \sqrt{\frac{\log n}{t}} \end{align}\] and hence, \(\|\widehat{s}_t(x)\|\lesssim \sqrt{\frac{\log n}{t}}\), with, \[\begin{align} \widehat{s}_{X_t}(x) = \frac{1}{c_t } \widehat{s}_{h(t)}(\frac{x}{c_t}) \lesssim \frac{1}{c_t} \sqrt{\frac{\log n}{\sigma_t^2/c_t^2}} = \sqrt{\frac{\log n}{\sigma_t^2}}. \end{align}\] Therefore, the condition for (31 ) is satisfied in our setting. We first derive a bound for \(\int_{\mathbb{R}^d}\mathbb{E}\big[ \| \widehat{s}_{X_t}(x)-s^\star_{X_t}(x)\|_2^2\big] p_{X_t}(x)\,\mathrm{d}x\) using the result in Theorem 1,

\[\begin{align} \int_{\mathbb{R}^d}\mathbb{E}\big[ \| \widehat{s}_{X_t}(x)-s^\star_{X_t}(x)\|_2^2\big] p_{X_t}(x)\,\mathrm{d}x &= \frac{1}{c_t^2}\int_{\mathbb{R}^d}\mathbb{E}\big[ \| \widehat{s}_{h(t)}(\frac{x}{c_t})-s^\star_{h(t)}(\frac{x}{c_t})\|_2^2\big] p_{X_t}(x)\,\mathrm{d}x \quad (\text{Equation (\ref{eq:score32between32VE32process32and32VP})})\\ &= \int \frac{1}{c_t^2}(c_t^2)^{d/2}\mathbb{E} \big[\| \widehat{s}_{h(t)}(y)-s^\star_{h(t)}(y)\|_2^2\big] p_{X_t}(c_t y)\,\mathrm{d}y \\ &= \frac{1}{c_t^2} \int \mathbb{E} \big[\| \widehat{s}_{h(t)}(y)-s^\star_{h(t)}(y)\|_2^2\big] p_{Z_{h(t)}}(y)\,\mathrm{d}y \\ &\lesssim \frac{1}{c_t^2} \frac{dM^3}{n} \Big( \frac{1}{h(t)}+\frac{\sigma^{k\vee2}}{h(t)^{{(k\vee2)}/2+1}}\Big) \big(\mathsf{poly}\log n+\log h(t)\big)\quad (\text{Equation (\ref{eq:score32estimation32error32for32VE})}). \end{align}\] Observe that, \[\begin{align} h'(t) = 2e^{2t} = 2/c_t^2 \end{align}\] and thus, \[\begin{align} \int_{T_j}^{T_{j+1}} \int_{\mathbb{R}^d}\mathbb{E}\big[ \| \widehat{s}_{X_t}(x)-s^\star_{X_t}(x)\|_2^2\big] p_{X_t}(x)\,\mathrm{d}x \,\mathrm{d}t &\lesssim \frac{dM^3}{n}\Big(\int_{T_j}^{T_{j+1}} \Big( \frac{1}{h(t)}+\frac{\sigma^{k\vee2}}{h(t)^{(k\vee2)/2+1}} \Big) h'(t) \,\mathrm{d}t\Big) \big(\mathsf{poly}\log n+\log h(t)\big)\\ &\lesssim \frac{dM^3}{n} \Big( \log \frac{h(T_{j+1})}{h(T_j)} + \frac{2\sigma^{k\vee2}}{k\vee2}\frac{1}{h(T_j)^{(k\vee2)/2}} \Big) \big(\mathsf{poly}\log n+T\big) . \end{align}\]

We then take the synthesized discretization as, \[\begin{align} T_{j+1} = 2T_j, \quad T \asymp \log n, \quad \tau\asymp n^{-\gamma} \end{align}\] for some \(\gamma>0\) that will be determined later. Then we can easily check that \(L\asymp \log n\) and, \[\begin{align} &\sum_{j=0}^{L-1} \sqrt{(1-e^{-4T_j})\cdot \int_{T_j}^{T_{j+1}} \int_{\mathbb{R}^d}\mathbb{E}\big[ \| \widehat{s}_{X_t}(x)-s^\star_{X_t}(x)\|_2^2\big] p_{X_t}(x) \,\mathrm{d}x \,\mathrm{d}t} \\ &\lesssim \frac{\sqrt{d}M^{3/2} \mathsf{poly}\log n}{\sqrt{n}} \Big( \sum_{j=0}^{L-1} \sqrt{(1-e^{-4T_j})\frac{2\sigma^{k\vee2}}{k\vee2}\frac{1}{(e^{2T_k}-1)^{(k\vee2)/2}}} + \mathsf{poly}\log n\Big)\\ &\lesssim \frac{\sqrt{d}M^{3/2}\mathsf{poly}\log n}{\sqrt{n}} \cdot \Big(\sum_{j=0}^{L-1}\sqrt{\frac{4\sigma^kT_j}{2^{k/2}T_j^{(k\vee2)/2}}} +1 \Big) \\ &\lesssim \begin{cases} \frac{\sqrt{d}M^{3/2}\mathsf{poly}\log n}{\sqrt{n}}\tau^{-\frac{k}{4}+\frac{1}{2}}, \quad &k\geq 2 \\ \frac{\sqrt{d}M^{3/2}\mathsf{poly}\log n}{\sqrt{n}}, & k=1 \end{cases}. \end{align}\]

Further take \(\delta = n^{-1}\) and take \(\tau = n^{-2/k}\), then it holds that, \[\begin{align} W_1(p^\star,p_{\widehat{Y}_{T-\tau}}) \lesssim \frac{dM^{3/2} \mathsf{poly}\log n}{n^{1/(k\vee 2)}}. \end{align}\]

8 Proof of Lemmas↩︎

As long as we could ensure that with high probability the data segmentation with \(n_0\) points are correct and each class contains at least \(k+1\) points, then we could recover the exact linear subspaces \(V_j\) and \(c_j(X) = \mathbb{1}_{\{ X \in V_j\}}\) almost surely for \(X \sim p^\star\). This gives the following Lemma 3.

We will further implement an algorithm that proves Lemma 3 and discuss its complexity in Appendix 8.1. This lemma is just for theoretical guarantee for our further score estimation under exact subspaces recovery.

With exact subspace recovery, when constructing the score component in (10 ), we only need to estimate \(s_t^\mathsf{low}\) and the following lemma gives a time-dependent rate that only relies on intrinsic dimension \(k_i\), using certain low-dimensional estimator (17 ).

Here we adopt the estimator from [15] with a further cut-off. This achieves a similar rate w.r.t \(t\) compared with [14], [15]. The proof is provided in Appendix 8.3.

8.1 Proof of Lemma 3.↩︎

Basically, subspace clustering for noise-free model aims to solve the following optimization problem iteratively, \[\begin{align} \min_{A} \sum_{i=1}^{n_0} \| X^{(i)}-AA^{\top}X^{(i)}\|_0 \end{align}\] Here the \(L_0\) norm is defined as, \[\begin{align} \| x\|_0 = \begin{cases} 0 , \quad x=0\\ 1, \quad \text{else} \end{cases} \end{align}\] As a non-convex and non-smooth optimization problem, it is basically an NP-hard problem.

With known upper bound of both intrinsic dimension \(k\) and the number of subspaces \(M\), the following algorithm is guaranteed to recover exact subspace recovery with high probability and under Assumption 1,

Define the event sets, \[\begin{align} \mathcal{E}_0 &\mathrel{\vcenter{:}}=\Big\{ \exists \text{ at least } k+1 \text{ samples on each subspace}\Big\} \\ \mathcal{E}_1(p) & \mathrel{\vcenter{:}}=\Big\{ x \in \text{span} \{ y_1.\cdots,y_p\}, V_{j} \nsubseteq\text{ span} \{ y_1,\cdots,y_p\}, \forall j\in [M] \Big\}, \quad \text{for }p\leq \widetilde{k} \end{align}\] To ensure that Algorithm 3 works with probability larger than \(1-Mn^{-10}\), we only need, \[\begin{align} \mathbb{P}[\mathcal{E}_0^c] &\lesssim Mn^{-10}\tag{32} \\ \mathbb{P}_{x,y_1,\cdots y_p \overset{\text{i.i.d}}{\sim} \mu^*}[\mathcal{E}_1(p)] &= 0, \quad \forall p\leq \widetilde{k} \tag{33} \end{align}\] Since under \(\mathcal{E}_1(p)^c\), samples from other subspaces will be excluded via finding the smallest \(p\), we could recover the exact subspace in each iteration.

Proof of (33 ).

We first conditioned on \(y_1,\cdots,y_p\) and \(V_j \nsubseteq\text{ span} \{ y_1,\cdots,y_p\}\) for all \(j\in [M]\), then, \[\begin{align} \mathbb{P} [X\in \text{span} \{ y_1.\cdots,y_p\}] &= \sum_{j=1}^M \int_{V_j \cap \text{span} \{ y_1.\cdots,y_p\} } p_j^\star(\!\,\mathrm{d}x) \\ &= 0 \quad (\text{Assumption \ref{assume:multi-modal}}) \end{align}\] Then we integrate this over \(y_1,\cdots,y_p \sim \mu^*\) and get \(\mathbb{P}[\mathcal{E}_1(p)] = 0\).

Proof of (32 ).

Basically, under Assumption 1, denote \(N_i(n_0), \forall i \in [M]\) as the random variable of the number of samples on \(V_i\) with sample size \(n_0\). Then, we could view \(N_i(n_0)\) as the sum of Bernoulli r.v with \(p=p_i^\star(V_i)\geq \frac{1}{c_{\mu}M}\) and thus, \[\begin{align} N_i(n_0) \geq \frac{n_0}{2c_{\mu}M}, \quad \text{with probability }\geq 1-2e^{-\frac{n_0}{2c_{\mu}^2M^2}} \end{align}\] here we apply Hoeffding inequality just like the proof of Claim 1. Finally, we could apply union bound and take, \[\begin{align} n_0 = O(c_{p}^2M^2(k+1) \log n) . \end{align}\]

The iteration complexity for Algorithm 3 is bounded by \(\begin{pmatrix} n_0\\ \widetilde{k}+1 \end{pmatrix}\cdot M\).

8.2 Proof of Lemma 1.↩︎

The following lemma helps prove this result,

We leave the proof of this lemma in the end of this section, and we first prove Lemma 1 as follows using Lemma 5. For \(x\in \mathcal{G}_t(i)\) \[\begin{align} \mathbb{E}\big[\big(w_t(i,x)-\widehat{w}_t(i,x) \big)^2\big] &= \mathbb{E} \Big[\big(\frac{q_t(i,x)}{p_t(x)} - \frac{\widehat{q}_t(i,x)}{\widehat{p}_t(x)}\big)^2 \Big] \\ &\lesssim \mathbb{E} \Big[ \big(\frac{q_t(i,x)-\widehat{q}_t(i,x)}{p_t(x)}\big)^2\Big] + \mathbb{E} \Big[ \big( \frac{\widehat{q}_t(i,x)}{\widehat{p}_t(x)}\big)^2 \cdot \frac{(\widehat{p}_t(x)-p_t(x))^2}{p_t^2(x)}\Big] \end{align}\]

• For the first term, \[\begin{align} \mathbb{E} \Big[ \big( \frac{q_t(i,x)-\widehat{q}_t(i,x)}{p_t(x)}\big)^2\Big] &\leq \frac{1}{p_t^2(x)} \cdot \frac{1}{N}(2\pi t)^{-d/2} e^{-\frac{1}{2t} \| x-\mathsf{proj}_i(x)\|_2^2} q_t(i,x). \quad (\text{Lemma \ref{lem:MSE32for32p95t32and32p95t40i41}}) \end{align}\]

• For the second term, \[\begin{align} \mathbb{E} \Big[ \big( \frac{\widehat{q}_t(i,x)}{\widehat{p}_t(x)}\big)^2 \cdot \frac{(\widehat{p}_t(x)-p_t(x))^2}{p_t^2(x)}\Big] &\leq \mathbb{E} \Big[ \big( \frac{\widehat{p}_t(x)-p_t(x)}{p_t(x)}\big)^2\Big] \\ &\leq \frac{1}{p_t^2(x)} \cdot \frac{1}{(2\pi t)^{d/2}N} \Big(\sum_{i=1}^M e^{-\frac{1}{2t}\|x-\mathsf{proj}_i(x)\|_2^2} \cdot q_t(i,x) \Big) \quad (\text{Lemma \ref{lem:MSE32for32p95t32and32p95t40i41}}) \end{align}\]

Therefore, it holds that, \[\begin{align} \mathbb{E}[\big(w_t(i,x)-\widehat{w}_t(i,x) \big)^2] \lesssim \frac{1}{p_t^2(x)} \cdot \frac{1}{(2\pi t)^{d/2}N} \Big(\sum_{i=1}^M e^{-\frac{1}{2t}\|x-\mathsf{proj}_i(x)\|_2^2} \cdot q_t(i,x) \Big) . \end{align}\]

Then we prove Lemma 5, and the proof idea is quite similar to the proof of Lemma 3 in [15].

Proof of Lemma 5.

• First prove that both are unbiased estimators, \[\begin{align} \mathbb{E}[\widehat{p}_t(x)] &= \mathbb{E}[\frac{1}{N}\sum_{j=1}^N \varphi_t(x-X^{(j)};d)] = \mathbb{E}_{X\sim p^\star}[\varphi_t(x-X;d)] = p_t(x) \\ \mathbb{E}[\widehat{q}_t(i,x)] &= \mathbb{E}_{X\sim p^\star} [\varphi_t(x-X;d) \mathbb{1}_{\{ c(X) = i\}}] = \mathbb{E}_{X\sim p^\star} [\varphi_t(x-X;d) \mathbb{1}_{\{ X \in V_i\}}] = \int_{V_i} \varphi_t(x-y;d) p_i^\star(\!\,\mathrm{d}y) \end{align}\] In the last equation, we apply Assumption 1.

• Then derive the mean squared error (MSE) for both estimators, i.e., both variance. To prove (34 ), observe that, \[\begin{align} \mathsf{Var}(\widehat{p}_t(x)) &= \frac{1}{N} \mathsf{Var}_{X \sim p^\star} (\varphi_t(x-X;d)) \leq \frac{1}{N} \mathbb{E}_{X\sim p^\star}[\varphi_t^2(x-X;d)] = \frac{1}{N}\sum_{i=1}^M \int_{V_i} \varphi^2_t(x-y;d) p_i^\star(\!\,\mathrm{d}y) . \end{align}\] Plug in the structure of \(V_i\): \[\begin{align} \int_{V_i} \varphi_t^2(x-y;d)p_i^\star(\!\,\mathrm{d}y) &= \int_{V_i} (2\pi t)^{-d} e^{-\frac{1}{t}\|x-y\|_2^2} p_i^\star(\!\,\mathrm{d}y) \\ &= (2\pi t)^{-d} e^{-\frac{1}{t} \| x-\mathsf{proj}_i(x)\|_2^2} \int_{V_i} e^{-\frac{1}{t} \| y-\mathsf{proj}_i(x)\|_2^2} p_i^\star(\!\,\mathrm{d}y) \end{align}\] Notice that, \[\begin{align} \int_{V_i} e^{-\frac{1}{t} \| y-\mathsf{proj}_i(x)\|_2^2} p_i^\star(\!\,\mathrm{d}y) \leq \int_{V_i} e^{-\frac{1}{2t} \| y-\mathsf{proj}_i(x)\|_2^2} p_i^\star(\!\,\mathrm{d}y) . \end{align}\] Hence, \[\begin{align} \int_{V_i} \varphi_t^2(x-y;d)p_i^\star(\!\,\mathrm{d}y) \leq (2\pi t)^{-d/2} e^{-\frac{1}{2t} \| x-\mathsf{proj}_i(x)\|_2^2} \cdot q_t(i,x) . \end{align}\] Therefore, \[\begin{align} \mathbb{E}\big[\big(p_t(x)-\widehat{p}_t(x)\big)^2\big] &\leq \frac{1}{N}\sum_{i=1}^M \int_{V_i} \varphi_t^2(x-y;d)p_i^\star(\!\,\mathrm{d}y) \\ &\leq \frac{1}{(2\pi t)^{d/2}N} \sum_{i=1}^M e^{-\frac{1}{2t} \| x-\mathsf{proj}_i(x)\|_2^2} \cdot q_t(i,x). \end{align}\]

• To prove (35 ), similarly, \[\begin{align} \mathbb{E}\big[\big( q_t(i,x) - \widehat{q}_t(i,x)\big)^2\big] \leq \frac{1}{N} \mathbb{E}_{X\sim p^\star} [\varphi^2_t(x-X;d)\mathbb{1}_{\{ X \in V_i\}} ] = \frac{1}{N} \sum_{j=1}^M \int_{V_j} \varphi_t^2(x-y;d) \mathbb{1}_{\{ y \in V_i\}} p_j^\star(\!\,\mathrm{d}y) \end{align}\] Under Assumption 1, it holds that, \[\begin{align} \mathbb{E}\big[\big( q_t(i,x) - \widehat{q}_t(i,x)\big)^2\big] \leq \frac{1}{N} \int_{V_i} \varphi_t^2(x-y;d) p_i^\star(\!\,\mathrm{d}y) &\leq \frac{1}{N}(2\pi t)^{-d} e^{-\frac{1}{t} \| x-\mathsf{proj}_i(x)\|_2^2} \int_{V_i} e^{-\frac{1}{t} \| y-\mathsf{proj}_i(x)\|_2^2} p_i^\star(\!\,\mathrm{d}y) \\ & \leq \frac{1}{N}(2\pi t)^{-d} e^{-\frac{1}{t} \| x-\mathsf{proj}_i(x)\|_2^2} \int_{V_i} e^{-\frac{1}{2t} \| y-\mathsf{proj}_i(x)\|_2^2} p_i^\star(\!\,\mathrm{d}y) \\ &\leq \frac{1}{N}(2\pi t)^{-d/2} e^{-\frac{1}{2t} \| x-\mathsf{proj}_i(x)\|_2^2} q_t(i,x). \end{align}\]

8.3 Proof of Lemma 4.↩︎

In the proof of this lemma, the target distribution \(\nu\) is a \(\sigma\)-subgaussian distribution in \(\mathbb{R}^k\), and we denote \(p_t\) as the density of \(\nu \ast \mathcal{N}(0,tI_k)\) and \(s_t(\cdot) \mathrel{\vcenter{:}}=\nabla \log p_t(x)\) as its score. \(\{X^{(i)}\}_{i=1}^N \overset{\text{i.i.d}}{\sim} \nu\) are N samples for score estimation.

We take the estimator based on (17 ) as, \[\begin{align} \widehat{s}_t(x) = \mathsf{clip}_R\Big( \frac{\nabla \widehat{g}_t(x)}{\widehat{g}_t(x)} \psi\Big(\widehat{g}_t(x);\frac{\log N}{N(2\pi t)^{k/2}}\Big)\Big) \quad \text{with } R \mathrel{\vcenter{:}}=\sqrt{\frac{2}{t} \log N}. \label{eq:low-dim32score32estimator} \end{align}\tag{36}\] Here the clip operator is defined as, \[\begin{align} \mathsf{clip}_r(z) \mathrel{\vcenter{:}}= \begin{cases} z, \quad &\|z\|_2 \leq r \\ \mathsf{proj}_{B_r(0)}(z) ,\quad &\text{else} \end{cases} \end{align}\] where \(\mathsf{proj}_{B_r(0)}(z) = r \cdot \frac{z}{\|z\|_2}\) means the projection of \(z\) on that ball; \(\psi(x;\eta) \mathrel{\vcenter{:}}=\mathbb{1}_{\{x\geq \eta \}}\) is the hard thresholding function; \(\widehat{g}_t(x) \mathrel{\vcenter{:}}=\frac{1}{N}\sum_{i=1}^N \varphi_t(X^{(i)}-x;k)\) is the kernel based density estimator. As a remark, the estimator before clipping is almost the same as that in [15], and we apply hard-thresholding here for simplicity since we use DDPM sampling procedure.

Then we show that the score estimator defined in (36 ) satisfies the conditions in Lemma 4.

In [15], they defined \[\begin{align} \mathcal{F}_t \mathrel{\vcenter{:}}=\Big\{ x: p_t(x)\geq \frac{c_{\eta} \log N}{N(2\pi t)^{k/2}}\Big\} \end{align}\] and it holds that, \[\begin{align} \int_{\mathcal{F}_t^c} p_t(x) \,\mathrm{d}x &\lesssim \Big( \frac{16}{\pi}\Big)^{k/2} \frac{1}{N} \Big( 1 + \frac{\sigma^k }{ t^{k/2}}\Big) (\log N)^{k/2+1} \tag{37} \\ \int_{\mathcal{F}_t^c} \|s_t(x)\|_2^2\cdot p_t(x) \,\mathrm{d}x &\lesssim \Big( \frac{16}{\pi}\Big)^{k/2} \frac{1}{N} \Big( \frac{1}{t} + \frac{\sigma^k }{ t^{k/2+1}}\Big) (\log N)^{k/2+2} \tag{38} \end{align}\] from Lemma 8 in [15]. Notice that from Lemma 4 in [15], it has been proven that, \[\begin{align} \| s_t(x)\|_2^2 \leq \frac{2}{t} \log \frac{1}{(2\pi t)^{k/2} p_t(x)} \end{align}\] Hence, for \(x\in \mathcal{F}_t\), one has, \[\begin{align} \| s_t(x)\|_2^2 \leq \frac{2}{t} \log \frac{N}{c_{\eta} \log N} \leq \frac{2}{t} \log N . \label{eq:score32bdd32on32F95t} \end{align}\tag{39}\] Here we use the constant selection of \(c_{\eta} \geq 2\) as in [15].

That’s the reason why we use the clip operator in (36 ). Therefore, with the estimator constructed in (36 ), it automatically satisfies the bounded condition and for its \(L^2\) error, \[\begin{align} \int_{\mathbb{R}^k} \mathbb{E}\big[\big\| \widehat{s}_t(x) - s_t(x)\big\|_2^2 \big] p_t(x) \,\mathrm{d}x&\leq \int_{\mathcal{F}_t} \mathbb{E}\big[\big\| \widehat{s}_t(x) - s_t(x)\big\|_2^2 \big] p_t(x) \,\mathrm{d}x \\ &+ \int_{\mathcal{F}_t^c} \mathbb{E}\big[\big\| \widehat{s}_t(x) - s_t(x)\big\|_2^2 \big] p_t(x) \,\mathrm{d}x \\ &\eqqcolon\chi_1 + \chi_2 \end{align}\]

• For \(\chi_1\). Due to (39 ), for \(x\in \mathcal{F}_t\), \(s_t(x) = \mathsf{clip}_R(s_t(x))\). What’s more, \(\mathsf{clip}_R(\cdot)\) is \(1\)-Lip continous, and thus, \[\begin{align} \chi_1 &\leq \int_{\mathcal{F}_t}\mathbb{E}\big[\big\|\frac{\nabla \widehat{g}_t(x)}{\widehat{g}_t(x)} \psi\Big(\widehat{g}_t(x);\frac{\log N}{N(2\pi t)^{k/2}}\Big) -s_t(x)\big\|_2^2 \big] \cdot p_t(x) \,\mathrm{d}x \\ & \leq \int_{\mathbb{R}^k}\mathbb{E}\big[\big\|\frac{\nabla \widehat{g}_t(x)}{\widehat{g}_t(x)} \psi\Big(\widehat{g}_t(x);\frac{\log N}{N(2\pi t)^{k/2}}\Big) -s_t(x)\big\|_2^2 \big] \cdot p_t(x) \,\mathrm{d}x \\ &\leq \frac{(4/\sqrt{\pi})^k}{N} \Big( \frac{1}{t}+\frac{\sigma^k}{t^{k/2+1}}\Big) (\log N)^{k/2+1} \end{align}\] The last inequality results from Proposition 1 in [15], which holds for hard thresholding \(\psi\) by simple modification.

• For \(\chi_2\). Notice that, \[\begin{align} \chi_2 &\mathrel{\vcenter{:}}=\int_{\mathcal{F}_t^c}\mathbb{E}\big[\big\| \widehat{s}_t(x) - s_t(x)\big\|_2^2 \big] p_t(x) \,\mathrm{d}x \\ &\leq \int_{\mathcal{F}_t^c} \Big( R^2 + \|s_t(x)\|_2^2\Big) p_t(x) \,\mathrm{d}x \\ &\leq R^2 \int_{\mathcal{F}_t^c} p_t(x) \,\mathrm{d}x + \int_{\mathcal{F}_t^c}\|s_t(x)\|_2^2 \cdot p_t(x) \,\mathrm{d}x \\ &\lesssim \Big( \frac{16}{\pi}\Big)^{k/2} \frac{1}{N} \Big( \frac{1}{t} + \frac{\sigma^k }{ t^{k/2+1}}\Big) (\log N)^{k/2+2} \quad (\text{Using }(\ref{eq:prob32bdd32on32F95t94c}) + (\ref{eq:score32bdd32on32F95t94c})) \end{align}\]

Therefore, it holds that, \[\begin{align} \int_{\mathbb{R}^k} \mathbb{E}\big[\big\| \widehat{s}_t(x) - s_t(x)\big\|_2^2 \big] p_t(x) \,\mathrm{d}x \lesssim \Big( \frac{4}{\sqrt{\pi}}\Big)^{k} \frac{1}{N} \Big( \frac{1}{t} + \frac{\sigma^k }{ t^{k/2+1}}\Big) (\log N)^{k/2+2} . \end{align}\]

8.4 Proof of Lemma 2↩︎

Observe that, conditioned on \(N_i\), we could obtain \(N_i\) i.i.d samples from normalized distribution \(p_i^\star\) to estimate its score function. We only need to prove (25 ). Notice that, \[\begin{align} q_t(i,x) = \int_{V_i} \varphi_t(x-y;d) p_i^\star(\!\,\mathrm{d}y) &= (2\pi t)^{-d/2} e^{-\frac{1}{2t}\| x-\mathsf{proj}_i(x)\|_2^2} \int_{V_i} e^{-\frac{1}{2t}\| \mathsf{proj}_i(x)-y\|_2^2} p_i^\star( \!\,\mathrm{d}y) \\ &= (2\pi t)^{-d/2} e^{-\frac{1}{2t}\| x-\mathsf{proj}_i(x)\|_2^2} p_i^\star(V_i) \int_{\mathbb{R}^k} e^{-\frac{1}{2t}\| z - A_i^{\top}x\|_2^2} p_i^{\mathsf{low}}(\! \,\mathrm{d}z) \\ &= (2\pi t)^{-(d-k_i)/2} e^{-\frac{1}{2t}\| x-\mathsf{proj}_i(x)\|_2^2} p_i^\star(V_i) \cdot p_t^{\mathsf{low}}(i,A_i^{\top}x) \end{align}\] here \(p_t^{\mathsf{low}}(i,\cdot)\) denotes the density function of \(p_i^{\mathsf{low}} \ast \mathcal{N}(0,tI_{k_i})\) in \(\mathbb{R}^{k_i}\). Hence, \[\begin{align} &\int_{\mathbb{R}^d} \mathbb{E}\big[\big\| \widehat{s}_t(i,x) - s_t(i,x)\big\|_2^2 \big| N_i \big] q_t(i,x) \,\mathrm{d}x \\ &\leq p_i^\star(V_i)\int_{\mathbb{R}^d} \mathbb{E}\big[\big\| \widehat{s}_t^{\mathsf{low}}(i,A_i^{\top}x) - s_t^{\mathsf{low}}(i,A_i^{\top}x)\big\|_2^2\big| N_i \big] (2\pi t)^{-(d-k_i)/2} e^{-\frac{1}{2t}\| x-A_iA_i^{\top}x\|_2^2} \cdot p_t^{\mathsf{low}}(i,A_i^{\top}x) \,\mathrm{d}x \end{align}\] Apply the same linear transform as in (29 ), \[\begin{align} &\int_{\mathbb{R}^d} \mathbb{E}\big[\big\| \widehat{s}_t(i,x) - s_t(i,x)\big\|_2^2 \big| N_i \big] q_t(i,x) \,\mathrm{d}x \\ &\leq p_i^{\star}(V_i) \int_{\mathbb{R}^d} \mathbb{E}\big[\big\| \widehat{s}_t^{\mathsf{low}}(i,z_{1:{k_i}}) - s_t^{\mathsf{low}}(i,z_{1:{k_i}})\big\|_2^2 \big| N_i \big] (2\pi t)^{-(d-k_i)/2} e^{-\frac{1}{2t}\| z_{k_i+1:d}\|_2^2} \cdot p_t^{\mathsf{low}}(i,z_{1:k_i}) \,\mathrm{d}z\\ &= p^\star(V_i) \int_{\mathbb{R}^{k_i}} \mathbb{E}\big[\big\| \widehat{s}_t^{\mathsf{low}}(i,z_{1:k_i}) - s_t^{\mathsf{low}}(i,z_{1:k_i})\big\|_2^2 \big| N_i \big] \cdot p_t^{\mathsf{low}}(i,z_{1:k_i}) \,\mathrm{d}z_{1:k_i} \quad (\text{Tonelli's Theorem})\\ &\lesssim p_i^{\star}(V_i) \cdot \frac{(4/\sqrt{\pi})^{k_i}}{N_i}\Big( \frac{1}{t}+\frac{\sigma^{k_i}}{t^{k_i/2+1}}\Big) \cdot \mathsf{poly}\log N \quad (\text{Lemma \ref{lem:low-dim32score32estimator}}) \end{align}\] Finally, \[\begin{align} \int_{\mathbb{R}^d} \mathbb{E}\big[\big\| \widehat{s}_t(i,x) - s_t(i,x)\big\|_2^2 \mathbb{1}_{\{N_i \geq n_i \}}\big] q_t(i,x) \,\mathrm{d}x &= \int_{\mathbb{R}^d} \mathbb{E}\Big[\mathbb{E}\big[\big\| \widehat{s}_t(i,x) - s_t(i,x)\big\|_2^2 \big| N_i \big] \cdot \mathbb{1}_{\{N_i \geq n_i \}}\Big] q_t(i,x) \,\mathrm{d}x\\ & = \mathbb{E}\Big[ \Big( \int_{\mathbb{R}^d} \mathbb{E}\big[\big\| \widehat{s}_t(i,x) - s_t(i,x)\big\|_2^2 \big| N_i \big] \cdot q_t(i,x) \,\mathrm{d}x\Big) \cdot \boldsymbol{1}_{\{N_i \geq n_i \}}\Big] \\ &\lesssim p_i^\star(V_i) \cdot \frac{(4/\sqrt{\pi})^{k_i}}{n_i}\Big( \frac{1}{t}+\frac{\sigma^{k_i}}{t^{k_i/2+1}}\Big) \cdot \mathsf{poly}\log N . \end{align}\]

9 Proof of Claims↩︎

9.1 Proof of Claim 1↩︎

For any \(i\in[M]\), we first bound the probability \(\mathbb{P}[N_i < \frac{N}{2c_{p}M}]\). Note that \(N_i = \sum_{j=1}^{N} \mathbb{1}_{\{ X^{(j)} \in V_i\}}\) is basically the sum of \(N\) i.i.d Bernoulli r.vs with parameter \(p^\star(V_i) \geq \frac{1}{c_{p}M}\) (Assumption 1). Applying Hoeffding’s inequality, we have, \[\begin{align} \mathbb{P}\Big[ N_i < \frac{N}{2c_{p}M} \Big] & = \mathbb{P}\Big[ N_i - N\cdot p^\star(V_i) < -\Big(N\cdot p^\star(V_i) - \frac{N}{2c_{p}M}\Big) \Big] \\ &\leq \exp \Big( -2N\cdot (\frac{1}{2c_{p}M})^2 \Big) = \exp \Big( -\frac{N}{2c_{p}^2M^2}\Big). \end{align}\] Then applying union bound over \(i\in[M]\), it holds that, \[\begin{align} \mathbb{P}\big[\mathcal{A}^c\big] = \mathbb{P}\Big[\exists i \in [M], N_i < \frac{N}{2c_{p}M}\Big] \leq Me^{-\frac{N}{2c_{p}^2M^2}}. \end{align}\]

9.2 Proof of Claim 2↩︎

For the target distribution \(p^\star\), it is supported on \(\cup_{i=1}^M V_i\). Hence for any \(\theta \in \mathbb{R}^d\) and \(\| \theta\|_2 =1\), denote r.v \(X \sim p^\star\) \[\begin{align} \mathbb{E}\Big[ \exp \Big( \frac{(X^{\top}\theta)^2}{\sigma^2}\Big)\Big] &= \sum_{i=1}^M \int_{V_i} \exp \Big( \frac{(x^{\top}\theta)^2}{\sigma^2} \Big) p_i^\star(\!\,\mathrm{d}x) \\ & = \sum_{i=1}^M p_i^\star(V_i) \int_{\mathbb{R}^{k_i}}\exp \Big( \frac{(z^{\top}A_i^{\top}\theta)^2}{\sigma^2} \Big) p_i^{\mathsf{low}}(\!\,\mathrm{d}z) \\ & \leq \sum_{i=1}^M p_i^\star(V_i) \int_{\mathbb{R}^{k_i}}\exp \Big( \frac{(z^{\top}A_i^{\top}\theta)^2}{\sigma_i^2} \Big) p_i^{\mathsf{low}}(\!\,\mathrm{d}z) \\ &\leq \sum_{i=1}^M p_i^\star(V_i) \cdot 2 = 2. \end{align}\] Here we use \(\| A_i^{\top}\theta \| \leq 1\) and Assumption 2. This shows that, \(p^\star\) is a sub-gaussian distribution in \(\mathbb{R}^d\) with parameter \(\sigma\). Therefore, it is straightforward that \(p^\star \ast \mathcal{N}(0,tI_d)\) is \(c\sqrt{\sigma^2 +t}\) sub-gaussian, for some absolute constant \(c>0\).

• For \(\int_{\mathcal{B}_t^c} p_t(x)\!\,\mathrm{d}x\), it holds that, \[\begin{align} \int_{\mathcal{B}_t^c} p_t(x)\!\,\mathrm{d}x &\leq \sum_{i=1}^M \mathbb{P}_{Z_t \sim p_t} \big[\|A_i^{\top}Z_t\|_2>B_t\big] \\ & \leq \sum_{i=1}^M \mathbb{P}_{Z_t \sim p_t} \big[\|A_i^{\top}Z_t\|_{\infty}>B_t/\sqrt{k_i}\big] \lesssim Mk \exp\big( \frac{cB_t^2}{k_i(\sigma^2+t)}\big) \lesssim \frac{Mk}{N^4}\quad (\text{Lemma \ref{lem:tail32bdd32for32subgaussian}}) \end{align}\] Recall the definition of \(\mathcal{B}_t\) in (28 ), and we take large enough constant \(C_B>0\) to ensure \(\exp\big( \frac{cB_t^2}{k_i(\sigma^2+t)}\big)\lesssim N^{-4}\).

• For \(\int_{\mathcal{B}_t^c} \|x\|_2^2\cdot p_t(x)\!\,\mathrm{d}x\), it holds that, \[\begin{align} \int_{\mathcal{B}_t^c} \|x\|_2^2\cdot p_t(x)\!\,\mathrm{d}x &= \mathbb{E}_{Z_t \sim p_t} [\| Z_t \|_2^2 \mathbb{1}_{\{ Z_t\notin \mathcal{B}_t\}}] \\ & \leq \sqrt{\mathbb{E}_{Z_t \sim p_t}[\|Z_t\|_2^4] \cdot \int_{\mathcal{B}_t^c} p_t(x) \!\,\mathrm{d}x} \quad (\text{C-S Ineq})\\ & \lesssim \sqrt{d^2(\sigma^2+t)^2 \frac{Mk}{N^4}}\\ &\lesssim \frac{d\sqrt{M}}{N^2} (\sigma^2+t). \end{align}\]

10 Auxiliary Lemmas↩︎

References↩︎