Score Approximation, Estimation and Distribution Recovery of Diffusion Models on Low-Dimensional Data

Minshuo Chen; Kaixuan Huang; Tuo Zhao; Mengdi Wang

Source-linked AI summary

Score Approximation, Estimation and Distribution Recovery of Diffusion Models on Low-Dimensional Data

Minshuo Chen, Kaixuan Huang, Tuo Zhao, Mengdi Wang

arXiv:2302.07194v1Feb 14, 2023cs.LGstat.ML

TL;DR

The paper studies whether diffusion models can approximate and estimate score functions, then recover distributions, for data supported on unknown low-dimensional linear subspaces. It analyzes an encoder-decoder neural score estimator and the resulting backward process, showing convergence rates governed by intrinsic rather than ambient dimension. The main scope boundary is the assumed linear-subspace data structure, alongside a tradeoff in the early stopping time between score estimation and distribution recovery.

Problem
The paper asks whether neural networks can learn score functions and whether diffusion models can recover distributions when data have low-dimensional geometric structure.
Method
The paper analyzes an encoder-decoder neural score estimator and a discretized backward process for data supported on unknown low-dimensional linear subspaces.
Results
The learned score accurately approximates and estimates the ground-truth score, while the generated distribution converges near the data distribution at rates depending on intrinsic dimension.
Takeaways & Limitations
The results indicate that diffusion models can capture unknown linear data geometry and avoid the curse of ambient dimensionality in this setting.
Takeaways & Limitations
The analysis assumes data are supported on a low-dimensional linear subspace, leaving more general manifold assumptions for future work.

Abstract

from arXiv · show

Diffusion models achieve state-of-the-art performance in various generation tasks. However, their theoretical foundations fall far behind. This paper studies score approximation, estimation, and distribution recovery of diffusion models, when data are supported on an unknown low-dimensional linear subspace. Our result provides sample complexity bounds for distribution estimation using diffusion models. We show that with a properly chosen neural network architecture, the score function can be both accurately approximated and efficiently estimated. Furthermore, the generated distribution based on the estimated score function captures the data geometric structures and converges to a close vicinity of the data distribution. The convergence rate depends on the subspace dimension, indicating that diffusion models can circumvent the curse of data ambient dimensionality.

1 Introduction

The paper develops an integrated theory for score approximation, score estimation, and distribution recovery when high-ambient-dimensional data lie on unknown low-dimensional linear subspaces.

Research questions: Diffusion models are studied through two questions: whether neural networks can learn structured-data score functions and whether learned scores recover the data distribution.The questions include architecture choice, sample complexity, geometric structure capture, and its effect on estimation.
Motivation: High-resolution data may have low-dimensional structure despite high ambient dimension, but whether diffusion models capture these structures remains unclear.Prior work often assumes an accurate score estimator and therefore does not address score approximation and learning directly.
Setting: The analysis assumes data x = Az supported on an unknown d-dimensional linear subspace in D dimensions, with d < D.The columns of A form an orthonormal basis, while z is the latent variable.
Contributions: The proposed integrated analysis provides distribution-estimation sample-complexity bounds for low-dimensional linear-subspace data.It addresses score learning and distribution recovery within the same framework.
Contributions: An encoder-decoder network with skip-layer connections approximates and estimates scores, with estimation rates depending on intrinsic dimension rather than ambient dimension.The stated estimation rate is e O(1/√t0 n^−1/(d+5)).
Results: The generated distribution converges to a close vicinity of the data distribution, and the analysis is free of the curse of ambient dimensionality.The introduction reports an e O(n^−1/(2(d+5))) total-variation rate in one direction and vanishing magnitude in the orthogonal direction.

2 Preliminaries

Diffusion models add Gaussian noise through a forward SDE and reverse the process using the score function; neural score matching provides a trainable approximation for practical generation.

