Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions

Abstract: Score-based diffusion models have demonstrated remarkable empirical success in learning high-dimensional 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 at most $\widetilde{O}(\varepsilon^{-k \vee 2})$ samples to achieve $\varepsilon$ 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.

• The paper introduces a kernel-based score estimator that exploits union-of-subspace structure to accurately learn low-dimensional multi-modal distributions.
• It establishes theoretical guarantees where sample complexity depends on intrinsic dimension instead of ambient dimension, achieving near-minimax optimal error bounds.
• The work relaxes traditional assumptions and offers a roadmap for extending these results to neural network-based score models and more complex geometries.

Diffusion Models Are Statistically Optimal for Learning Low-Dimensional Multi-Modal Distributions

Background and Motivation

Score-based diffusion models (SBDMs) have demonstrated remarkable empirical efficacy in modeling complex distributions with low-dimensional and multi-modal structure. Despite these successes, theoretical guarantees for their sample complexity—particularly in structured, heterogeneous data regimes—remain limited. Prior analyses have typically imposed strong regularity assumptions (e.g., globally smooth score functions, bounded densities, or log-concavity), which hinder applicability to realistic, multi-modal datasets and fail to account for intrinsic structure. Moreover, minimax-optimality results for general distributions suffer from the curse of dimensionality, scaling sample complexity with ambient dimension $d$, rather than the more relevant intrinsic dimension $k$.

This paper addresses these limitations by rigorously analyzing diffusion models for distributions supported on unions of low-dimensional linear subspaces, establishing statistical optimality in sample complexity with minimal distributional assumptions. The context encompasses many practical scenarios where distinct classes or clusters reside in disparate subspaces—a geometric configuration frequently observed in high-dimensional data modalities (e.g., image and signal domains).

Problem Formulation

The target distribution $p^\star$ is supported on a union of $M$ linear subspaces $\{V_i\}_{i=1}^M$ within $\mathbb{R}^d$, with maximal intrinsic dimension $k = \max_i k_i$. Each restriction $p^\star|_{V_i}$ is assumed subgaussian. The task is to learn $p^\star$, generating samples that are $\varepsilon$-close in the $k$0 Wasserstein distance.

The score-based diffusion model consists of a two-stage process:

• Forward process: Progressive corruption of data via Gaussian noise (e.g., an Ornstein-Uhlenbeck process).
• Reverse process: Iterative denoising, guided by learned time-indexed score functions $k$1.

The primary statistical bottleneck is accurate score estimation. The model must reconcile score estimation error (induced by finite samples) and propagation error in the reverse-time sampling dynamics.

Kernel-Based Score Estimation Algorithm

The core contribution is a novel kernel-based, regularized score estimator for the smoothed distribution $k$2. The estimator leverages the union-of-subspaces structure through:

• Subspace recovery: Clustering input data into $k$3 subspaces via classical algorithms (e.g., sparse or greedy subspace clustering), requiring only $k$4 samples.
• Component-wise estimation: For each subspace, the score function of $k$5 is decomposed into normal and tangent components; the tangent component is further reduced to a $k$6-dimensional score function.
• Kernel density estimation + thresholding: Estimation of each low-dimensional score is performed via Gaussian kernel density estimation, then regularized via clipping and hard thresholding to mitigate instability in low-density regions.

The resultant estimator achieves low estimation error, governed solely by intrinsic dimension $k$7.

Theoretical Guarantees

Strong theoretical results are established for both score estimation and sampling error:

• Score estimation error:

$k$8

for $k$9. This bound depends only on $p^\star$0, not $p^\star$1, and holds under minimal assumptions (subgaussianity, no smoothness, boundedness, or log-concavity required).

• Sampling error in Wasserstein distance:

Diffusion samplers achieve minimax-optimal convergence rate, matching the information-theoretic lower bound for $p^\star$2-dimensional distribution learning:

$p^\star$3

with $p^\star$4 only appearing linearly in the prefactor, and up to logarithmic factors.

Figure 1: Empirical $p^\star$5 score estimation error versus diffusion time $p^\star$6, demonstrating intrinsic dimension-driven scaling even in high ambient dimension $p^\star$7.

The empirical results shown in Figure 1 corroborate theoretical predictions: when modeling synthetic distributions with $p^\star$8, multi-modal unions of $p^\star$9-dimensional subspaces, the score estimation error exhibits a scaling dictated by $M$0, not $M$1. This validates the theoretical claim that optimality is determined by intrinsic structure.

Implications and Comparison to Prior Work

Contrast with prior bounds: Previous results for general smooth-density classes incurred sample complexity scaling as $M$2 for $M$3-Hölder densities, which rapidly becomes infeasible for large $M$4. This work demonstrates that, for union-of-subspace models (and with only subgaussian tails), sample complexity is governed by $M$5, enabling efficient learning even in very high-dimensional ambient spaces.

Relaxed assumptions: The analysis applies under substantially weaker conditions than prior optimality results—no bounded-density, log-concave, or smooth score assumptions. Furthermore, it accommodates multi-modality and well-separated components (with vanishing density between modes), a setting excluded from earlier theoretical frameworks.

Algorithmic design: Although kernel score estimation serves as a proof device, the decomposition highlights explicit low-dimensional approximation targets for neural network-based score approximators. This paves the way for extending statistical guarantees to practical architectures.

Discretization and iteration complexity: Recent work has shown that sampling complexity (number of reverse-process iterations) also adapts to intrinsic dimension, and discretization error does not degrade the end-to-end statistical rate once score estimation error is controlled.

Future Directions

Several avenues emerge:

• Neural network-based scores: Extending the analysis to NN-based ERM estimators entails rigorous quantification of both approximation and generalization errors, with complexity controlled by intrinsic dimension.
• Manifolds and noisy unions: Realistic data often resides near (rather than on) unions of subspaces or manifolds. The normal-tangent decomposition remains applicable, albeit with additional technical complexity in controlling ambient noise.
• End-to-end Wasserstein bounds: Integrating statistical and numerical (discretization) error into sharp, non-asymptotic convergence guarantees remains a pertinent challenge.

Conclusion

This paper establishes that score-based diffusion models achieve near-minimax optimality in learning low-dimensional multi-modal distributions, with sample complexity scaling in the intrinsic dimension. The results deepen theoretical understanding of SBDMs' statistical efficiency, validate empirical phenomena in high-dimensional generative modeling, and provide a roadmap for extending guarantees to practical (NN-based) models and more complex data geometries (2605.30153).