Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions Jingda Wu 1 Changxiao Cai 1
arXiv:2605.30153v1 [stat.ML] 28 May 2026
Abstract
1. Introduction Score-based diffusion models (Sohl-Dickstein et al., 2015; Song & Ermon, 2019) have achieved state-of-the-art performance across a wide range of generative modeling applications, including image and video generation (Ho et al., 2020; Ho et al., 2022), signal processing (Song et al., 2021), and language modeling (Austin et al., 2021; Nie et al., 2025). 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.
Score-based diffusion models have demonstrated remarkable empirical success in learning highdimensional distributions, particularly those exhibiting low-dimensional and multi-modal structures. However, theoretical understanding of their statistical efficiency remains limited. Existing theories typically rely on strong regularity assumptions, such as uniformly bounded densities or globally smooth score functions, which fail to capture such intrinsic structures. In this work, we study the sample complexity of diffusion models for learning distributions supported on a union of low-dimensional subspaces. Assuming that the data distribution within each subspace is subgaussian, we show that diffusion models require e −k∨2 ) (up to some logat most the order of O(ε arithmic factor) samples to achieve ε sampling error in 1-Wasserstein distance, where k is the intrinsic dimension. This near-optimal convergence rate depends only on the intrinsic dimension and significantly improves upon prior theoretical guarantees that suffer from the curse of dimensionality. Notably, our analysis applies to a broad collection of distributions without imposing smoothness, bounded-density, or log-concavity assumptions. Overall, our results show that diffusion models can statistically adapt to intrinsic low-dimensional structure while naturally accommodating multi-modal data, offering a rigorous theoretical justification for their success in complex high-dimensional learning tasks.
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. Leveraging intrinsic structures. Motivated by this, a growing body of statistical theory has been developed to understand the sample complexity of diffusion models (Shah et al., 2023; Oko et al., 2023; Chen et al., 2024; Cole & Lu, 2024; Li et al., 2024; Dou et al., 2024). For a broad class of d-dimensional distributions with βHölder smooth densities (without assuming smooth scores or log-concave/uniformly-bounded densities), state-of-theart theory shows that both DDPM (Zhang et al., 2024) and DDIM (Cai & Li, 2025) require on the order of (up to logarithmic factors)
1
Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, USA. Correspondence to: Changxiao Cai <[email protected]>. Proceedings of the 43 rd International Conference on Machine Learning, Seoul, South Korea. PMLR 306, 2026. Copyright 2026 by the author(s).
ε− 1
d+2β β
(1)
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
training samples to generate an output within ε 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.
sion. In addition, we assume that the restriction of the target distribution to each subspace is σ-subgaussian. Under these two assumptions, we construct a kernel-based regularized score estimator sbt for the score function of the Gaussian-smoothed distribution pt := p⋆ ∗N (0, tId ) for any t > 0. Given n samples drawn from the target distribution p⋆ , we establish a finite-sample L2 score estimation error bound: 1 1 1 ⋆ 2 e st (X) − st (X)∥2 = O + (k∨2)/2+1 . E ∥b n t t
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 (Chen et al., 2023; Tang & Yang, 2024; Azangulov et al., 2024; Yakovlev & Puchkin, 2025), a factor model (Chen et al., 2025a), or certain dependence structures (Fan et al., 2025). 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.
Here the expectation is taken over both the training data and X ∼ pt , where we write a ∨ b := max{a, b}. Building on this score estimation guarantee, we prove that diffusion samplers require at most the order of (up to logarithmic factors) ε−(k∨2) training samples to generate a sample that is ε-close in 1-Wasserstein distance to the target distribution p⋆ . 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 kdimensional distribution (Chewi et al., 2024). 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.
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 (Vidal, 2011; Brown et al., 2022). As a result, existing theory still falls short of explaining the empirical effectiveness of diffusion models when learning low-dimensional, multi-modal distributions.
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.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.
1.2. Related works Statistical theory for diffusion models. Recent work has begun to provide finite-sample guarantees for diffusionbased sampling by studying statistical bounds for score estimation error. Under strong density regularity assumptions (e.g., boundedness on compact domains), Oko et al. (2023) showed that neural-network-based ERM score estimators lead to minimax-optimal rates in both TV and W1 distances when used with reverse SDE samplers. Using nonparametric
Concretely, we consider a target data distribution p⋆ supported on a union of subspaces (UoS), i.e., supp(p⋆ ) ⊆ ∪M i=1 Vi , where each Vi is a linear subspace with dimension ki . We denote by k := maxi∈[M ] ki the maximum intrinsic dimen2
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
span({αi }ki=1 ) the linear space spanned by these vectors.
constructions under density lower-bound conditions, Dou et al. (2024) derived minimax-optimal score estimation rates and corresponding sampling guarantees. More recently, for subgaussian targets with β-Hölder smooth densities, Zhang et al. (2024) established minimax-optimality for DDPM using truncated kernel score estimators, and Cai & Li (2025) 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.
For a probability distribution p and a random vector X, we write X ∼ p to mean that X follows the distribution p. We denote by supp(p) the support of probability measure p, i.e., the smallest closed set S such that p(S) = 1. In addition, let 1{·} denote the indicator function. For random vectors X, Y , we define the 1-Wasserstein distance between their distributions pX and pY by ZZ W1 (pX , pY ) := inf ∥x − y∥ γ(dx, dy), γ∈Γ(pX ,pY )
Adaptation to low-dimensional structures. For distributions supported on a linear subspace, Chen et al. (2023) established convergence rates governed by the subspace dimension under smooth score assumptions. For manifoldsupported targets, Azangulov et al. (2024) 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, Fan et al. (2025) obtained minimax-optimal rates for diffusion learning under structured dependence (exponentialinteraction) models. In addition, Wang et al. (2024) analyzed mixtures of low-rank Gaussians, focusing on the special case of orthogonal subspaces. Boffi et al. (2025) 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.
where Γ(pX , pY ) denotes the set of couplings of pX and pY . For probability distributions P, Q, we denote their convolution by P ∗ Q. Finally, we use 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⋆ and gradually adds Gaussian noise. A popular choice is the Ornstein-Uhlenbeck (OU) process (Song et al., 2020): √ dXt = −Xt dt + 2 dBt , with X0 ∼ p⋆ . (2)
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 (Li & Yan, 2024; Liang et al., 2025; Potaptchik et al., 2024; Huang et al., 2024). 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 (Li & Cai, 2025; Zhao & Cai, 2026; Cai & Li, 2026; Chen et al., 2025b; Dmitriev et al., 2026).
Here (Bt )t∈[0,T ] is a standard Brownian motion in Rd . A key property of this OU process is that the conditional distribution of Xt given X0 remains Gaussian for all t. More precisely, one can verify that d
Xt | X0 = ct X0 + σt Wt (3) √ where ct := e−t , σt := 1 − e−2t , and Wt ∼ N (0, Id ) is independent of X0 . In particular, the parameter t fully determines the noise level of the forward process. As t becomes sufficiently large, ct approaches zero and the distribution of Xt becomes close to the standard Gaussian distribution N (0, Id ).
1.3. Notation For a, b ∈ R, we denote a ∨ b := max{a, b} and a ∧ b := min{a, b}. For positive integer M , let [M ] := {1, · · · , M }. For random vector X, we use pX to denote its distribution or probability density function, depending on the context. For any vector x ∈ Rd , we denote ∥ · ∥p as its p-norm, i.e., Pd ∥x∥p := ( i=1 |xi |p )1/p , and write ∥x∥∞ := maxi |xi |. We use ∥ · ∥ to denote the 2-norm for simplicity. For any vector x ∈ Rd and any i, j ∈ [d] with i < j, we denote by xi:j ∈ Rj−i+1 the subvector consisting of the i-th through j-th entries of x. For vectors {αi }ki=1 , we denote by
Reverse process. Running the forward dynamics backward in time transforms Gaussian noise into samples from p⋆ , forming the basis of diffusion-based sampling. For the OU process in (2), its time-reversal SDE is given by Y0 ∼ pXT ,
√ dYt = Yt + 2∇ log pXT −t (Yt ) dt + 2 dB t .
(4)
Here pXT −t denotes the density of the forward process (2) at time T −t and {B t }t∈[0,T ] is a standard Brownian motion in 3
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Algorithm 1 Sampling procedure
2.2. Assumptions
1: Input: Early stopping time τ > 0, end time T > 0,
In this section, we introduce the assumptions imposed on the target distribution p⋆ .
score estimator sbXt for t ∈ [τ, T ]. 2: Sample y ∼ N (0, Id ). 3: Solve the reverse SDE: √ dYbt = Ybt + 2b sXT −t (Ybt ) dt + 2 dBt
First, to capture low-dimensional, multi-modal structure, we assume that the support of p⋆ is contained in a finite union of low-dimensional linear subspaces. Assumption 1 (Union of low-dimensional subspaces). There exist linear subspaces V1 , V2 , . . . , VM ⊆ Rd , with dimension dim(Vi ) = ki , such that
(6)
for t ∈ [0, T − τ ] with Yb0 = y. bT −τ . 4: Output: generated sample Y
supp(p⋆ ) ⊆ ∪M i=1 Vi . Rd . By classical time-reversal results for SDEs (Anderson, d 1982), this process satisfies YT −t = Xt for all t ∈ [0, T ].
Moreover, p⋆ assigns zero probability to intersections between different subspaces, i.e.,
The crucial ingredient in the reverse dynamics is the score function of the marginal distributions of the forward process. For a random vector X ∈ Rd with density pX , its score function is given by s⋆X (x) := ∇ log pX (x) =
∇pX (x) . pX (x)
p⋆ (Vi ∩ Vj ) = 0,
1 , cp M
p⋆ (Vi ) ≥
Since these scores are unknown in practice, they must be estimated from training samples {X (i) }ni=1 drawn from p⋆ .
for some constant cp > 0.
Sampling procedure. Since XT → N (0, Id ) as T → ∞, diffusion-based sampling can be implemented by initializing the reverse dynamics from N (0, Id ) and replacing the true score s⋆Xt with a learned estimator sbXt . The resulting procedure is summarized in Algorithm 1. We introduce an early stopping time τ > 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.
For each subspace Vi , let Ai ∈ Rd×ki be a matrix whose columns form an orthogonal basis of Vi : Vi = span(col(Ai )),
(7)
with
σ2 h(t) := 2t . ct
A⊤ i Ai = Iki .
Denote by proji : Rd → Vi the projection onto Vi , given by
Here {Bt }t≥0 also denotes the Brownian motion and thus Zt follows the distribution p⋆ ∗ N (0, tId ). It is straightforward to verify that the score functions of Xt and Zt satisfy 1 ⋆ x s ( ) ct Zh(t) ct
(10)
For each i ∈ [M ], let p⋆i := p⋆ |Vi denote the restriction of the target distribution p⋆ to subspace Vi . By Assumption PM 1, we can decompose the target distribution as p⋆ = i=1 p⋆i .
Score estimation reduction. To estimate the score function s⋆Xt , it is often more convenient to construct score estimator for the following variance-exploding (VE) process
s⋆Xt (x) =
∀i ∈ [M ]
This assumption provides a tractable abstraction for lowdimensional, multi-modal distributions, where different modes may concentrate on different subspaces. Such unionof-subspaces structure has been widely used in the modeling of heterogeneous high-dimensional data (Wang et al., 2024) and has also been observed empirically in real-world datasets (Brown et al., 2022; Kamkari et al., 2024).
d
with Z0 ∼ p⋆ .
(9)
Finally, each subspace has non-trivial mass:
(5)
dZt = dBt ,
∀i ̸= j.
proji (x) = Ai A⊤ i x. Remark 1. Our framework can be extended naturally to distributions concentrated near a union of low-dimensional subspaces. In this paper, we focus on the noiseless setting to isolate the essential roles of low-dimensional structure and multi-modality, without introducing the additional technical complications caused by ambient noise. Extensions to noisy settings are discussed in Section 6.
(8)
As a result, it suffices to estimate the score function of Zt for any t > 0 and then define sbXt (x) := c1t sbZh(t) ( cxt ) as the estimator of s⋆Xt (x). For notational simplicity, we denote by s⋆t := ∇ log pZt the score function of Zt , and let sbt denote its estimator. The derivation of (8) is provided in Appendix A.2.
Next, we impose a mild subgaussian assumption on the target distribution within each subspace. 4
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Assumption 2 (Subgaussian within each subspace). Let plow be the normalized push-forward distribution of p⋆i onto i ki R under A⊤ i : := law(A⊤ plow i i Z),
p⋆ Z∼ ⋆ i . p (Vi )
where we define R wt (i, x) :=
Vi
φt (x − y; d)p⋆i (dy) pt (x)
=:
qt (i, x) , (12b) pt (x)
⋆ − x−y t φt (x − y; d)pi (dy) . st (i, x) := Vi R φ (x − y; d)p⋆i (dy) Vi t
R
(11)
(12c)
We assume that plow is σi -subgaussian, that is, for any unit i vector θ ∈ Rki with ∥θ∥2 = 1, E exp (X ⊤ θ/σi )2 ≤ 2, X ∼ plow i .
Intuitively, wt (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 st (i, x) is the score of the corresponding smoothed component.
We denote σ := maxi∈[M ] σi .
Our key observation is that each score component st (i, x) ∈ Rd 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 ki -dimensional score function (Chen et al., 2023): 1 ⊤ st (i, x) = − x − proji (x) + Ai slow (13) t (i, Ai x), t
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.
where proji (x) := Ai A⊤ i x is the projection of x onto subspace Vi , and slow (i, ·) : Rki → Rki is the score function of t the ki -dimensional smoothed distribution plow i ∗ N (0, tIki ) (see (11)) on the subspace Vi . This decomposition reduces score estimation to a low-dimensional problem: once Vi is identified, estimating st (i, ·) is governed by the difficulty of estimating the low-dimensional score slow t (i, ·) in dimension ki , rather than the ambient dimension d.
3.1. Algorithm Given n i.i.d. samples {X (i) }ni=1 drawn from the target distribution p⋆ , our goal is to build a score estimator sbt that learns the score function s⋆t of the perturbed data distribution pt = p⋆ ∗ N (0, tId ) for any time t > 0.
Motivated by this observation, we propose a two-step score estimation procedure. We first use the data to estimate the subspaces {Ai }M i=1 and construct a classification function c : Rd → [M ] that assigns points to subspaces (such that c(x) = i if and only if x ∈ Vi ). 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 n0 samples are used for subspace recovery and the remaining N = n − n0 samples are used for score estimation.
Motivation. Observe that the density pt can be written as Z pt (x) = p⋆ ∗ N (0, tId ) (x) = φt (x − y; d)p⋆ (dy), where φt (x; d) := (2πt)−d/2 exp −∥x∥22 /(2t) is the density of N (0, tId ) in Rd . Recall that under the UoS assumpPM tion, we can decompose p⋆ = i=1 p⋆i , yielding pt (x) =
M Z X
φt (x − y; d)p⋆i (dy).
Vi
i=1
In what follows, we describe the proposed score estimator in reverse order.
The gradient of pt can then be computed as ∇pt (x) =
M Z X i=1
−
Vi
Score estimator. We begin by presenting the score estimator assuming access to subspace estimates {Ai }M i=1 and a classification function c(·).
x−y φt (x − y; d)p⋆i (dy). t
Therefore, the score function s⋆t = ∇pt /pt of pt admits the following mixture-type decomposition: s⋆t (x) = =:
M X
1 p (x) t i=1 M X
Z − Vi
Inspired by the score decomposition in (12), we construct the score estimator as a weighted combination of score components associated with each subspace:
x−y φt (x − y; d)p⋆i (dy) t
sbt (x) :=
M X
w bt (i, x)b st (i, x),
(14)
i=1
wt (i, x) · st (i, x),
(12a)
where w bt (i, x) and sbt (i, x) estimate wt (i, x) in (12b) and st (i, x) in (12c), respectively.
i=1
5
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
and for any i ∈ [M ],
• Score component estimator sbt (i, x). In light of the lowdimensional structure in (13), it suffices to learn an low estimator sblow t for each ki -dimensional score st (i, ·), and then construct the estimator for the d-dimensional score st (i, x) associated with Vi as sbt (i, x) := −
N
qbt (i, x) :=
(18b)
x − proji (x) ⊤ + Ai sblow t (i, Ai x). (15) t
We then define the weight estimator as w bt (i, x) :=
(j) To build sblow ∈ Vi } det , let Ci := {j ∈ [N ] : X
note the index set of samples belonging to subspace Vi . Since slow t (i, ·) is the score of the smoothed lowdimensional distribution plow i ∗ N (0, tIki ), we first estimate its density using the Gaussian kernel estimator gbt (i, x) :=
1 X (j) φt (x − A⊤ ; ki ), i X |Ci |
1 X φt (x − X (j) ; d)1{c(X (j) )=i} . N j=1
qbt (i, x) 1{x∈Gt (i)} , pbt (x)
where Gt (i) is a set given by n o Gt (i) := x : ∥x − proji (x)∥2 ≤ Rt (i) ,
(19)
(20)
p with Rt (i) = CR td log(N dtki /2 ) for some universal constant CR > 0.
(16)
j∈Ci
In a word, the weight estimator w bt (i, x) is a plugin estimator for the true weight (12b), up to some where φt (x; ki ) denotes the density of N (0, tIki ). We low-probability set under pt . The indicator function then define the ki -dimensional score estimator as 1{x∈Gt (i)} is introduced for technical convenience in ∇b gt (i, x) log N the analysis and could be removed with a shaper argulow sbt (i, x) := clipR ψ gbt (i, x); . gbt (i, x) N (2πt)ki /2 ment. (17) Subspace recovery. Finally, we briefly discuss how the Here, ψ(x; η) := 1{x ≥ η} is a thresholding function, subspace estimates {Ai }M i=1 and classification function c(·) and the clip operator is defined by can be obtained from training data. This is a classical subspace clustering problem. Under standard identifiability ( z, ∥z∥2 ≤ r; and separation conditions, and assuming known bounds on clipr (z) := z the number of subspaces M and the maximal intrinsic di∥z∥ · r, otherwise. mension k, several polynomial-time methods can recover p the underlying subspaces and cluster assignments, such We set the clipping radius to be R = 2 log N/t. as sparse subspace clustering (Elhamifar & Vidal, 2013), In words, we first form the plug-in estimator ∇b gt /b gt thresholding-based subspace clustering (Heckel & Bölcskei, using the kernel density estimator (16). We then ap2015) and greedy subspace clustering (Park et al., 2014). ply a thresholding rule ψ(b gt ; ηt ), which regularizes From a statistical perspective, this geometric recovery step this ratio according to the estimated density level gbt is typically less demanding than learning the full target disand the threshold ηt = N −1 (2πt)−ki /2 log N that detribution. pends on the sample size N and time t. Specifically, in low-density regions where ∇b gt /b gt is unstable due 3.2. Theoretical guarantees to small denominators and limited data, the resulting We now state our theoretical guarantees for the proposed score estimator sbt is set to zero. This regularization is score estimator and the resulting sampler. important not only for controlling the subsequent estimation error, but also for improving generalization by We first present the L2 error for the proposed score estimator preventing the estimator from closely fitting empirical sbt in (14). The proof can be found in Appendix A.1. artifacts. Theorem 1. Suppose the target distribution p⋆ satisfies Assumptions 1 and 2. Under the event of exact subspace • Weight estimator w bt (i, x). As for the mixture weight recovery and t ≤ N O(1) , the L2 -error of the score estimator wt (i, x) defined in (12b), we first construct Gaussian in (14) using N samples satisfies kernel density estimators for pt (x) and qt (i, x): E ∥b st (X) − s⋆t (X)∥22 N X 1 pbt (x) := φt (x − X (j) ; d), (18a) dM 3 1 σ k∨2 N j=1 ≤ Cscore + (k∨2)/2+1 poly log N N t t 6
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
in 1-Wasserstein distance, it suffices to have ε−(k∨2) samples, up to logarithmic factors, thereby breaking the curse of dimensionality that plagues prior results (Wibisono et al., 2024; Zhang et al., 2024; Dou et al., 2024; Cai & Li, 2025).
for some constant Cscore > 0 independent of N , d, M and t. The expectation here is taken over the i.i.d. training samples {X (i) }N i=1 used for score estimation and X ∼ pt . In words, Theorem 1 shows that the convergence rate of the L2 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 (Wibisono et al., 2024; Zhang et al., 2024; Dou et al., 2024; Cai & Li, 2025),which do not exploit intrinsic low-dimensional data structures and therefore suffer from the curse of dimensionality.
• 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.
Moreover, we note that exact subspace recovery can be achieved with high probability using n0 = Csc M 2 k log n samples for subspace clustering, for a sufficiently large constant Csc > 0. This sample size is negligible compared with the remaining N = n − n0 samples used for score estimation, provided that n is sufficiently large.
Remark 2. We believe that the linear dependence on d in the prefactor of (21) is likely a proof artifact and may be improved through a sharper analysis. Determining whether this dependence is intrinsic or can be removed is an interesting direction for future work. Remark 3. Prior theory (Cai & Li, 2025) suggests that, once the score estimation error is controlled, the discretization error of practical diffusion samplers does not affect the final statistical rate of the sampling error. Accordingly, this work focuses on the main statistical bottleneck, namely score estimation, by analyzing the idealized continuous-time reverse process. Meanwhile, establishing sharp Wasserstein discretization bounds under mild distributional conditions remains an important direction for future work.
We next translate the resulting score estimation guarantee into a sampling guarantee for the diffusion sampler. The proof is deferred to Appendix A.2. Theorem 2. Suppose the target distribution p⋆ satisfies Assumptions 1 and 2. Let n0 = Csc M 2 k log n for some large constant Csc > 0 and N = n − n0 . Then for sufficiently large n, the output YbT −τ of Algorithm 1, using the score estimator in (14) constructed from N samples with T = log n and τ = n−2/k , satisfies E W1 (p⋆ , pYb
T −τ
1 ) ≤ CdM 3/2 n− k∨2 poly log n
(21)
for some constant C > 0 independent of n, d and M . Here the expectation is taken over the samples {X (i) }ni=1 .
4. Analysis This section provides the proof sketches for Theorems 1–2.
Theorem 2 provides the convergence rate of the W1 sampling error for diffusion sampling equipped with the proposed score estimator. By exploiting the intrinsic lowdimensional 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.
Proof sketch of Theorem 1. In light of the expressions of the true score (12) and the score estimator (14), the L2 -error decomposes as: E ∥b st (Zt ) − s⋆t (Zt )∥22 M Z X 2 ≲ E wt (i, x) − w bt (i, x) ∥b st (i, x)∥22 pt (x)dx i=1
Several remarks are in order: + • (Near-)minimax optimality. The minimax risk of es1 timating a k-dimensional density scales as n− k∨2 (Chewi et al., 2024). Since sampling is always harder than density estimation, Theorem 2 shows that our sampling algorithm is minimax optimal (up to logarithmic factors).
M Z X
wt2 (i, x) E ∥st (i, x) − sbt (i, x)∥22 pt (x)dx,
i=1
(22) where the expectation in the first line is taken over both Zt ∼ pt and the i.i.d. samples {X (i) }N i=1 , 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 w bt (i, x) and score estimator sbt (i, x).
• Sample complexity. The error bound in (21) demonstrates that in order to achieve an ε-accurate sampling 7
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Lemma 2. For any fixed ni , the score estimator sbt (i, x) in (15) satisfies Z E ∥b st (i, x) − st (i, x)∥22 1{Ni ≥ni } qt (i, x) dx
For the weight estimator, we recall the true weight wt (i, x) = qt (i, x)/pt (x) in (12b). The mean squared error is controlled in the following lemma, with proof deferred to Appendix B.2.
Rd
Lemma 1. For any x ∈ Gt (i), the weight estimator in (19) satisfies
≤ Cki
2 E wt (i, x) − w bt (i, x) M 1 1 X − 1 ∥x−proji (x)∥22 ≲ 2 e 2t qt (i, x) , d/2 pt (x) t N i=1
Combining (23)–(25) completes the proof of Theorem 1.
Remark 4. The above mean squared error bound depends on the point x, density pt (x) and the geometric structure. This enables us to obtain a tight L2 error bound and further get rid of the dependence on ambient dimension d.
Proof sketch of Theorem 2. As we will show in Lemma 3 in Appendix B, n0 = Csc M 2 k log n samples suffice to recover the subspaces exactly with probability at least 1 − M n−10 , for a sufficiently large constant Csc > 0. Thus, for large n, the remaining sample size for score estimation satisfies N = n − n0 ≥ 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⋆ and the generated distribution pYb to the integral of the score estimation T −τ error over time via the following stability bound (Oko et al., 2023; Azangulov et al., 2024; Tang & Yang, 2024): E W1 (p⋆ , pYb ) T −τ √ √ ≲ d τ + δ + e−T + s Z L−1 X 1 Tj+1 σTj+1 log sXt (X) − s⋆Xt (X)∥22 dt E ∥b δ Tj j=0
With Lemma 1 in hand, we can now bound the first term in (22) √ associated with weight estimation. It is easy to see pt is σ 2 + t-subgaussian. Define Bt := x ∈ Rd : ∥A⊤ i x∥2 ≤ Bt , ∀i ∈ [M ] , p with Bt := CB k(σ 2 + t) log N for some universal constant CB > 0. One can show that the estimation error within Bt dominates since Btc is a low probability region w.r.t. pt . This allows us to apply Lemma 1 to derive the following bound that only depends on ki : Z E (wt (i, x) − w bt (i, x))2 pt (x) dx M Z X j=1
e =O
⊤ 2 1 1 e− 2t ∥x−Aj Aj x∥2 dx d/2 Bt N t
X M
(σ 2 + t)kj /2 . N tkj /2 j=1
(24)
for some constant Cki > 0 only depending on ki . In addition, for any x, the ℓ2 -norm of sbt (i, x) is bounded by r 2 ∥x − proji (x)∥2 ∥b st (i, x)∥2 ≲ + log N . (25) t t
where the expectation is taken over the i.i.d. samples {X (i) }N i=1 .
≲
p⋆i (Vi ) 1 σ ki + k /2+1 poly log N ni t t i
(26) for 0 < T0 = τ < T1 < · · · < TL = T and any δ > 0. Here the expectation is taken over the randomness of samples and X ∼ pXt .
(23)
By Theorem 1, together with N ≥ n/2 and the score relation in (8), one can show that Z E ∥b sXt (x) − s⋆Xt (x)∥22 pXt (x) dx Rd 3 1 σ k∨2 e dM + . =O n h(t) h(t)(k∨2)/2+1
Regarding the second term in (22) associated with score estimation, notice that wt (i, x) = qt (i, x)/pt (x) ≤ 1. Thus, it suffices to bound Z wt2 (i, x)E ∥st (i, x) − sbt (i, x)∥22 pt (x) dx Z ≤ E ∥st (i, x) − sbt (i, x)∥22 qt (i, x) dx.
Observe that h(t) = 2e2t = 2/c2t , and thus Z Tj+1 Z E ∥b sXt (x) − s⋆Xt (x)∥22 pXt (x) dx dt
PN To this end, let Ni := j=1 1{c(X (j) )=i} denote the sample size that we use to estimate the score on subspace Vi . The following lemma provides a mean squared error bound for the score estimator sbt (i, x) in (15). The proof is deferred to Appendix B.4.
Rd
Tj
dM 3 h(Tj+1 ) 2σ k∨2 1 e =O log + . n h(Tj ) k ∨ 2 h(Tj )(k∨2)/2
8
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
We then choose a dyadic partition of the time interval by setting T ≍ log n,
Tj+1 = 2Tj ,
τ ≍ n−γ
for some γ > 0 that will be specified below. Plugging these into (26) yields s Z Tj+1 L−1 X E ∥b sXt (X) − s⋆Xt (X)∥22 dt (1 − e−4Tj ) Tj
j=0
(
√ ≲
dM
3/2
poly log n
k 1 √1 τ − 4 + 2 , n √1 , n
k≥2 k=1 Figure 1. Empirical L2 -score error versus diffusion time t.
Finally, taking δ = n−1 and τ = n−2/k , we conclude E W1 (p⋆ , pYb
T −τ
dM 3/2 ) ≲ 1/(k∨2) poly log n. n
5. Numerical results
that diffusion-based sampling can learn the target distribue −(k∨2) ) samples, where k represents tion using at most O(ε 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.
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 L2 -score estimation guarantee in Theorem 1.
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.
We consider a synthetic target distribution in Rd 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 L2 -score estimation error by Monte Carlo using 10,000 independent samples from pt and average the result over 20 independent training datasets generated from the same target distribution. Figure 1 plots the empirical L2 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 lowdimensional, multi-modal data commonly observed in practice. We construct a kernel-based score estimator and prove 9
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Acknowledgements
Chen, S., Cong, K., and Li, J. Optimal inference schedules for masked diffusion models. arXiv preprint arXiv:2511.04647, 2025b.
C. Cai is supported in part by the NSF grant DMS-2515333.
Chewi, S., Niles-Weed, J., and Rigollet, P. Statistical optimal transport. arXiv preprint arXiv:2407.18163, 3, 2024.
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.
Cole, F. and Lu, Y. Score-based generative models break the curse of dimensionality in learning a family of sub-gaussian probability distributions. arXiv preprint arXiv:2402.08082, 2024.
References
Dmitriev, D., Huang, Z., and Wei, Y. Efficient sampling with discrete diffusion models: Sharp and adaptive guarantees. arXiv preprint arXiv:2602.15008, 2026.
Anderson, B. D. Reverse-time diffusion equation models. Stochastic Processes and their Applications, 12(3):313– 326, 1982.
Dou, Z., Kotekal, S., Xu, Z., and Zhou, H. H. From optimal score matching to optimal sampling. arXiv preprint arXiv:2409.07032, 2024.
Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and Van Den Berg, R. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34:17981–17993, 2021.
Elhamifar, E. and Vidal, R. Sparse subspace clustering: Algorithm, theory, and applications. IEEE transactions on pattern analysis and machine intelligence, 35(11): 2765–2781, 2013.
Azangulov, I., Deligiannidis, G., and Rousseau, J. Convergence of diffusion models under the manifold hypothesis in high-dimensions. arXiv preprint arXiv:2409.18804, 2024.
Fan, J., Gu, Y., and Li, X. Optimal estimation of a factorizable density using diffusion models with relu neural networks. arXiv preprint arXiv:2510.03994, 2025.
Boffi, N., Jacot, A., Tu, S., and Ziemann, I. Shallow diffusion networks provably learn hidden low-dimensional structure. In International Conference on Learning Representations, volume 2025, pp. 52889–52923, 2025.
Heckel, R. and Bölcskei, H. Robust subspace clustering via thresholding. IEEE transactions on information theory, 61(11):6320–6342, 2015.
Brown, B. C., Caterini, A. L., Ross, B. L., Cresswell, J. C., and Loaiza-Ganem, G. Verifying the union of manifolds hypothesis for image data. arXiv preprint arXiv:2207.02862, 2022.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020.
Cai, C. and Li, G. Minimax optimality of the probability flow ode for diffusion models. arXiv preprint arXiv:2503.09583, 2025.
Ho, J., Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., and Fleet, D. J. Video diffusion models. Advances in neural information processing systems, 35:8633–8646, 2022.
Cai, C. and Li, G. Confidence-based decoding is provably efficient for diffusion language models. arXiv preprint arXiv:2603.22248, 2026.
Huang, Z., Wei, Y., and Chen, Y. Denoising diffusion probabilistic models are optimally adaptive to unknown low dimensionality. arXiv preprint arXiv:2410.18784, 2024.
Chen, M., Huang, K., Zhao, T., and Wang, M. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data. In International Conference on Machine Learning, pp. 4672–4712. PMLR, 2023.
Kamkari, H., Ross, B. L., Hosseinzadeh, R., Cresswell, J. C., and Loaiza-Ganem, G. A geometric view of data complexity: Efficient local intrinsic dimension estimation with diffusion models. Advances in Neural Information Processing Systems, 37:38307–38354, 2024.
Chen, M., Xu, R., Xu, Y., and Zhang, R. Diffusion factor models: Generating high-dimensional returns with factor structure. arXiv preprint arXiv:2504.06566, 2025a.
Li, G. and Cai, C. Breaking AR’s sampling bottleneck: Provable acceleration via diffusion language models. Advances in Neural Information Processing Systems, 38: 11700–11725, 2025.
Chen, S., Kontonis, V., and Shah, K. Learning general gaussian mixtures with efficient score matching. arXiv preprint arXiv:2404.18893, 2024. 10
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Li, G. and Yan, Y. Adapting to unknown low-dimensional structures in score-based diffusion models. Advances in Neural Information Processing Systems, 37:126297– 126331, 2024.
Vidal, R. Subspace clustering. IEEE Signal Processing Magazine, 28(2):52–68, 2011. Wang, P., Zhang, H., Zhang, Z., Chen, S., Ma, Y., and Qu, Q. Diffusion models learn low-dimensional distributions via subspace clustering. arXiv preprint arXiv:2409.02426, 2024.
Li, S., Chen, S., and Li, Q. A good score does not lead to a good generative model. arXiv preprint arXiv:2401.04856, 2024. Liang, J., Huang, Z., and Chen, Y. Low-dimensional adaptation of diffusion models: Convergence in total variation. arXiv preprint arXiv:2501.12982, 2025.
Wibisono, A., Wu, Y., and Yang, K. Y. Optimal score estimation via empirical bayes smoothing. In The Thirty Seventh Annual Conference on Learning Theory, pp. 4958–4991. PMLR, 2024.
Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., and Li, C. Large language diffusion models. arXiv preprint arXiv:2502.09992, 2025.
Yakovlev, K. and Puchkin, N. Generalization error bound for denoising score matching under relaxed manifold assumption. arXiv preprint arXiv:2502.13662, 2025.
Oko, K., Akiyama, S., and Suzuki, T. Diffusion models are minimax optimal distribution estimators. In International Conference on Machine Learning, pp. 26517– 26582. PMLR, 2023.
Zhang, K., Yin, C. H., Liang, F., and Liu, J. Minimax optimality of score-based diffusion models: Beyond the density lower bound assumptions. arXiv preprint arXiv:2402.15602, 2024.
Park, D., Caramanis, C., and Sanghavi, S. Greedy subspace clustering. Advances in neural information processing systems, 27, 2014.
Zhao, Y. and Cai, C. Adaptation to intrinsic dependence in diffusion language models. arXiv preprint arXiv:2602.20126, 2026.
Potaptchik, P., Azangulov, I., and Deligiannidis, G. Linear convergence of diffusion models under the manifold hypothesis. arXiv preprint arXiv:2410.09046, 2024. Shah, K., Chen, S., and Klivans, A. Learning mixtures of gaussians using the ddpm objective. Advances in Neural Information Processing Systems, 36:19636–19649, 2023. Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learning, pp. 2256–2265. pmlr, 2015. Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019. Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020. Song, Y., Shen, L., Xing, L., and Ermon, S. Solving inverse problems in medical imaging with score-based generative models. arXiv preprint arXiv:2111.08005, 2021. Tang, R. and Yang, Y. Adaptivity of diffusion models to manifold structures. In International Conference on Artificial Intelligence and Statistics, pp. 1648–1656. PMLR, 2024. Vershynin, R. High-dimensional probability. University of California, Irvine, 10(11):31, 2020. 11
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
A. Proof of theorems A.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 A as o n N A := Ni ≥ , ∀i ∈ [M ] . 2cp M PN Here Ni := j=1 1{c(X (j) )=i} is a random number of samples on Vi . The following claim tells us that, event A happens with high probability. Claim 1. Under Assumption 1, it holds that, −
P[Ac ] ≤ M e
N 2 2c2 pM
.
The proof follows from concentration inequality and can be found in Appendix C.1. We first consider the error on event A, where we have enough samples on each subspace. Notice the score decomposition (12) and the score estimator in (14), the L2 estimation error can be written as Z h Z M i X 2 E s⋆t (x) − sbt (x) 2 1A pt (x)dx = E[∥ wt (i, x)st (i, x) − w bt (i, x)b st (i, x) ∥22 1A ]pt (x)dx i=1
≤M
M Z X
h i 2 E wt (i, x)st (i, x) − w bt (i, x)b st (i, x) 2 1A pt (x)dx
(C-S Ineq)
i=1
=: M
M X
Li .
i=1
We further decompose Li using the error of weight estimator and the error of score estimator respectively, Z h i 2 Li := E wt (i, x)st (i, x) − w bt (i, x)b st (i, x) 2 1A pt (x)dx Z h Z h i i 2 ≲ E wt (i, x) − w bt (i, x) · ∥b st (i, x)∥22 1A pt (x)dx + E wt (i, x))2 · ∥st (i, x) − sbt (i, x)∥22 1A pt (x)dx =: Li,1 + Li,2 . Bound of Li,1 .
We utilize a set Bt here and further decompose Li,1 as Z h i 2 := Li,1 E wt (i, x) − w bt (i, x) · ∥b st (i, x)∥22 1A pt (x)dx Z h i 2 ≤ E wt (i, x) − w bt (i, x) · ∥b st (i, x)∥22 1A pt (x)dx G (i)∩Bt Z t h i 2 + E wt (i, x) − w bt (i, x) · ∥b st (i, x)∥22 1A pt (x)dx G (i)c ∩B Z t h t i 2 + E wt (i, x) − w bt (i, x) · ∥b st (i, x)∥22 1A pt (x)dx Btc
=: κ1 + κ2 + κ3 , here we recall that Gt (i) is the regularization set defined in (20) and Bt is defined as, n o p Bt := x ∈ Rd : ∥A⊤ with Bt := CB k(σ 2 + t) log(N ). i x∥2 ≤ Bt , ∀i ∈ [M ] , for certain absolute constant CB > 0. The following claim tells us that the probability outside Bt is negligible and see Appendix C.2 for its proof. 12
(27)
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Claim 2. Under Assumption 2, for the set Bt , it holds that, Z Mk pt (x)dx ≲ 4 c N Bt √ Z d M 2 2 (σ + t). ∥x∥2 · pt (x)dx ≲ N4 Btc • For κ1 , notice that for x ∈ Gt (i), Rt (i) ∥b st (i, x)∥2 ≤ + t
r
2 log N t
using Lemma 2. Hence, Z h i 2 κ1 := E wt (i, x) − w bt (i, x) · ∥b st (i, x)∥22 1A pt (x)dx Gt (i)∩Bt
≲
≤
≤ ≲
R (i)2 t
t2 R (i)2 t
t2 R (i)2 t
t2
1 + · t
M X 1 1 − 2t ∥x−projj (x)∥22 e · q (j, x) dx t d/2 N · p (x) t Gt (i)∩Bt (2πt) j=1
Z
(Lemma 1)
M Z ⊤ 2 1 1 X 1 + e− 2t ∥x−Aj Aj x∥2 dx · d/2 t N (2πt) j=1 Bt
+
M k 1 X 2kj Bt j · t N (2πt)kj /2 j=1
d M σk · 1 + k/2 · poly log N + log t t N t
In the third inequality, we apply the following changing variable technique in integration and then use Tonelli’s Theorem, ! A⊤ j x z = Tj x := P ⊤ x (28) V⊥ j
here columns of PVj⊥ denotes an orthogonal basis of Vj⊥ and this is an orthogonal transform with 2 ⊤ 2 ∥x − Aj A⊤ j x∥2 = ∥PV ⊥ x∥2 . j
• For κ2 , we have, Z κ2 := Gt (i)c ∩Bt
h i 2 E wt (i, x) − w bt (i, x) · ∥b st (i, x)∥22 1A pt (x)dx
qt2 (i, x) E[∥b st (i, x)∥22 ] · pt (x)dx 2 Gt (i)c ∩Bt pt (x) Z 1 ∥x − A A⊤ x∥2 i i 2 ≲ + · qt (i, x)dx 2 t t c Gt (i) ∩Bt Z
=
(Lemma 2 + qt (i, x) ≤ pt (x))
Since it holds that, Z qt (i, x) :=
1
⊤
2
φt (x − y; d)p⋆i (dy) ≤ p⋆i (Vi ) · (2πt)−d/2 e− 2t ∥x−Ai Ai x∥2 ,
Vi
13
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
we apply the similar linear transform as (28), Z 1 ∥z 2 2 1 ki +1:d ∥2 κ2 ≲ p⋆i (Vi ) · (2πt)−d/2 e− 2t ∥zki +1:d ∥2 dz 1∥z∥≤Bt · 1∥zki +1:d ∥2 ≥Rt (i) + 2 t t Rd Z R (i)2 1 d · R (i)2 + d2 tC t t 1 · exp − dz1:ki ≲ p⋆i (Vi )(2πt)−ki /2 1∥z1:ki ∥≤Bt + t t2 d · tC1 Rki d2 R (i)2 2ki Btki ⋆ d · Rt (i)2 t pi (Vi ) · exp − ≲ + 2 k /2 t t d · tC1 (2πt) i d σ ki d σk ≲ 1 + k /2 · poly log N ≲ 1 + k/2 · poly log N. Nt Nt t i t
(Lemma 6)
Here C1 > 0 is a universal constant which is related with the sub-gaussian norm of standard Gaussian distribution. • For κ3 , Z κ3 := Btc
Z ≤ Btc
Z ≲
h i 2 E wt (i, x) − w bt (i, x) · ∥b st (i, x)∥22 1A pt (x)dx h i E ∥b st (i, x)∥22 pt (x)dx ∥x − A A⊤ x∥2 i i t2
Btc
Z ≤ Btc
∥x∥2 2
t2
+
2
+
1 · pt (x)dx t
(Lemma 2)
1 pt (x)dx t
dM (σ 2 + t) poly log N (Claim 2) t2 N 2 d σ k∨2 ≲ 1 + (k∨2)/2 poly log N. tN t ≲
• For Li,1 , we could sum them up, Li,1 ≤ κ1 + κ2 + κ3 σ k∨2 dM 1 + (k∨2)/2 · poly log N + log t . ≲ Nt t Bound of Li,2 . Z Li,2 :=
h i E wt (i, x))2 · ∥st (i, x) − sbt (i, x)∥22 · 1A pt (x)dx
i qt2 (i, x) h 2 E s (i, x) − s b (i, x) · 1 pt (x)dx t t A 2 p2 (x) Z ht i 2 ≤ E st (i, x) − sbt (i, x) 2 · 1{Ni ≥N/2cp M } qt (i, x)dx (since qt (i, x) ≤ pt (x)) √ σ ki cp M (4/ π)ki 1 + k /2+1 poly log N (Lemma 2). ≲ p⋆ (Vi ) N t t i Z
=
Bound of error on A. Therefore, Z M X ⋆ 2 E[∥st (x) − sbt (x)∥2 1A ]pt (x)dx ≤ M (Li,1 + Li,2 ) i=1
dM 3 σ k∨2 ≲ 1 + (k∨2)/2 · poly log N + log t . Nt t Here for simplicity, we omit constant terms cp and terms that are only related with intrinsic dimension k. 14
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Bound of error outside Ac . Z
E[∥s⋆t (x) − sbt (x)∥22 1Ac ]pt (x)dx ≲
Z
∥st (x)∥22 pt (x)dx · P[Ac ]
1 ∥x∥22 + 2 pt (x)dx · P[Ac ] (Lemma 2) t t d d(σ 2 + t) − 2N 2 2cp M · 2M e ≲ + t t2
+
Z
√ Here we apply Lemma 11 in (Cai & Li, 2025) and that p⋆ ∗ N (0, tId ) is c σ 2 + t subgaussian r.v as we have proven in −
Appendix C.2. Notice that, for large enough N , e Z
Summary. Z
N 2 2c2 pM
≤
M 2 c2p and thus, N
E[∥s⋆t (x) − sbt (x)∥22 1Ac ]pt (x)dx ≲
σ2 dM 3 1+ . Nt t
In summary, our analysis above shows that,
E ∥b st (x) − s⋆t (x)∥22 pt (x)dx = ≲
Z
E[∥s⋆t (x) − sbt (x)∥22 1A ]pt (x)dx +
Z
E[∥s⋆t (x) − sbt (x)∥22 1Ac ]pt (x)dx
dM 3 σ k∨2 1 + (k∨2)/2 · poly log N + log t . Nt t
This proves Theorem 1. A.2. Proof of Theorem 2 Proof of (8).
Recall the forward process (3), it holds that, Z pXt (x) = Z =
pXt |X0 (x|y)p⋆ (dy) 2πσt2
−d/2 − 2σ12 ∥x−ct y∥22 ⋆ p (dy) e t
Hence, ∇x pXt (x) = (2πσt2
−d/2
·
Z
−
x − ct y − 2σ12 ∥x−ct y∥22 ⋆ e t p (dy) σt2
Therefore, its score function, ∇x pXt (x) pXt (x) R x−ct y − 2σ12 ∥x−ct y∥22 ⋆ p (dy) − σ2 e t t = R − 2σ12 ∥x−ct y∥22 e t p⋆ (dy)
s⋆Xt (x) =
Similarly, for any t > 0 and the VE process (7), it holds that, R s⋆Zt (x) =
1 2 − x−y e− 2t ∥x−y∥2 p⋆ (dy) t R − 1 ∥x−y∥2 2 p⋆ (dy) e 2t 15
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Therefore, R s⋆Xt (ct x) =
− ct (x−y) σt2 R
e
−
− c2t ∥x−y∥2 2 2 ⋆ e 2σt p (dy)
c2 t ∥x−y∥2 2 2σ 2 t
p⋆ (dy)
1 ⋆ (x) s ct Zh(t) 1 =: s⋆h(t) (x) ct =
σ2
−2t
2t Here h(t) := c2t = 1−e e−2t = e − 1. t
Proof of Theorem 2. We denote E as the event set of exact subspace recovery using n0 samples. From Lemma 3, we could take n0 = O(M 2 k log n) and thus P[E c ] ≲ M n−10 . Hence, for large n such that n0 ≤ 0.5n, we have the sample size for score estimation: N ≥ 0.5n. Therefore, for the score estimation error of VE process, we have, E ∥b st (X) − s⋆t (X)∥22 = E ∥b st (X) − s⋆t (X)∥22 · 1E + E ∥b st (X) − s⋆t (X)∥22 · 1E c q d log n σ k∨2 dM 3 1 + (k∨2)/2 · poly log n + log t + · M n−10 + E[∥s⋆t (X)∥42 ] · M n−10 ≲ nt t t σ k∨2 dM 3 1 + (k∨2)/2 · poly log n + log t ≲ (29) nt t R2 (i)
In the first inequality, we apply Theorem 1 and that ∥b st (x)∥2 ≲ tt2 Lemma 11 in (Cai & Li, 2025) for the moment bound of true scores.
n + logt n ≲ d log ; in the last inequality, we apply t
Now we consider the sampling error results from score estimation error. Notice that 1-Wasserstein distance between target distribution and generated distribution YbT −τ using Algorithm 1 has the following control,
E W1 (p⋆ , pYb
T −τ
L−1 X √ √ ) ≲ d τ+
j=0
s log δ −1 · σT2 j+1
Z Tj+1 Z Tj
Rd
E ∥b sXt (x) − s⋆Xt (x)∥22 pXt (x) dx dt + δ + e−T (30)
for certain 0 < T0 = τ < T1 < · · · < TL = T and any δ > 0. This bound can be found as (8) in Azangulov et al. (2024), which cites Lemma D.7 in Oko et al. (2023) as the proof, and also can be found as Lemma B.2 in Tang & Yang (2024). Furthermore, the score estimator we define in (14) satisfies, r w bt (i, x) ̸= 0 =⇒ x ∈ Gt (i) =⇒ ∥b st (i, x)∥ ≲ and hence, ∥b st (x)∥ ≲
q
log n t
log n t , with,
1 x 1 sbXt (x) = sbh(t) ( ) ≲ ct ct ct
s
log n = σt2 /c2t
s
log n . σt2
R Therefore, the condition for (30) is satisfied in our setting. We first derive a bound for Rd E ∥b sXt (x) −s⋆Xt (x)∥22 pXt (x) dx using the result in Theorem 1, 16
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Z
Z 1 x x E ∥b E ∥b sXt (x) − s⋆Xt (x)∥22 pXt (x) dx = 2 sh(t) ( ) − s⋆h(t) ( )∥22 pXt (x) dx (Equation (8)) c ct ct d Rd Zt R 1 2 d/2 (ct ) E ∥b sh(t) (y) − s⋆h(t) (y)∥22 pXt (ct y) dy = 2 ct Z 1 = 2 E ∥b sh(t) (y) − s⋆h(t) (y)∥22 pZh(t) (y) dy ct σ k∨2 1 dM 3 1 + ≲ 2 poly log n + log h(t) (Equation (29)). (k∨2)/2+1 ct n h(t) h(t)
Observe that, h′ (t) = 2e2t = 2/c2t and thus, Z Tj+1 Z Tj
Z dM 3 Tj+1 1 σ k∨2 ′ E ∥b sXt (x) − s⋆Xt (x)∥22 pXt (x) dx dt ≲ + h (t) dt poly log n + log h(t) (k∨2)/2+1 n h(t) h(t) Tj Rd 3 1 h(Tj+1 ) 2σ k∨2 dM + log poly log n + T . ≲ n h(Tj ) k ∨ 2 h(Tj )(k∨2)/2
We then take the synthesized discretization as, Tj+1 = 2Tj ,
T ≍ log n,
τ ≍ n−γ
for some γ > 0 that will be determined later. Then we can easily check that L ≍ log n and, s Z Tj+1 Z L−1 X E ∥b sXt (x) − s⋆Xt (x)∥22 pXt (x) dx dt (1 − e−4Tj ) · √ ≲ √ ≲
≲
Rd
Tj
j=0
L−1
dM 3/2 poly log n X √ n j=0
s
L−1
dM 3/2 poly log n X √ · n j=0
(1 − e−4Tj ) s
4σ k Tj
+1 (k∨2)/2
2k/2 Tj
√ dM 3/2√poly log n τ − k4 + 12 ,
k≥2
k=1
√
2σ k∨2 1 + poly log n 2T (k∨2)/2 k ∨ 2 (e k − 1)
n dM 3/2 √poly log n , n
.
Further take δ = n−1 and take τ = n−2/k , then it holds that, W1 (p⋆ , pYb
T −τ
)≲
dM 3/2 poly log n . n1/(k∨2)
B. Proof of Lemmas As long as we could ensure that with high probability the data segmentation with n0 points are correct and each class contains at least k + 1 points, then we could recover the exact linear subspaces Vj and cj (X) = 1{X∈Vj } almost surely for X ∼ p⋆ . This gives the following Lemma 3. Lemma 3 (Subspace Clustering). Under Assumption 1, there exists an algorithm that uses the upper bound of M and k and ensures exact recovery of linear subspaces and hence the function cj with high probability. That is, for any large n, define n o E := ∀i ∈ [M ], ∃ji ∈ [M ] s.t. Vji = Vbi 17
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Here Vbi denotes the i-th estimated subspace of this algorithm Alg(n0 ). Then P E c ≲ M n−10 0 , and here we take n0 = O(c2p M 2 (k + 1) log n). with the randomness taken over samples {X (i) }ni=1
We will further implement an algorithm that proves Lemma 3 and discuss its complexity in Appendix B.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 (12), we only need to estimate slow and the t following lemma gives a time-dependent rate that only relies on intrinsic dimension ki , using certain low-dimensional estimator (17). Lemma 4 (Low-dim score estimator). For any σ sub-gaussian distribution ν in Rk , denote pt as the density of ν ⋆ N (0, tIk ) and st = ∇ log pt as its score. Suppose that {X (i) }N i=1 are N i.i.d samples from ν. Then for any t > 0, we could construct a kernel-based score estimator sbt using (17) that satisfies, 1. Time-dependent L2 estimation error. Z 4 k 1 1 σk 2 E sbt (x) − st (x) 2 pt (x) dx ≲ √ + k/2+1 (log N )k/2+2 π N t t Rk with expectation taken over samples {X (i) }N i=1 . 2. Bounded estimator. r ∥b st (x)∥2 ≤
2 log N t
Here we adopt the estimator from Cai & Li (2025) with a further cut-off. This achieves a similar rate w.r.t t compared with Zhang et al. (2024); Cai & Li (2025). The proof is provided in Appendix B.3. B.1. Proof of Lemma 3. Basically, subspace clustering for noise-free model aims to solve the following optimization problem iteratively, min A
n0 X
∥X (i) − AA⊤ X (i) ∥0
i=1
Here the L0 norm is defined as, ( 0, ∥x∥0 = 1,
x=0 else
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, n o E0 := ∃ at least k + 1 samples on each subspace n o E1 (p) := x ∈ span{y1 . · · · , yp }, Vj ⊈ span{y1 , · · · , yp }, ∀j ∈ [M ] ,
for p ≤ e k
To ensure that Algorithm 2 works with probability larger than 1 − M n−10 , we only need, P[E0c ] ≲ M n−10 P
i.i.d
x,y1 ,···yp ∼ µ∗
[E1 (p)] = 0,
∀p ≤ e k
(31a) (31b)
Since under E1 (p)c , samples from other subspaces will be excluded via finding the smallest p, we could recover the exact subspace in each iteration. 18
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Algorithm 2 Exact subspace recovery 0 f ≥ M , upper bound of the intrinsic dimension Require: Samples {X (i) }ni=1 , upper bound of the number of subspaces M e k≥k f do 1: for j=1,..,M 2: Iterate over all separations of remained samples into 2 categories with one having e k + 1 samples until finding a case that these e k + 1 samples are linearly dependent. 3: Find the smallest p, such as there exists p + 1 samples from these e k + 1 samples that are linearly dependent. 4: Define Vbj as the span of these p + 1 points. 5: Delete those samples that are on Vbj . 6: end for
Proof of (31b).
We first conditioned on y1 , · · · , yp and Vj ⊈ span{y1 , · · · , yp } for all j ∈ [M ], then, P[X ∈ span{y1 . · · · , yp }] =
M Z X Vj ∩span{y1 .··· ,yp }
j=1
=0
p⋆j (dx)
(Assumption 1)
Then we integrate this over y1 , · · · , yp ∼ µ∗ and get P[E1 (p)] = 0. Proof of (31a). Basically, under Assumption 1, denote Ni (n0 ), ∀i ∈ [M ] as the random variable of the number of samples on Vi with sample size n0 . Then, we could view Ni (n0 ) as the sum of Bernoulli r.v with p = p⋆i (Vi ) ≥ cµ1M and thus, Ni (n0 ) ≥
n0 , 2cµ M
with probability ≥ 1 − 2e
−
n0 2 2c2 µM
here we apply Hoeffding inequality just like the proof of Claim 1. Finally, we could apply union bound and take, n0 = O(c2p M 2 (k + 1) log n). n0 The iteration complexity for Algorithm 2 is bounded by e · M. k+1 B.2. Proof of Lemma 1. The following lemma helps prove this result, Lemma 5 (MSE for pbt (x) and qbt (i, x)). For kernel-based estimators pbt (x) and qbt (i, x) in (18a), (18b), it holds that, 1. They are unbiased estimators for pt (x) and qt (i, x) respectively under Assumption 1, i.e., E[b pt (x)] = pt (x) E[b qt (i, x)] = qt (i, x),
∀i ∈ [M ]
2. Under Assumption 1, we have the following point-wise MSE bound, 2 E pbt (x) − pt (x) ≤ 2 E qbt (i, x) − qt (i, x) ≤
M X 2 1 1 e− 2t ∥x−proji (x)∥2 · qt (i, x) (2πt)d/2 N i=1
(33a)
2 1 1 e− 2t ∥x−proji (x)∥2 qt (i, x), d/2 (2πt) N
(33b)
19
∀i ∈ [M ]
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
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 ∈ Gt (i) h q (i, x) qb (i, x) i 2 2 t t E wt (i, x) − w bt (i, x) =E − pt (x) pbt (x) h q (i, x) − qb (i, x) i h qb (i, x) (b pt (x) − pt (x))2 i 2 2 t t t ≲E +E · pt (x) pbt (x) p2t (x) • For the first term, h q (i, x) − qb (i, x) i 2 1 1 1 2 t t E ≤ 2 · (2πt)−d/2 e− 2t ∥x−proji (x)∥2 qt (i, x). pt (x) pt (x) N
(Lemma 5)
• For the second term, h qb (i, x) (b h pb (x) − p (x) i pt (x) − pt (x))2 i 2 2 t t t E · ≤ E pbt (x) p2t (x) pt (x) M X 1 1 1 ∥x−proji (x)∥22 − 2t ≤ 2 · · q (i, x) e t pt (x) (2πt)d/2 N i=1
(Lemma 5)
Therefore, it holds that, 2 E[ wt (i, x) − w bt (i, x) ] ≲
M X 1 1 1 − 2t ∥x−proji (x)∥22 e · · q (i, x) . t p2t (x) (2πt)d/2 N i=1
Then we prove Lemma 5, and the proof idea is quite similar to the proof of Lemma 3 in Cai & Li (2025). Proof of Lemma 5. • First prove that both are unbiased estimators, N
E[b pt (x)] = E[
1 X φt (x − X (j) ; d)] = EX∼p⋆ [φt (x − X; d)] = pt (x) N j=1 Z
E[b qt (i, x)] = EX∼p⋆ [φt (x − X; d)1{c(X)=i} ] = EX∼p⋆ [φt (x − X; d)1{X∈Vi } ] =
φt (x − y; d)p⋆i (dy)
Vi
In the last equation, we apply Assumption 1. • Then derive the mean squared error (MSE) for both estimators, i.e., both variance. To prove (33a), observe that, Var(b pt (x)) =
M Z 1 1 1 X VarX∼p⋆ (φt (x − X; d)) ≤ EX∼p⋆ [φ2t (x − X; d)] = φ2 (x − y; d)p⋆i (dy). N N N i=1 Vi t
Plug in the structure of Vi : Z Z 2 1 φ2t (x − y; d)p⋆i (dy) = (2πt)−d e− t ∥x−y∥2 p⋆i (dy) Vi Vi Z 2 2 1 1 = (2πt)−d e− t ∥x−proji (x)∥2 e− t ∥y−proji (x)∥2 p⋆i (dy) Vi
Notice that, Z
2 1 e− t ∥y−proji (x)∥2 p⋆i (dy) ≤
Vi
Z Vi
20
1
2
e− 2t ∥y−proji (x)∥2 p⋆i (dy).
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Hence, Z
2
1
φ2t (x − y; d)p⋆i (dy) ≤ (2πt)−d/2 e− 2t ∥x−proji (x)∥2 · qt (i, x).
Vi
Therefore, M 2 1 X E pt (x) − pbt (x) ≤ N i=1
≤
Z
φ2t (x − y; d)p⋆i (dy)
Vi
M X 2 1 1 e− 2t ∥x−proji (x)∥2 · qt (i, x). d/2 (2πt) N i=1
• To prove (33b), similarly, M Z 2 1 X 1 φ2 (x − y; d)1{y∈Vi } p⋆j (dy) E qt (i, x) − qbt (i, x) ≤ EX∼p⋆ [φ2t (x − X; d)1{X∈Vi } ] = N N j=1 Vj t
Under Assumption 1, it holds that, Z Z 2 2 2 1 1 1 1 φ2t (x − y; d)p⋆i (dy) ≤ (2πt)−d e− t ∥x−proji (x)∥2 E qt (i, x) − qbt (i, x) ≤ e− t ∥y−proji (x)∥2 p⋆i (dy) N Vi N V Z i 2 2 1 1 1 ≤ (2πt)−d e− t ∥x−proji (x)∥2 e− 2t ∥y−proji (x)∥2 p⋆i (dy) N Vi 2 1 1 ≤ (2πt)−d/2 e− 2t ∥x−proji (x)∥2 qt (i, x). N B.3. Proof of Lemma 4. In the proof of this lemma, the target distribution ν is a σ-subgaussian distribution in Rk , and we denote pt as the density of i.i.d ν ∗ N (0, tIk ) and st (·) := ∇ log pt (x) as its score. {X (i) }N i=1 ∼ ν are N samples for score estimation. We take the estimator based on (17) as, ∇b gt (x) log N sbt (x) = clipR ψ gbt (x); gbt (x) N (2πt)k/2
r with R :=
2 log N . t
(34)
Here the clip operator is defined as, ( z, clipr (z) := projBr (0) (z),
∥z∥2 ≤ r else
z where projBr (0) (z) = r · ∥z∥ means the projection of z on that ball; ψ(x; η) := 1{x≥η} is the hard thresholding function; 2 P N 1 (i) gbt (x) := N i=1 φt (X − x; k) is the kernel based density estimator. As a remark, the estimator before clipping is almost the same as that in Cai & Li (2025), and we apply hard-thresholding here for simplicity since we use DDPM sampling procedure.
Then we show that the score estimator defined in (34) satisfies the conditions in Lemma 4. In Cai & Li (2025), they defined n cη log N o Ft := x : pt (x) ≥ N (2πt)k/2 and it holds that, σk (log N )k/2+1 π N tk/2 Ftc Z 16 k/2 1 1 σk ∥st (x)∥22 · pt (x) dx ≲ + k/2+1 (log N )k/2+2 π N t t Ftc Z
pt (x) dx ≲
16 k/2 1
21
1+
(35a) (35b)
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
from Lemma 8 in Cai & Li (2025). Notice that from Lemma 4 in Cai & Li (2025), it has been proven that, 1 2 ∥st (x)∥22 ≤ log k/2 t (2πt) pt (x) Hence, for x ∈ Ft , one has, 2 N 2 log ≤ log N. t cη log N t
∥st (x)∥22 ≤
(36)
Here we use the constant selection of cη ≥ 2 as in Cai & Li (2025). That’s the reason why we use the clip operator in (34). Therefore, with the estimator constructed in (34), it automatically satisfies the bounded condition and for its L2 error, Z Z 2 2 E sbt (x) − st (x) 2 pt (x) dx ≤ E sbt (x) − st (x) 2 pt (x) dx k R F Z t 2 + E sbt (x) − st (x) 2 pt (x) dx Ftc
=: χ1 + χ2 • For χ1 . Due to (36), for x ∈ Ft , st (x) = clipR (st (x)). What’s more, clipR (·) is 1-Lip continous, and thus, Z ∇b gt (x) log N 2 χ1 ≤ E ψ gbt (x); − st (x) 2 · pt (x) dx k/2 gbt (x) N (2πt) F Z t ∇b log N gt (x) 2 ψ gbt (x); − st (x) 2 · pt (x) dx ≤ E k/2 gbt (x) N (2πt) Rk √ σk (4/ π)k 1 + k/2+1 (log N )k/2+1 ≤ N t t The last inequality results from Proposition 1 in Cai & Li (2025), which holds for hard thresholding ψ by simple modification. • For χ2 . Notice that, Z
2 E sbt (x) − st (x) 2 pt (x) dx
χ2 := Ftc
Z
≤ Ftc
≤R
2
R2 + ∥st (x)∥22 pt (x) dx
Z
Z pt (x) dx +
Ftc
≲
Ftc
16 k/2 1 1 π
N
t
+
t
∥st (x)∥22 · pt (x) dx
σk (log N )k/2+2 k/2+1
(Using (35a) + (35b))
Therefore, it holds that, Z
4 k 1 1 σk 2 + k/2+1 (log N )k/2+2 . E sbt (x) − st (x) 2 pt (x) dx ≲ √ π N t t Rk
B.4. Proof of Lemma 2 Observe that, conditioned on Ni , we could obtain Ni i.i.d samples from normalized distribution p⋆i to estimate its score function. We only need to prove (24). Notice that, Z Z 2 1 1 ⋆ −d/2 − 2t ∥x−proji (x)∥22 qt (i, x) = φt (x − y; d)pi (dy) = (2πt) e e− 2t ∥proji (x)−y∥2 p⋆i (dy) Vi Vi Z 2 ⊤ 2 1 1 −d/2 − 2t ∥x−proji (x)∥2 ⋆ = (2πt) e pi (Vi ) e− 2t ∥z−Ai x∥2 plow i (dz) Rk 1 −(d−ki )/2 − 2t ∥x−proji (x)∥22 ⋆ ⊤ = (2πt) e pi (Vi ) · plow t (i, Ai x)
22
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions ki low here plow t (i, ·) denotes the density function of pi ∗ N (0, tIki ) in R . Hence,
Z
2 E sbt (i, x) − st (i, x) 2 Ni qt (i, x) dx Rd Z ⊤ 2 1 2 ⊤ low ⊤ −(d−ki )/2 − 2t ⊤ ≤ p⋆i (Vi ) e ∥x−Ai Ai x∥2 · plow E sblow t (i, Ai x) − st (i, Ai x) 2 Ni (2πt) t (i, Ai x) dx Rd
Apply the same linear transform as in (28),
Z
2 E sbt (i, x) − st (i, x) 2 Ni qt (i, x) dx Rd Z 2 1 2 ⋆ low −(d−ki )/2 − 2t ≤ pi (Vi ) e ∥zki +1:d ∥2 · plow E sblow t (i, z1:ki ) − st (i, z1:ki ) 2 Ni (2πt) t (i, z1:ki ) dz d ZR low 2 low E sblow = p⋆ (Vi ) (Tonelli’s Theorem) t (i, z1:ki ) − st (i, z1:ki ) 2 Ni · pt (i, z1:ki ) dz1:ki Rki √ σ ki (4/ π)ki 1 + k /2+1 · poly log N (Lemma 4) ≲ p⋆i (Vi ) · Ni t t i
Finally,
Z Rd
Z
i h 2 E E sbt (i, x) − st (i, x) 2 Ni · 1{Ni ≥ni } qt (i, x) dx Rd i h Z 2 E sbt (i, x) − st (i, x) 2 Ni · qt (i, x) dx · 1{Ni ≥ni } =E Rd √ (4/ π)ki 1 σ ki ⋆ ≲ pi (Vi ) · + k /2+1 · poly log N. ni t t i
2 E sbt (i, x) − st (i, x) 2 1{Ni ≥ni } qt (i, x) dx =
C. Proof of Claims C.1. Proof of Claim 1 PN For any i ∈ [M ], we first bound the probability P[Ni < 2cN ]. Note that Ni = j=1 1{X (j) ∈Vi } is basically the sum of N pM i.i.d Bernoulli r.vs with parameter p⋆ (Vi ) ≥ cp1M (Assumption 1). Applying Hoeffding’s inequality, we have,
h P Ni <
h N i N i = P Ni − N · p⋆ (Vi ) < − N · p⋆ (Vi ) − 2cp M 2cp M 1 2 N ≤ exp − 2N · ( ) = exp − 2 2 . 2cp M 2cp M
Then applying union bound over i ∈ [M ], it holds that,
h P Ac = P ∃i ∈ [M ], Ni < 23
− 2N 2 N i ≤ M e 2cp M . 2cp M
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
C.2. Proof of Claim 2 d ⋆ For the target distribution p⋆ , it is supported on ∪M i=1 Vi . Hence for any θ ∈ R and ∥θ∥2 = 1, denote r.v X ∼ p M Z h (X ⊤ θ)2 i X (x⊤ θ)2 = p⋆i (dx) E exp exp 2 2 σ σ i=1 Vi Z M (z ⊤ A⊤ θ)2 X i = plow p⋆i (Vi ) exp i (dz) 2 ki σ R i=1 Z M (z ⊤ A⊤ θ)2 X i ≤ plow p⋆i (Vi ) exp i (dz) 2 ki σ R i i=1
≤
M X
p⋆i (Vi ) · 2 = 2.
i=1 ⋆ d Here we use ∥A⊤ i θ∥ ≤ 1 and Assumption 2. This shows √ that, p is a sub-gaussian distribution in R with parameter σ. ⋆ 2 Therefore, it is straightforward that p ∗ N (0, tId ) is c σ + t sub-gaussian, for some absolute constant c > 0.
• For
R
Btc
pt (x)dx, it holds that, Z pt (x)dx ≤ Btc
M X
PZt ∼pt ∥A⊤ i Zt ∥2 > Bt
i=1
≤
M X
p PZt ∼pt ∥A⊤ i Zt ∥∞ > Bt / ki ≲ M k exp
i=1
cBt2 M k ≲ 4 ki (σ 2 + t) N
Recall the definition of Bt in (27), and we take large enough constant CB > 0 to ensure exp • For
R
Btc
(Lemma 6)
cBt2 −4 . ki (σ 2 +t) ≲ N
∥x∥22 · pt (x)dx, it holds that, Z Btc
∥x∥22 · pt (x)dx = EZt ∼pt [∥Zt ∥22 1{Zt ∈B / t}] s ≤
EZt ∼pt [∥Zt ∥42 ] ·
Z pt (x)dx
(C-S Ineq)
Btc
r
Mk d2 (σ 2 + t)2 4 N √ d M 2 ≲ (σ + t). N2
≲
D. Auxiliary Lemmas Lemma 6 (Tail bound for subgaussian random vectors). Let ν be a σ-subgaussian distribution in R, i.e, n o σ = ∥X∥ψ2 := inf t > 0 : E exp X 2 /t2 ≤ 2 , X ∼ ν.
(37)
Then we have the following tail bound for X ∼ ν, any r ≥ 0 and some absolute constant c > 0, P[|X| ≥ r] ≤ 2 exp −
cr2 σ2
E[|X|2 1{|X|≥r} ] ≤ 2(r2 + σ 2 /c) exp − 24
(38) cr2 . σ2
(39)
Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions
Proof. • Proof of (38). This is a classical result for sub-gaussian r.v. We can find the proof in Proposition 2.5.2 in (Vershynin, 2020). • Proof of (39). For any r ≥ 0, E |X|2 1{|X|≥r} =
Z
x2 1{|x|≥r} · p(x)dx
R
=
Z Z |x| R
2zdz 1{|x|≥r} p(x)dx
0
Z Z +∞ =2
z · 1{|x|≥z} 1{|x|≥r} p(x)dzdx R 0 Z +∞ Z
(Tonelli’s Theorem)
1{|x|≥r∨z} · p(x)dx zdz (Tonelli’s Theorem) 0 R Z +∞ ≤ r2 P[|X| ≥ r] + 2 z · P[|X| ≥ z]dz Zr ∞ 2 2 ≤ r2 P[|X| ≥ r] + 4 z · e−cz /σ dz =2
r
= 2(r2 + σ 2 /c) exp −
25
cr2 . σ2