Forward and backward SDEs: The forward process progressively adds noise, while the backward process denoises data to generate samples from an initial Gaussian state.The forward process transforms the data distribution toward white noise, and the backward process reverses that corruption.
Forward and backward SDEs: The Ornstein-Uhlenbeck forward SDE shrinks data and adds Gaussian noise, with conditional distribution Xt|X0 given by N(α(t)X0, h(t)ID).Here α(t) and h(t) determine the mean shrinkage and noise variance.
Forward and backward SDEs: The reverse-time SDE uses ∇log pt, the score function, and matches the time-reversed forward process under mild conditions.In practice, the terminal distribution is replaced by a standard Gaussian and the score by a neural estimator.
Score matching: Score matching replaces the unavailable score target with an equivalent denoising objective based on the Gaussian forward transition kernel.The noise-gradient target has an analytical form because the transition kernel is Gaussian.
Score matching: Empirical score matching samples noisy observations and times over [t0, T], avoiding [0, t0] to prevent score blowup and stabilize training.The early stopping time t0 is later selected based on sample size.

3 Score decomposition

For data on a low-dimensional linear subspace, the score separates into on-support and orthogonal components with different roles in distribution learning and subspace recovery.

Score decomposition: The score decomposition enables efficient approximation and estimation by exploiting distinct behaviors of the on-support and orthogonal components.The analysis uses these properties to design the subsequent score-learning guarantees.
Assumption: Under x = Az, the data lie on a d-dimensional unknown linear subspace with orthonormal basis matrix A and latent variable z.The latent variable follows a distribution with a density on R^d.
Score components: The on-support score lies in the column span of A and depends on projected data, while the orthogonal score captures directions outside the subspace.The orthogonal component becomes singular as t approaches zero because h(t) approaches zero.
Backward dynamics: The orthogonal backward dynamics eventually have vanishing drift coefficients, so the data support is recovered.The on-support dynamics retain information needed for latent distribution recovery.
Gaussian example: Only the on-support score carries the latent covariance information and guides distribution recovery in the Gaussian example.In that example, the on-support score is linear in x, while the orthogonal score blows up near t = 0.
Regularity: The on-support score is Lipschitz continuous, a property used to establish approximation and estimation results.The paper identifies this regularity alongside the distinct component behavior.

4 Score approximation and estimation

The paper develops an encoder-decoder score-network architecture for low-dimensional subspace data, proving approximation and estimation guarantees whose key rates depend on intrinsic dimension. The analysis also exposes a tradeoff in the diffusion start time t0: smaller t0 worsens score estimation bounds, while larger t0 harms distribution recovery.

Architecture: The SNN score network uses an encoder-decoder structure with a shortcut connection and linear encoder-decoder maps.Its form is sV,θ(x, t) = 1 h(t)V fθ(V ⊤x, t) − 1 h(t)x, where V⊤ encodes and V decodes.
Score approximation: Theorem 1 establishes universal score-function approximation under the L2 norm over the unbounded domain RD.The proof uses truncation, while the architecture makes network size depend on intrinsic dimension d rather than ambient dimension D.
Score approximation: The architecture preserves approximation power while enforcing Lipschitz constraints on the score network.This regularity is essential for bounding distribution-recovery error later in the paper.
Score estimation: Theorem 2 gives an explicit sample-complexity guarantee for score matching, with convergence depending on intrinsic dimension d.For sufficiently large n, the squared L2 estimation error converges at rate e O(1 t0 n− 2 d+5 ), up to factors depending on β, log D, d, log t0, and τ.
Score estimation: The estimation bound creates a t0 tradeoff: it becomes vacuous as t0 approaches zero for fixed n, but larger t0 worsens distribution recovery.The first effect follows from the score-function blowup as t0 approaches zero.

5 Distribution estimation

The section establishes distribution-estimation guarantees for diffusion models with estimated scores, covering subspace recovery, latent-distribution error, and orthogonal behavior. The results expose a tradeoff in early stopping while showing that generated samples preserve the data geometry.

Theorem 3: Theorem 3 analyzes subspace recovery, latent distribution estimation, and generated-distribution behavior in the orthogonal space.The theorem uses an estimated score network and specifies choices of T and t0 under a covariance-eigenvalue condition.
Vanishing in the orthogonal space: The orthogonal pushforward of the continuous-time generated distribution is Gaussian with covariance bounded above by c t0 I.As t0 approaches zero, the orthogonal generated distribution degenerates toward a point mass at the origin.
Subspace recovery error: The learned subspace accurately matches the unknown data subspace, with error proportional to score-estimation error and dependent on the latent covariance's minimum eigenvalue.The latent distribution must span every direction of the true subspace for estimation.
Tradeoff on t0: The latent distribution error increases as t0 decreases because score-estimation error is amplified, while the Wasserstein bias decreases.Balancing these effects motivates setting t0 = n^(-1−δ(n))/(d+5).
Vanishing in the orthogonal space: The generated distribution therefore approximately vanishes in directions orthogonal to the recovered subspace, capturing the data's low-dimensional structure.This behavior follows from accurate subspace recovery and the orthogonal-process analysis.

6 Proof sketch of main results

The proof sketches construct a Lipschitz neural score approximator on a truncated domain, control statistical, truncation, and approximation errors, and compare learned and ground-truth backward processes. The distribution result follows by transferring score error to latent-process divergence.

Proof sketch of Theorem 1: The truncation argument controls the unbounded input domain using the sub-Gaussian tail of the forward-process distribution.Lipschitzness bounds the target score outside the hypercube, while the tail controls the resulting error.
Proof sketch of Theorem 1: Theorem 1 constructs a piecewise-linear score approximator on a bounded hypercube and sets it to zero outside the truncation region.The construction locally approximates the target using values at hypercube centers and trapezoid functions to detect regions.
Proof sketch of Theorem 2: Theorem 2 decomposes estimation error into finite-sample, truncation, and network-approximation terms, bounding them with concentration and approximation arguments.The empirical-loss analysis uses covering numbers, while the constructed approximator controls the approximation component.
Proof sketch of Theorem 3: The distribution proof compares learned and ground-truth latent backward processes and uses Girsanov's theorem to bound their KL divergence through drift errors.The continuous-time argument is supplemented by discretization error for the implemented backward process.
Proof sketch of Theorem 3: The latent score estimator is formed by projecting the network output into latent coordinates, enabling comparison with the latent backward process.The estimator is U^T fθ(Uz, t)−z, scaled by 1/h(t).

7 Conclusion and discussion

The paper shows that appropriately designed neural networks support accurate score approximation and estimation for low-dimensional linear-subspace data, with distribution-learning rates depending on intrinsic rather than ambient dimension. Its end-to-end scheme captures both the unknown subspace and the data distribution, while the analysis remains limited to the linear-subspace setting.

Conclusion: The paper establishes accurate score approximation and estimation, followed by efficient distribution learning from the estimated score.The estimation rate depends on intrinsic dimension and is free of the curse of ambient dimensionality.
Linear subspace assumption: The analysis assumes data lie in an unknown low-dimensional linear subspace, leaving more general manifold assumptions for future work.The authors emphasize that even the linear-subspace case requires non-trivial analysis.
End-to-end distribution learning: Unlike a PCA-like two-step pipeline, the proposed scheme learns the unknown linear structure and data distribution end to end.The paper presents this as a more realistic learning scheme for understanding diffusion-model behavior.

A Omitted proofs in Section 3

The omitted proofs derive the score decomposition and its regularity by exploiting latent coordinates and the Gaussian forward process. They show that the on-support score depends only on projected data while the orthogonal component becomes singular near zero noise.

Score decomposition: The on-support score depends only on the projected data A^T x, reducing its effective input dimension to the intrinsic dimension d.Writing z′ = A^T x allows the score analysis to proceed in latent coordinates.
Score decomposition: The forward-process marginal remains Gaussian, with conditional structure determined by the latent variable and the projected observation.This Gaussian representation supplies the density identities used in the score derivation.
Regularity: The proof verifies Lipschitz continuity by separately controlling spatial cases and the time derivative of the resulting score expression.The spatial bounds are combined before handling time regularity.

B.1 Proof of Theorem 1

The proof constructs a neural-network approximator for the score by restricting the unbounded input domain, using piecewise-linear partitioning, and controlling truncation and approximation errors. The resulting score approximation has an L2 error bound over the diffusion-time interval.

Approximation setup: The proof reduces score approximation to approximating g(z′, t) in the latent d-dimensional coordinates and then maps the result back through the subspace matrix.The constructed score approximator is ¯sV,θ = 1/h(t) A¯fθ(A⊤x, t) − 1/h(t)x.
Approximation setup: The unbounded latent input domain is truncated to a hypercube, with the network set to zero outside it while tail behavior controls the resulting L2 error.On S × [t0, T], the network achieves L∞ approximation; outside S, it is set to zero.
Neural approximation: A partition of unity built from coordinatewise trapezoid functions enables piecewise-linear approximation on the truncated space.The domain is partitioned into hypercubes and time intervals, with trapezoid functions used to form the partition.
Neural approximation: The approximating coordinate functions can be implemented by ReLU networks, and concatenating them yields the full vector-valued approximator.The construction also tracks Lipschitz continuity in the latent input and time variables.
Error bound: Choosing the truncation radius and network configuration produces the stated L2 approximation bound for the score over t ∈ [t0, T].The proof first bounds truncation and approximation terms, then substitutes the radius into the network parameters and score expression.

B.2 Proof of Theorem 2

The proof derives a generalization bound for the learned score network by decomposing estimation error into concentration, truncation, and approximation terms. Covering-number bounds for the subspace matrix and neural network yield the theorem’s intrinsic-dimension-dependent rate.

Error decomposition: The population loss is decomposed into concentration, truncation, and approximation components for the learned score network.The approximation component is explicitly separated into statistical error and approximation error.
Statistical error: The concentration term is controlled using a bounded induced loss class, empirical-process concentration, and a covering argument.The proof bounds the class uniformly, constructs an ι-cover, and applies a concentration lemma with high probability.
Truncation error: The truncation term is controlled by choosing a radius R so that sub-Gaussian or Gaussian tails make the omitted region sufficiently unlikely.The same radius also ensures the observed samples lie inside the truncation region with high probability.
Approximation error: The score approximation term is bounded using the constructed approximator from Theorem 1, while the remaining objective-function difference is an absolute constant.This transfers the approximation result into the statistical estimation analysis.
Final bound: The final rate is obtained by covering both the orthonormal subspace matrix and the neural-network function class, then substituting these bounds into the oracle inequality.After omitting factors in d, β, τ, log D, and log t0, the resulting expression gives the bound in Theorem 2.

C.1 Subspace Error and Latent Score Matching Error

This section analyzes subspace mismatch and latent score matching, then establishes latent distribution-error bounds for undiscretized and discretized backward diffusions. The analysis uses score decompositions, Gaussian-process properties, and assumptions on latent regularity and tails.

Latent decomposition: The diffusion score is decomposed into components aligned with the unknown data subspace and its orthogonal complement, allowing distribution recovery to be analyzed in latent coordinates.The backward process is correspondingly decomposed into latent and orthogonal components.
Distribution recovery: The undiscretized and discretized backward SDEs receive separate latent distribution-error bounds under sub-Gaussian and Lipschitz assumptions.The discretized analysis additionally accounts for the step size through an extra discretization error term.
Distribution recovery: The analysis uses a Gaussian reference process whose terminal covariance is bounded by et0 in the undiscretized setting and by e^(t0 + η) in the discretized setting.These bounds characterize the initialization mismatch used in the distribution-recovery argument.
Subspace error: Subspace mismatch contributes an error controlled by the distance between the estimated and true subspaces and the second moment of the latent score-related quantity.The projection residual is bounded through orthonormal-matrix inequalities.
Latent score matching: The latent score and network are compared through an error decomposition involving representation mismatch, network approximation, and score matching terms.An orthonormal alignment matrix is introduced to compare the estimated and true latent representations.

D.2 Proof of Lemma 4, Undiscretized Setting

The proof establishes distributional control for the backward SDE by verifying Novikov’s condition, applying Girsanov’s theorem, and handling initialization and discretization errors. Gaussian-process calculations provide the required moment and covariance bounds.

Undiscretized setting: Novikov’s condition is verified for the undiscretized backward process using Lipschitz score bounds and sub-Gaussian initialization.The proof controls exponential moments involving the process and its Brownian component.
Path-space comparison: Girsanov’s theorem converts the score discrepancy into a KL-divergence bound between the laws of the ground-truth and learned backward paths.The resulting path-space comparison is then related to marginal distribution error.
Initialization error: Different initial distributions are handled through marginalization, data processing, Pinsker’s inequality, and exponential convergence of the Ornstein–Uhlenbeck process.This separates path-law discrepancy from the error caused by starting the learned process from a Gaussian distribution.
Discretized setting: The discretized setting satisfies an analogous Novikov condition, but its analysis includes an additional error term proportional to the discretization step.The proof explicitly identifies this term as ϵdis(T − t0).
Gaussian reference process: The Gaussian reference process remains zero-mean with covariance bounded by et0 in the continuous analysis and e^(t0 + η) after discretization.The discretized covariance follows from a variance recursion over the time steps.