Diffusion Models Learn Low-Dimensional Distributions via Subspace Clustering

## Abstract

Abstract Recent empirical studies have demonstrated that diffusion models can effectively learn the image distribution and generate new samples. Remarkably, these models can achieve this even with a small number of training samples despite a large image dimension, circumventing the curse of dimensionality. In this work, we provide theoretical insights into this phenomenon by leveraging key empirical observations: (i) the low intrinsic dimensionality of image data, (ii) a union of manifold structure of image data, and (iii) the low-rank property of the denoising autoencoder in trained diffusion models. These observations motivate us to assume the underlying data distribution of image data as a mixture of low-rank Gaussians and to parameterize the denoising autoencoder as a low-rank model according to the score function of the assumed distribution. With these setups, we rigorously show that optimizing the training loss of diffusion models is equivalent to solving the canonical subspace clustering problem over the training samples. Based on this equivalence, we further show that the minimal number of samples required to learn the underlying distribution scales linearly with the intrinsic dimensions under the above data and model assumptions. This insight sheds light on why diffusion models can break the curse of dimensionality and exhibit the phase transition in learning distributions. Moreover, we empirically establish a correspondence between the subspaces and the semantic representations of image data, facilitating image editing. We validate these results with corroborated experimental results on both simulated distributions and image datasets. The code is available at https://github.com/huijieZH/Diffusion-Model-Generalizability .

## 1 Introduction

Generative modeling is a fundamental task in deep learning, which aims to learn a data distribution from training data to generate new samples. Recently, diffusion models have emerged as a new family of generative models, demonstrating remarkable performance across diverse domains, including image generation [, , , ], video content generation [, ], speech and audio synthesis [, ], fine-tuning [, , ]
and solving inverse problem [, , , ]. In general, diffusion models learn a data distribution from training samples through a process that imitates the non-equilibrium thermodynamic diffusion process [, , ]. Specifically, the training and sampling of diffusion models involve two stages: (i) a forward diffusion process where Gaussian noise is incrementally added to training samples at each time step, and (ii) a backward sampling process where the noise is progressively removed through a neural network that is trained to approximate the score function at all time steps. As described in prior works [, ], the generative capability of diffusion models lies in their ability to learn the score function of the data distribution, i.e., the gradient of the logarithm of the probability density function (*pdf*). We refer the reader to [, , ] for a more comprehensive introduction and survey on diffusion models.

Despite the significant empirical success and recent advances in understanding sampling convergence [, , ], distribution approximation [, ], memorization [, , , ], and generalization [, , ] of diffusion models, the fundamental working mechanisms remain poorly understood. One of the key questions is

*When and why can diffusion models learn the underlying data distribution without suffering from the curse of dimensionality?*

At first glance, the answer might seem quite straightforward. If the diffusion model can learn the empirical distribution of the training data that closely approximates the underlying data distribution, then the puzzle is solved! However, it has been shown in [] that the number of samples for an empirical distribution to approximate the underlying data distribution across all time steps could grow exponentially with respect to (*w.r.t.*) the data dimension, leading to the curse of dimensionality. In contrast, recent empirical results in [, ] showed that the number of training samples for a diffusion model to learn the underlying distribution is much smaller than the worst-case scenario, breaking the curse of dimensionality. Therefore, there is a significant gap between theory and practice.

In this work, we aim to address the above question of learning the underlying distribution via diffusion models by leveraging low-dimensional models. Our key observations are as follows: (i) The intrinsic dimensionality of real image data is significantly lower than the ambient dimension, a fact well-supported by extensive empirical evidence [, , ]; (ii) Image data lies on a disjoint union of manifolds of varying intrinsic dimensions, as empirically verified in [, ]; (iii) We empirically observe that the denoising autoencoder [, , ] of diffusion models trained on real-world image datasets exhibit low-rank structures (see Figure 3). Based on these observations, we conduct a theoretical investigation of distribution learning through diffusion models by assuming that (i) the underlying data distribution is a mixture of low-rank Gaussians (MoLRG, see Definition 1), and (ii) the denoising autoencoder is parameterized according to the score function of the MoLRG. These assumptions have been carefully examined and validated by our experiments on real image datasets.

### 1.1 Our Contributions

In this work, our contributions can be summarized as follows:

- •
Equivalence between training diffusion models and subspace clustering. Under the above setup, we show that the training loss of diffusion models based on the denoising autoencoder is equivalent to the *unsupervised* subspace clustering problem [, , ] (see Theorem 3). This equivalence implies that training diffusion models is essentially learning low-dimensional manifolds of the data distribution.
- •
Understanding phase transition from failure to success in learning distributions. By leveraging the above equivalence and the data model, we further show that if the number of samples exceed the intrinsic dimension of the subspaces, the optimal solutions of the training loss can recover the underlying distribution. This explains why diffusion models can break the curse of dimensionality. Conversely, if the number of samples is insufficient, it may learn an incorrect distribution. This highlights a phase transition from failure to success in learning the underlying distributions as the sample size increases.
- •
Correspondence between semantic representations and the subspaces. Interestingly, we find that the discovered low-dimensional subspaces in pre-trained diffusion models possess semantic meanings for natural images (see Figure 1). This motivates us to propose a training-free method to edit images on a frozen-trained diffusion model.

We also conduct extensive numerical experiments on both synthetic and real-world data sets to verify our assumptions and validate our theory. More broadly, the theoretical insights we gained in this work provide practical guidance as follows. First, we have shown that the number of samples for learning the underlying distribution via diffusion models scales proportionally with its intrinsic dimension. This insight allows us to improve training efficiency by quantifying the number of required training samples. Second, the identified correspondence between semantic representations and subspaces provides valuable guidance on controlling data generation. By manipulating the semantic representations within these subspaces, we can achieve more precise and targeted data generation.

### 1.2 Related Works

Now, we review recent works on diffusion models closely related to our study and discuss their connections to our work.

Figure: Figure 1: Correspondence between the singular vectors of the Jacobian of the DAE and semantic image attributes.
Refer to caption: x1.png

#### Learning a mixture of Gaussians via diffusion models.

Recent works have extensively studied distribution learning and generalizability of diffusion models for learning a mixture of full-rank Gaussian (MoG) model [, , , , ]. Specifically, they assumed that there exist centers $\bm{\mu}_{1},\dots,\bm{\mu}_{K}\in\mathbb{R}^{n}$ such that image data approximately follows from the following distribution:

$$ $\displaystyle\bm{x}\sim\sum_{k=1}^{K}\pi_{k}\mathcal{N}(\bm{\mu}_{k},\bm{I}_{n }),$ (1) $$

where $\pi_{k}\geq 0$ is the mixing proportion of the $k$-th mixture component satisfying $\sum_{k=1}^{K}\pi_{k}=1$. Shah et al. [] showed that gradient descent for optimizing denoising diffusion probabilistic objective can efficiently locate the ground-truth centers of MoG. Later, Chen et al. [] extended the above analysis without assuming the separation condition on the underlying mixture components. Gatmiry et al. [] employed diffusion models to learn the MoG model and provided a time and sample complexity analysis to achieve a targeted total variation error. Recently, Wu et al. [] provided a theoretical analysis of diffusion guidance in the MoG model. Moreover, Cole and Lu [] analyzed the approximation and generalization of score-based diffusion models for learning sub-Gaussian distributions, which includes a mixture of Guassians as an example.

Notably, the MoLRG model studied in our work is distinct from the above MoG model that is widely studied in the literature. Specifically, the MoG model consists of multiple Gaussians with varying means and covariance spanning the full-dimensional space (see Equation 1), while a MoLRG comprises multiple Gaussians with zero mean and low-rank covariance (see Equation 7), lying in a union of low-dimensional subspaces. As such, the MoLRG model, inspired by the inherent low-dimensionality of image datasets [, , ], offers a deeper insight into how diffusion models can learn underlying distributions in practice without suffering from the curse of dimensionality. Along the direction of studying diffusion models for low-dimensional data, Chen et al. [] studied a general low-dimensional data distribution and showed that the underlying distribution can be recovered via diffusion models. In contrast, we specifically studied a mixture of low-dimensional Gaussian distributions, enabling us to obtain a sharp bound for learning the underlying distribution.

#### Memorization and generalization in diffusion models.

Recently, extensive studies [, , ] empirically revealed that diffusion models learn the score function across two distinct regimes — memorization (i.e., learning the empirical distribution) and generalization (i.e., learning the underlying distribution) — depending on the training dataset size vs. the model capacity. Notably, for a model with a fixed number of parameters, there is a phase transition from memorization to generalization as the number of training samples increases [, ]. In the memorization regime, recent works [, ] showed that diffusion models learn the empirical optimal denoiser when the network is highly over-parameterized. Yoon et al. [] argued that diffusion models tend to generalize when they fail to memorize training data. In the generalization regime, extensive works [, , ] studied the generalization performance of diffusion models in different settings. Interestingly, Kadkhodaie et al. [] observed that two distinct diffusion models can learn nearly identical score functions even if they are trained on sufficiently large separate, non-overlapping portions of the same dataset so that they can generate almost the same output from the same pure noise. Notably, most existing studies on the memorization and generalization of diffusion models are empirical. In contrast, our work provides rigorous theoretical explanations for these intriguing experimental observations based on the MoLRG model. We demonstrate that diffusion models learn the underlying data distribution with the number of training samples scaling linearly with the intrinsic dimension of the data distribution. Our theory reveals a phase transition from failure to success in learning the underlying distribution as the number of training samples increases, which sheds light on the phase transition from memorization to generalization.

#### Subspace clustering.

Subspace clustering is a fundamental problem in unsupervised learning, which aims to identify and group data points that lie in a union of low-dimensional subspaces in a high-dimensional space [, , ]. Over the past years, a substantial body of literature has explored various approaches to the algorithmic development and theoretical analysis of subspace clustering. These include techniques such as sparse representation [, , ], low-rank representation [, , ], and spectral clustering [, ]. In this work, we present a new interpretation of diffusion models from the perspective of subspace clustering. This is the first time that diffusion models have been analyzed through this lens, offering new insights into how these models can effectively learn complex data distributions by leveraging the intrinsic low-dimensional subspaces within the data.

#### Notation.

We write matrices in bold capital letters like $\bm{A}$, vectors in bold lower-case letters like $\bm{a}$, and scalars in plain letters like $a$. Given a matrix, we use $\|\bm{A}\|$ to denote its largest singular value (i.e., spectral norm), $\sigma_{i}(\bm{A})$ its $i$-th largest singular value, and $a_{ij}$ its $(i,j)$-th entry, $\mathrm{rank}(\bm{A})$ its rank, $||\bm{A}||_{F}$ its Frobenius norm. Given a vector $\bm{a}$, we use $\|\bm{a}\|$ to denote its Euclidean norm and $a_{i}$ its $i$-th entry. Let $\mathcal{O}^{n\times d}$ denote the set of all $n\times d$ orthonromal matrices. We simply write the score function $\nabla_{\bm{x}}\log p(\bm{x})$ of a distribution with probability density function (pdf) $p(\bm{x})$ as $\nabla\log p(\bm{x})$.

#### Organization.

In Section 2, we introduce the preliminaries on diffusion models and some assumptions on the data and model. In Section 3, we present the main results of this paper. In Section 4, we conduct numerical experiments to support our theory and demonstrate its practical implications. Finally, in Section 5, we summarize our work and discuss potential directions for future research.

## 2 Problem Setup

In this work, we consider an image dataset consisting of samples $\{\bm{x}^{(i)}\}_{i=1}^{N}\subseteq\mathbb{R}^{n}$, where each data point is *i.i.d.* sampled from an underlying data distribution $p_{\rm data}(\bm{x})$. Instead of learning this pdf directly, score-based diffusion models aim to learn the score of this distribution from the training samples.

### 2.1 Preliminaries on Score-Based Diffusion Models

#### Forward and reverse SDEs of diffusion models.

In general, diffusion models consist of forward and reverse processes indexed by a continuous time variable $t\in[0,1]$. Specifically, the forward process progressively injects noise into the data. This process can be described by the following stochastic differential equation (SDE):

$$ $\displaystyle\mathrm{d}\bm{x}_{t}=f(t)\bm{x}_{t}\mathrm{d}t+g(t)\mathrm{d}\bm{ w}_{t},$ (2) $$

where $\bm{x}_{0}\sim p_{\rm data}$, the scalar functions $f(t),g(t):\mathbb{R}\to\mathbb{R}$ respectively denote the drift and diffusion coefficients,(^1^11In general, the functions $f(t)$ and $g(t)$ are chosen such that (i) $\bm{x}_{t}$ for all $t$ close to 0 0 approximately follows the data distribution $p_{\rm data}$ and (ii) $\bm{x}_{t}$ for all $t$ close to $1$ is nearly a standard Gaussian distribution; see, e.g., the settings in [, , ].) and $\{\bm{w}_{t}\}_{t\in[0,1]}$ is the standard Wiener process. For ease of exposition, let $p_{t}(\bm{x})$ denote the pdf of $\bm{x}_{t}$ and $p_{t}(\bm{x}_{t}|\bm{x}_{0})$ the transition kernel from $\bm{x}_{0}$ to $\bm{x}_{t}$. According to (2), we have

$$ $\displaystyle p_{t}(\bm{x}_{t}|\bm{x}_{0})=\mathcal{N}(\bm{x}_{t};s_{t}\bm{x}_ {0},s^{2}_{t}\sigma^{2}_{t}\bm{I}_{n}),\ \text{where}\ s_{t}=\exp\left(\int_{0 }^{t}f(\xi)\mathrm{d}\xi\right),\sigma_{t}=\sqrt{\int_{0}^{t}\frac{g^{2}(\xi)} {s^{2}({\xi})}\mathrm{d}\xi},$ (3) $$

where $s_{t}:=s(t)$ and $\sigma_{t}:=\sigma(t)$ for simplicity. The reverse process gradually removes the noise from $\bm{x}_{1}$ via the following reverse-time SDE:

$$ $\displaystyle\mathrm{d}\bm{x}_{t}=\left[f(t)\bm{x}_{t}-g^{2}(t)\nabla\log p_{t }(\bm{x}_{t})\right]\mathrm{d}t+g(t)\mathrm{d}\bar{\bm{w}}_{t},$ (4) $$

where $\{\bar{\bm{w}_{t}}\}_{t\in[0,1]}$ is another standard Wiener process, independent of $\{\bm{w}_{t}\}$, running backward in time from $t=1$ to $t=0$. It is worth noting that if $\bm{x}_{1}$ and $\nabla\log p_{t}$ are provided, the reverse process has exactly the same distribution as the forward process at each time $t\geq 0$ [].

#### Training loss of diffusion models.

Unfortunately, the score function $\nabla\log p_{t}$ is usually unknown, as it depends on the underlying data distribution $p_{\rm data}$. To enable data generation via the reverse SDE (4), a common approach is to estimate the score function $\nabla\log p_{t}$ using the training samples $\{\bm{x}^{(i)}\}_{i=1}^{N}$ based on the scoring matching [, ]. Because of
the equivalence between the score function $\nabla\log p_{t}(\bm{x}_{t})$ and the posterior mean $\mathbb{E}\left[\bm{x}_{0}|\bm{x}_{t}\right]$, i.e.,

$$ $\displaystyle s_{t}\mathbb{E}\left[\bm{x}_{0}|\bm{x}_{t}\right]=\bm{x}_{t}+s_{ t}^{2}\sigma_{t}^{2}\nabla\log p_{t}(\bm{x}_{t}),$ (5) $$

according to Tweedie’s formula and (3), an alternative approach to estimate the score function $\nabla\log p_{t}$ is to estimate the posterior mean $\mathbb{E}\left[\bm{x}_{0}|\bm{x}_{t}\right]$. Consequently, extensive works [, , , , ] have considered training a time-dependent function $\bm{x}_{\bm{\theta}}(\cdot,t):\mathbb{R}^{n}\times[0,1]\to\mathbb{R}^{n}$, known as denoising autoencoder (DAE), parameterized by a neural network with parameters $\bm{\theta}$ to estimate the posterior mean $\mathbb{E}\left[\bm{x}_{0}|\bm{x}_{t}\right]$. To determine the parameters $\bm{\theta}$, we can minimize the following empirical loss:

$$ $\displaystyle\min_{\bm{\theta}}\ \ell(\bm{\theta}):=\frac{1}{N}\sum_{i=1}^{N} \int_{0}^{1}\lambda_{t}\mathbb{E}_{\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I}_ {n})}\left[\left\|\bm{x}_{\bm{\theta}}(s_{t}\bm{x}^{(i)}+\gamma_{t}\bm{ \epsilon},t)-\bm{x}^{(i)}\right\|^{2}\right]\mathrm{d}t,$ (6) $$

where $\lambda_{t}:[0,1]\to\mathbb{R}^{+}$ is a weighting function and $\gamma_{t}:=s_{t}\sigma_{t}$. As shown in [], training the DAE is equivalent to performing explicit or implicit score matching under mild conditions. We refer the reader to Section A.1 for the relationship between this loss and the score-matching loss in [, ].

### 2.2 Low-Dimensional Data and Models

#### Mixture of low-rank Gaussian data distribution.

Although real-world image datasets are high dimensional in terms of pixel count and overall data volume, extensive empirical works [, , , ] suggest that their intrinsic dimensions are much lower. For instance, Pope et al. [] employed a kernelized nearest neighbor method to estimate the intrinsic dimensionality of various datasets, including MNIST [], CIFAR-10 [], and ImageNet []. Their findings indicate that even for complex datasets like ImageNet, the intrinsic dimensionality is approximately $40$, which is significantly lower than its ambient dimension.

Figure: Figure 2: An illustration of training samples that are generated according to the MoLRG model.
Refer to caption: x2.png

Recently, Brown et al. [], Kamkari et al. [] empirically validated the union of manifolds hypothesis, demonstrating that high-dimensional image data often lies on a disjoint union of manifolds instead of a single manifold. Notably, the union of subspaces, a special case of this hypothesis, has been widely used in the subspace clustering literature [, , ]. These observations motivate us to model the underlying data distribution as a mixture of low-rank Gaussians, where the data points are generated from a mixture of several Gaussian distributions with zero mean and low-rank covariance matrices. We formally define the MoLRG distribution as follows:

###### Definition 1 (Mixtures of Low-Rank Gaussians) .

We say that a random vector $\bm{x}\in\mathbb{R}^{n}$ follows a mixture of $K$ low-rank Gaussian distribution with parameters $\{\pi_{k}\}_{k=1}^{K}$ and $\{\bm{U}_{k}^{\star}\}_{k=1}^{K}$ if

$$ $\displaystyle\bm{x}\sim\sum_{k=1}^{K}\pi_{k}\mathcal{N}(\bm{0},\bm{U}_{k}^{ \star}\bm{U}_{k}^{\star T}),$ (7) $$

where $\bm{U}_{k}^{\star}\in\mathcal{O}^{n\times d_{k}}$ denotes the orthonormal basis of the $k$-th component and $\pi_{k}\geq 0$ is the mixing proportion of the $k$-th mixture component satisfying $\sum_{k=1}^{K}\pi_{k}=1$.

Before we proceed, we make some remarks on this data model. First, to study how diffusion models learn the underlying data distribution, many recent works have studied a mixture of full-rank Gaussian distributions (see Eq. (1)); see, e.g., [, , , , ]. However, compared to this model, a MoLRG is a more suitable model for capturing the low-dimensionality in image data distribution. Second, Brown et al. [], Kamkari et al. [] conducted extensive numerical experiments to validate that image datasets such as MNIST and ImageNet approximately lie on a union of low-dimensional manifolds. Because a nonlinear manifold can be well approximated by its tangent space (i.e., a linear subspace) in a local neighborhood, the MoLRG model, which represents data as a union of linear subspace, serves a good local approximation of a union of manifolds. Finally, assuming Gaussian distributions in each subspace in the MoLRG model is to enesure theoretical tractability while approximating the real-world image distributions, making it a practical starting point for theoretical studies on real-world image datasets. Now, we compute the ground-truth posterior mean $\mathbb{E}\left[\bm{x}_{0}|\bm{x}_{t}\right]$ when $\bm{x}_{0}$ satisfies the MoLRG model as follows.

###### Lemma 1 .

Suppose that $\bm{x}_{0}$ satisfies the MoLRG model. For each time $t>0$, it holds that

$$ $\displaystyle\mathbb{E}\left[\bm{x}_{0}|\bm{x}_{t}\right]=\dfrac{s_{t}}{s_{t}^ {2}+\gamma_{t}^{2}}\frac{\sum_{k=1}^{K}\exp\left(\phi_{t}\|\bm{U}_{k}^{\star T }\bm{x}_{t}\|^{2}\right)\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\bm{x}_{t}}{\sum _{k=1}^{K}\exp\left(\phi_{t}\|\bm{U}_{k}^{\star T}\bm{x}_{t}\|^{2}\right)},\ \text{where}\ \phi_{t}:=\frac{s_{t}^{2}}{2\gamma_{t}^{2}(s_{t}^{2}+\gamma_{t}^ {2})}.$ (8) $$

We defer the proof of this lemma to Section A.2. Notably, this lemma shows that the ground-truth posterior mean is a convex combination of projections of $\bm{x}_{t}$ onto each subspace, where the weights are soft-max functions. Moreover, this lemma provides guidance on the network parameterization of the DAE $\bm{x}_{\bm{\theta}}(\cdot,t)$ as discussed below.

Figure: (a) Real image datasets
Refer to caption: extracted/5832094/figure/jacobian-dataset-edm.png

###### Definition 1 (Mixtures of Low-Rank Gaussians) .

###### Lemma 1 .

#### Low-rank network parameterization.

In this work, we empirically observed that the DAE $\bm{x}_{\bm{\theta}}(\cdot,t)$ trained on real-world image datasets exhibits a low-dimensional structure. Specifically, when we train diffusion models with the U-Net architecture [] on various image datasets, it is observed that the numerical rank of the Jacobian of the DAE, i.e., $\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)$, is substantially lower than the ambient dimension in most time steps; see Figure 3(a). Additionally, this pattern of low dimensionality appears to be consistent across different datasets with different noise levels $t$.

When training diffusion models with U-Net on the samples generated according to the MoLRG model, the Jacobian of the DAE also exhibits a similar low-rank pattern, as illustrated in Figure 3(b). For the theoretical study based upon MoLRG, the above observations motivate us to consider a low-rank parameterization of the network. According to the ground-truth posterior mean of the MoLRG  model in Lemma 1, a natural parameterization for the DAE is

$$ $\displaystyle\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)=\dfrac{s_{t}}{s_{t}^{2}+\gamma _{t}^{2}}\sum_{k=1}^{K}w_{k}(\bm{\theta};\bm{x}_{t})\bm{U}_{k}\bm{U}_{k}^{T} \bm{x}_{t},\ \text{where}\ w_{k}(\bm{\theta};\bm{x}_{t})=\frac{\exp\left(\phi_ {t}\|\bm{U}_{k}^{T}\bm{x}_{t}\|^{2}\right)}{\sum_{l=1}^{K}\exp\left(\phi_{t}\| \bm{U}_{l}^{T}\bm{x}_{t}\|^{2}\right)}$ (9) $$

and the network parameters $\bm{\theta}=\{\bm{U}_{k}\}_{k=1}^{K}$ satisfy $\bm{U}_{k}\in\mathcal{O}^{n\times d_{k}}$. Although this approach may seem idealized, it offers several practical insights:

- •
Resemblance to the U-Net architecture: If we consider a single low-rank Gaussian, the network parameterization for
approximating the score function takes the form $\bm{x}-s_{t}/(s_{t}^{2}+\gamma_{t}^{2})\bm{U}\bm{U}^{T}\bm{x}$, which resembles the structure of a practical U-Net with a linear encoder, decoder, and skip connections. This provides theoretical insights into why U-Net is preferred for training diffusion models.
- •
Capturing low-rank property of the DAE: As shown in Figure 3, the Jacobian of the DAE of diffusion models trained on real image data is low-rank for most time steps, and our parameterization captures this low-rank property.
- •
Providing practical guidance: The theoretical insights gained from this work offer practical guidance. First, to learn the underlying distribution, the number of samples should be proportional to its intrinsic dimension. In practice, this informs us on how to use a minimal number of samples to train diffusion models to achieve generalization. Second, in Section 4.2, we empirically showed that
the subspace basis has semantic meanings. Based upon this, our results shed light on how to control the data generation by manipulating the semantic representation within the subspaces.

Similar simplifications have been widely used for theoretical analysis in various ideal data distributions; see, e.g., [, , , ]. Notably, under this specific network parameterization in (9), learning the score function $\nabla\log p_{t}(\bm{x}_{t})$ reduces to learning the network parameters $\bm{\theta}$ in (9) according to Lemma 1 and Eq. (5).

## 3 Main Results

Based on the setups in Section 2.2, we are ready to conduct a theoretical analysis of distribution learning using diffusion models. Specifically, when the underlying data distribution $p_{\rm data}(\bm{x})$ is modeled as MoLRG and the DAE is approximately parameterized as Eq. (9), we show that

Note that $\min_{\bm{\theta}\in\bm{\Theta}}\ell(\bm{\theta})$ is equivalent to $\min_{\bm{\theta}\in\bm{\Theta}}f(\bm{\theta})$ when we have $f(\bm{\theta})=a\ell(\bm{\theta})+b$, where $a>0$ and $b$ are absolute constants that do not depend on $\bm{\theta}$. This equivalence applies to our results in both Theorem 1 and Theorem 3.

### 3.1 A Warm-Up Study: A Single Low-rank Gaussian Case

To begin, we start from a simple case that the underlying distribution $p_{\rm data}$ is a *single* low-rank Gaussian. Specifically, the training samples $\{\bm{x}^{(i)}\}_{i=1}^{N}\subseteq\mathbb{R}^{n}$ are generated according to

$$ $\displaystyle\bm{x}^{(i)}=\bm{U}^{\star}\bm{a}_{i}+\bm{e}_{i},$ (10) $$

where $\bm{U}^{\star}\in\mathcal{O}^{n\times d}$ denotes an orthonormal basis, $\bm{a}_{i}\overset{i.i.d.}{\sim}\mathcal{N}(\bm{0},\bm{I}_{d})$ is coefficients for each $i\in[N]$, and $\bm{e}_{i}\in\mathbb{R}^{n}$ is noise for all $i\in[N]$.(^2^22Since real-world images inherently contain noise due to various factors, such as sensor limitation, environment conditions, and transition error, it is reasonable to add a noise term to this model.) According to (9), we parameterize the DAE into

$$ $\displaystyle\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)=\frac{s_{t}}{s_{t}^{2}+\gamma_ {t}^{2}}\bm{U}\bm{U}^{T}\bm{x}_{t},$ (11) $$

where $\bm{\theta}=\bm{U}\in\mathcal{O}^{n\times d}$. Equipped with the above setup, we can show the following result.

###### Theorem 1 .

Suppose that the DAE $\bm{x}_{\bm{\theta}}(\cdot,t)$ in Problem (6) is parameterized into (11) for each $t\in[0,1]$. Then, Problem (6) is equivalent to the following principal component analysis (PCA) problem:

$$ $\displaystyle\max_{\bm{U}\in\mathbb{R}^{n\times d}}\sum_{i=1}^{N}\|\bm{U}^{T} \bm{x}^{(i)}\|^{2}\qquad\mathrm{s.t.}\quad\bm{U}^{T}\bm{U}=\bm{I}_{d}.$ (12) $$

We defer the proof to Section B.1. In the single low-rank Gaussian model, Theorem 1 demonstrates that training diffusion models with a DAE of the form (11) to learn this distribution is equivalent to performing PCA on the training samples. Leveraging this equivalence, we can further characterize the number of samples required for learning underlying distribution under the data model (10).

###### Theorem 2 .

Consider the setting of Theorem 1. Suppose that the training samples $\{\bm{x}^{(i)}\}_{i=1}^{N}$ are generated according to the noisy single low-rank Gaussian model defined in (10). Let $\hat{\bm{U}}$ denote an optimal solution of Problem (6). The following statements hold:

- i)
If $N\geq d$, it holds with probability at least $1-1/2^{N-d+1}-\exp\left(-c_{2}N\right)$ that any optimal solution $\hat{\bm{U}}$ satisfies
$\displaystyle\left\|\hat{\bm{U}}\hat{\bm{U}}^{T}-\bm{U}^{\star}\bm{U}^{\star T
}\right\|_{F}\leq\frac{c_{1}\sqrt{\sum_{i=1}^{N}\|\bm{e}_{i}\|^{2}}}{\sqrt{N}-
\sqrt{d-1}},$
(13)
where $c_{1},c_{2}>0$ are constants that depend polynomially only on the Gaussian moment.
- ii)
If $N<d$, there exists an optimal solution $\hat{\bm{U}}\in\mathcal{O}^{n\times d}$ such that with probability at least $1-1/2^{d-N+1}-\exp\left(-c_{2}^{\prime}d\right)$,
$\displaystyle\left\|\hat{\bm{U}}\hat{\bm{U}}^{T}-\bm{U}^{\star}\bm{U}^{\star T
}\right\|_{F}\geq\sqrt{2\min\{d-N,n-d\}}-\frac{c_{1}^{\prime}\sqrt{\sum_{i=1}^
{N}\|\bm{e}_{i}\|^{2}}}{\sqrt{d}-\sqrt{N-1}},$
(14)
where $c_{1}^{\prime},c_{2}^{\prime}>0$ are constants that depend polynomially only on the Gaussian moment.

###### Theorem 1 .

###### Theorem 2 .

#### Remark 1.

We defer the proof to Section B.2. Building on the equivalence in Theorem 1 and the DAE parameterization (11), Theorem 2 clearly shows a phase transition from failure to success in learning the underlying distribution as the number of training samples increases. This phase transition is further corroborated by our experiments in Figures 4(a) and 4(b).

- •
When the number of training samples is larger than the dimension of the subspace, i.e., $N\geq d$, any optimal solution $\hat{\bm{U}}$ recovers the basis of underlying subspace with an approximation error depending on the noise level. This, together with the fact that learning the underlying distributions corresponds to learning the parameters $\bm{U}$ under the parameterization (11), implies that optimizing the training loss of diffusion models can learn the underlying distribution. Conversely, when $N<d$, optimizing the training loss fails to learn the underlying distribution.
- •
Note that our theory cannot explain why diffusion models memorize training data (i.e., learning the empirical distribution). This is because the parameterization (11) is not as sufficiently over-parameterized as architectures like U-Net. We plan to explore this over-parameterized setting in future work to better understand how diffusion models achieve memorization and to extend our theoretical insights accordingly.

### 3.2 From Single Low-Rank Gaussian to Mixtures of Low-Rank Gaussians

In this subsection, we extend the above study to the MoLRG distribution. In particular, we consider a noisy version of the MoLRG model as defined Definition 1. Specifically, the training samples are generated by

$$ $\displaystyle\bm{x}^{(i)}=\bm{U}_{k}^{\star}\bm{a}_{i}+\bm{e}_{i}\;\text{with probability}\;\pi_{k},\;\ \forall i\in[N],$ (15) $$

where $\bm{U}_{k}^{\star}\in\mathcal{O}^{n\times d_{k}}$ denotes an orthonormal basis for each $k\in[K]$, $\bm{a}_{i}\overset{i.i.d.}{\sim}\mathcal{N}(\bm{0},\bm{I}_{d_{k}})$ is coefficients, and $\bm{e}_{i}\in\mathbb{R}^{n}$ is noise for each $i\in[N]$. As argued by [], image data lies on a disjoint union of manifolds. This motivates us to assume that the basis matrices of subspaces satisfy $\bm{U}_{k}^{\star T}\bm{U}_{l}^{\star}=\bm{0}$ for each $k\neq l$. To simplify our analysis, we assume that $d_{1}=\dots=d_{K}=d$ and the mixing weights satisfy $\pi_{1}=\dots=\pi_{K}=1/K$. Moreover, we consider a hard-max counterpart of Eq. (9) for the DAE parameterization as follows:

$$ $\displaystyle\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)=\dfrac{s_{t}}{s_{t}^{2}+\gamma _{t}^{2}}\sum_{k=1}^{K}\hat{w}_{k}(\bm{\theta},\bm{x}_{0})\bm{U}_{k}\bm{U}_{k} ^{T}\bm{x}_{t},$ (16) $$

where $\bm{\theta}=\{\bm{U}_{k}\}_{k=1}^{K}$ and the weights $\{\hat{w}_{k}(\bm{\theta};\bm{x}_{0})\}_{k=1}^{K}$ are set as

$$ $\displaystyle\hat{w}_{k}(\bm{\theta};\bm{x}_{0})=\begin{cases}1,\ &\text{if}\ k=k_{0},\\ 0,\ &\text{otherwise},\end{cases}$ (17) $$

where $k_{0}\in[K]$ is an index satisfying $\|\bm{U}_{k_{0}}^{T}\bm{x}_{0}\|\geq\|\bm{U}_{l}^{T}\bm{x}_{0}\|$ for all $l\neq k_{0}\in[K]$. We should point out that we use two key approximations here. First, the soft-max weights $\{w_{k}(\bm{\theta},\bm{x}_{t})\}$ in Eq. (9) are approximated by the hard-max weights $\{\hat{w}_{k}(\bm{\theta};\bm{x}_{0})\}_{k=1}^{K}$. Second, $\|\bm{U}_{k}^{T}\bm{x}_{t}\|$ is approximated by its expectation, i.e., $\mathbb{E}_{\bm{\epsilon}}[\|\bm{U}_{k}^{T}\bm{x}_{t}\|^{2}]=\mathbb{E}_{\bm{
\epsilon}}\left[\|\bm{U}_{k}^{T}(s_{t}\bm{x}_{0}+\gamma_{t}\bm{\epsilon})\|^{2
}\right]=s_{t}^{2}\|\bm{U}_{k}^{T}\bm{x}_{0}\|^{2}+\gamma_{t}^{2}d$. We refer the reader to Section C.1 for more details on these approximation. Now, we are ready to show the following theorem.

###### Theorem 3 .

Suppose that the DAE $\bm{x}_{\bm{\theta}}(\cdot,t)$ in Problem (6) is parameterized into (16) for each $t\in[0,1]$, where $\hat{w}_{k}(\bm{\theta},\bm{x}_{0})$ is defined in (17) for each $k\in[K]$. Then, Problem (6) is equivalent to the following subspace clustering problem:

$$ $\displaystyle\max_{\bm{\theta}}\frac{1}{N}\sum_{k=1}^{K}\sum_{i\in C_{k}(\bm{ \theta})}\|\bm{U}_{k}^{T}\bm{x}^{(i)}\|^{2}\qquad\mathrm{s.t.}\quad\left[\bm{U }_{1},\dots,\bm{U}_{K}\right]\in\mathcal{O}^{n\times dK},$ (18) $$

where $C_{k}(\bm{\theta}):=\left\{i\in[N]:\|\bm{U}_{k}^{T}\bm{x}^{(i)}\|\geq\|\bm{U}_
{l}^{T}\bm{x}^{(i)}\|,\ \forall l\neq k\right\}$ for each $k\in[K]$.

We defer the proof to Section C.2. When the DAE is parameterized into (16), Theorem 3 demonstrates that optimizing the training loss of diffusion models is equivalent to solving the subspace clustering problem [, , ]. Problem (18) seeks to maximize the sum of the squared norms of each point’s projection onto its assigned subspace, subject to the constraint $C_{k}(\bm{\theta})$. Moreover, the equivalence enables us to characterize the required minimum number of samples for learning the underlying MoLRG distribution.

###### Theorem 4 .

Consider the setting of Theorem 3. Suppose that the training samples $\{\bm{x}^{(i)}\}_{i=1}^{N}$ are generated by the MoLRG distribution in Definition 1. Suppose $d\gtrsim\log N$ and $\|\bm{e}_{i}\|\lesssim\sqrt{d/N}$ for all $i\in[N]$. Let $\{\hat{\bm{U}}_{k}\}_{k=1}^{K}$ denote an optimal solution of Problem (6) and $N_{k}$ denote the number of samples from the $k$-th Gaussian component. Then, the following statements hold:

- (i)
If $N_{k}\geq d$ for each $k\in[K]$, there exists a permutation $\Pi:[K]\to[K]$ such that with probability at least $1-2K^{2}N^{-1}-\sum_{k=1}^{K}\left(1/2^{N_{k}-d+1}+\exp\left(-c_{2}N_{k}\right
)\right)$ for each $k\in[K]$,
$\displaystyle\left\|\hat{\bm{U}}_{\Pi(k)}\hat{\bm{U}}_{\Pi(k)}^{T}-\bm{U}_{k}^
{\star}\bm{U}_{k}^{\star T}\right\|_{F}\leq\frac{c_{1}\sqrt{\sum_{i=1}^{N}\|
\bm{e}_{i}\|^{2}}}{\sqrt{N}_{k}-\sqrt{d-1}},$
(19)
where $c_{1},c_{2}>0$ are constants that depend polynomially only on the Gaussian moment.
- (ii)
If $N_{k}<d$ for some $k\in[K]$, there exists a permutation $\Pi:[K]\to[K]$ and $k\in[K]$ such that with probability at least $1-2K^{2}N^{-1}-\sum_{k=1}^{K}\left(1/2^{d-N_{k}+1}+\exp\left(-c_{2}^{\prime}N_
{k}\right)\right)$,
$\displaystyle\left\|\hat{\bm{U}}_{\Pi(k)}\hat{\bm{U}}_{\Pi(k)}^{T}-\bm{U}_{k}^
{\star}\bm{U}_{k}^{\star T}\right\|_{F}\geq\sqrt{2\min\{d-N_{k},n-d\}}-\frac{c
_{1}^{\prime}\sqrt{\sum_{i=1}^{N}\|\bm{e}_{i}\|^{2}}}{\sqrt{d}-\sqrt{N_{k}-1}},$
(20)
where $c_{1}^{\prime},c_{2}^{\prime}>0$ are constants that depend polynomially only on the Gaussian

###### Theorem 3 .

###### Theorem 4 .

#### Remark 2.

We defer the proof to Section C.3. We discuss the implications of our results below.

- •
*Phase transition in learning the underlying distribution.* This theorem demonstrates that when the number of samples in each subspace exceeds the dimension of the subspace and the noise is bounded, the optimal solution of the training loss (6) under the parameterization (16) can recover the underlying subspaces up to the noise level. Conversely, when the number of samples is insufficient, there exists an optimal solution that may recover wrong subspaces; see Figures 4(c) and 4(d).
- •
*Connections to the phase transition from memorization to generalization.* We should clarify the difference between the phase transition described in Theorems 2 & 4 and the phase transition from memorization to generalization. Our phase transition refers to the shift from failure to success of learning the underlying distribution as the number of training samples increase, whereas the latter concerns the shift from memorizing data to generalizing from it as the number of training samples increases. Nevertheless, our theory still sheds light on the minimal number of samples required for diffusion models to enter the generalized regime.
- •
*Semantic meanings of the subspaces.* More interestingly, we empirically discovered a correspondence between the subspace bases and the semantics of images, as demonstrated by the experiments in Section 4.2. Specifically, when we move $\bm{x}_{t}$ along the direction of a singular vector of the Jacobian of the DAE trained on real image data with step size, the corresponding semantics of the generated images changes proportionally to the step size; see Figures 6 and 9(a, c). This enables controlled image generation without the need for additional training.

## 4 Experiments & Practical Implications

In this section, we first investigate phase transitions of diffusion models in learning distributions under both theoretical and practical settings in Section 4.1. Next, we demonstrate the practical implications of our work by exploring the correspondence between low-dimensional subspaces and semantic representations for controllable image editing in Section 4.2. Detailed experimental setups are postponed to Appendix E.

### 4.1 Phase Transition in Learning Distributions

In this subsection, we conduct experiments on both synthetic and real datasets to study the phase transition of diffusion models in learning distributions.

#### Learning the MoLRG distribution with the theoretical parameterizations.

To begin, we optimize the training loss (6) with the theoretical parameterization (9), where the data samples are generated by the MoLRG distribution. First, we apply stochastic gradient descent (see Algorithm 1) to solve Problem (6) with the DAE parameterized as (9). For comparison, according to Theorem 1 (resp., Theorem 3), we apply a singular value decomposition (resp., subspace clustering []) to solve Problem (12) (resp, Problem (18)).

We conduct three sets of experiments, where the data samples are respectively generated according to the single low-rank Gaussian distribution (10) with $K=1$ and a mixture of low-rank Gaussian distributions (15) with $K=2,3$. In each set, we set the total dimension $n=48$ and let the subspace dimension $d$ and the number of training samples $N$ vary from $2$ to $8$ and $2$ to $15$ with increments of $1$, respectively. For every pair of $d$ and $N$, we generate 20 instances, run the above methods, and calculate the successful rate of recovering the underlying subspaces. Then, the simulation results are visualized in Figure 4 and Figure 8. It is observed that all these methods exhibit a phase transition from failure to success in learning the subspaces as the number of training samples increases, which supports the results in Theorems 2 and 4.

Figure: (a) PCA
Refer to caption: x3.png

#### Learning the MoLRG distribution with U-Net.

Next, we optimize the training loss (6) with parameterizing the DAE $\bm{x}_{\bm{\theta}}(\cdot,t)$ using U-Net, detailed experiment settings are in Section E.2. Since U-Net is highly over-parameterized with about 54 million parameters and is different from the posterior mean in Lemma 1,
we measure the generalization ability of U-Net via generalization (GL) score [] defined as follows:

$$ $\displaystyle\text{GL score}=\frac{\mathcal{D}(\bm{x}^{(i)}_{\text{gen}})}{ \mathcal{D}(\bm{x}_{\texttt{MoLRG}}^{(i)})},\quad\mathcal{D}(\bm{x}^{(i)}) \coloneqq\sum_{j=1}^{N}\min_{j\neq i}||\bm{x}^{(i)}-\bm{x}^{(j)}||,$ (21) $$

where $\{\bm{x}_{\texttt{MoLRG}}^{(i)}\}_{i=1}^{N}$ are samples generated from the MoLRG distribution and $\{\bm{x}^{(i)}_{\text{gen}}\}_{i=1}^{N}$ are new samples generated by the trained U-Net. Intuitively, $\mathcal{D}(\bm{x}_{\text{gen}}^{(i)})$ reflects the uniformity of samples in the space: its value is small when the generated samples cluster around the training data, while the value is large when generated samples disperse in the entire space. Therefore, the trained diffsion model is in memorization regime when $D(\bm{x}^{(i)}_{\text{gen}})\ll\mathcal{D}(\bm{x}^{(i)}_{\texttt{MoLRG}})$ and the GL score is close to 0, while it is in generalization regime when $D(\bm{x}^{(i)}_{\text{gen}})\approx\mathcal{D}(\bm{x}^{(i)}_{\texttt{MoLRG}})$ and the GL score is close to 1.

In the experiments, we generate the data samples using the MoLRG distribution with $K=2$, $n=48$, and $d_{k}\in\{3,4,5,6\}$. Then, we plot the GL score against the $N_{k}/d_{k}$ for each $d_{k}$ in Figure 5(a). It is observed that for a fixed $d_{k}$, the generalization performance of diffusion models improves as the number of training samples increases. Notably, for different values of $d_{k}$, the plot of the GL score against the $N_{k}/d_{k}$ remains approximately consistent. This observation indicates that the phase transition curve for U-Net learning the MoLRG distribution depends on the ratio $N_{k}/d_{k}$ rather than on $N_{k}$ and $d_{k}$ individually. When $N_{k}/d_{k}\approx 60$, GL score $\approx 1.0$ suggesting that U-Net generalizes when $N_{k}\geq 60d_{k}$. This linear relationship for the phase transition differs from $N_{k}\geq d_{k}$ in Theorem 4 due to training with U-Net instead of the optimal network parameterization in Eq. (9). Nevertheless, Theorem 2 and Theorem 4 still provide valuable insights into learning distributions via diffusion models by demonstrating a similar phase transition phenomenon and confirming a linear relationship between $N_{k}$ and $d_{k}$.

Figure: (a) MoLRG distribution
Refer to caption: extracted/5832094/figure/phase_transition_UNet_MoLRG.png

#### Learning real image data distributions with U-Net.

Finally, we train diffusion models using U-Net on real image datasets AFHQ, CelebA, FFHQ, and CIFAR-10. The detailed experiment settings are deferred to Section E.3. According to [], we define the generalization (GL) score on real-world image dataset as follows:

$$ $\displaystyle\text{GL score}:=1-\mathbb{P}\left(\max_{i\in[N]}\left[\mathcal{M }_{\text{SSCD}}(\bm{x},\bm{y}_{i})\right]>0.6\right).$ $$

Here, the SSCD similarity is first introduced in [] to measure the replication between image pair $(\bm{x}_{1},\bm{x}_{2})$, which is defined as follows:

$$ $\displaystyle\mathcal{M}_{\text{SSCD}}(\bm{x}_{1},\bm{x}_{2})=\dfrac{\text{ SSCD}(\bm{x}_{1})\cdot\text{SSCD}(\bm{x}_{2})}{||\text{SSCD}(\bm{x}_{1})||_{2} \cdot||\text{SSCD}(\bm{x}_{2})||_{2}}$ $$

where $\text{SSCD}(\cdot)$ represents a neural descriptor for copy detection of images. We empirically sample 10K initial noises to estimate the probability. Intuitively, GL score measures the dissimilarity between the generated sample $\bm{x}$ and all $N$ samples $\bm{y}_{i}$ from the training dataset $\{\bm{y}_{i}\}_{i=1}^{N}$. Higher GL score indicates stronger generalizability. For each data set, we train U-Net and plot the GL score against the number of training samples in Figure 5(b).

The phase transition in the real dataset is illustrated in Figure 5(b). As observed, the order in which the samples need to generalize follows the relationship: AFHQ $>$ CelebA $>$ FFHQ $\approx$ CIFAR-10. Additionally, from our previous observations in Figure 3, the relationship of the intrinsic dimensions for these datasets is: AFHQ $>$ FFHQ $>$ CelebA $\approx$ CIFAR-10. Both AFHQ and CelebA align well with our theoretical analysis, which indicates that more samples are required for the model to generalize as the intrinsic dimension increases.

### 4.2 Semantic Meanings of Low-Dimensional Subspaces

Figure: (a)
Refer to caption: x7.png

In this subsection, we conduct experiments to verify the correspondence between the low-dimensional subspaces of the data distribution and the semantics of images on real datasets. We denote the Jacobian of the DAE $\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)$ by $\bm{J}_{t}:=\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)\in\mathbb{R}
^{n\times n}$ and let $\bm{J}_{t}=\bm{U}\bm{\Sigma}\bm{V}^{T}$ be an singular value decomposition (SVD) of $\bm{J}_{t}$, where
$r=\mathrm{rank}(\bm{J}_{t})$, $\bm{U}=[\bm{u}_{1},\cdots,\bm{u}_{r}]\in\mathcal{O}^{n\times r}$, $\bm{V}=[\bm{v}_{1},\cdots,\bm{v}_{r}]\in\mathcal{O}^{n\times r}$, and $\bm{\Sigma}=\mathrm{diag}(\sigma_{1},\dots,\sigma_{r})$ with $\sigma_{1}\geq\dots\geq\sigma_{r}$ being the singular values. To validate the semantic meaning of the basis vectors $\bm{v}_{i}$, we vary the value of $\alpha$ from negative to positive and visualize the resulting changes in the generated images.

In the experiments, we use a pre-trained diffusion denoising probabilistic model (DDPM) [] on the MetFaces dataset []. We randomly select an image $\bm{x}_{0}$ from this dataset and use the reverse process of the diffusion denoising implicit model (DDIM) [] to generate $\bm{x}_{t}$ at $t=0.7T$, where $T$ denote the total number of time steps. We respectively choose the changed direction as the leading right singular vectors $\bm{v}_{1},\bm{v}_{3},\bm{v}_{4},\bm{v}_{5},\bm{v}_{6}$ and use $\tilde{\bm{x}}_{t}=\bm{x}_{t}+\alpha\bm{v}_{i}$ to generate new images with $\alpha\in[-4,4]$ shown in Figure 6. It is observed that these singular vectors enable different semantic edits in terms of gender, hairstyle, and color of the image. For comparison, we generate a random unit vector $\bm{s}$ and move $\bm{x}_{t}$ along the direction of $\bm{s}$, where the editing strength $\alpha$ is the same as the semantic edits column-wise. The results are shown in the last column of Figure 6. Moving along random directions provides minimal semantic changes in the generated images, indicating that the low-dimensional subspace spanned by $\bm{V}$ is non-trivial and corresponds to semantic meaningful image attributes. More experimental results can be found in Figure 9 in Section E.3.

## 5 Conclusion & Discussion

In this work, we studied the training loss of diffusion models to investigate when and why diffusion models can learn the underlying distribution without suffering from the curse of dimensionality. Motivated by extensive empirical observations, we assumed that the underlying data distribution is a mixture of low-rank Gaussians. Specifically, we showed that minimizing the training loss is equivalent to solving the subspace clustering problem under proper network parameterization. Based on this equivalence, we further showed that the optimal solutions to the training loss can recover the underlying subspaces when the number of samples scales linearly with the intrinsic dimensionality of the data distribution. Moreover, we established the correspondence between the subspaces and semantic representations of image data. Since our studied network parameterization is not sufficiently over-parameterized, a future direction is to extend our analysis to an over-parameterized case to fully explain the transition from memorization to generalization.

## Acknowledgements

P.W., H.Z., Z.Z., S.C., and Q.Q. acknowledge support from NSF CAREER CCF-2143904, NSF CCF-2212066, NSF CCF-2212326, NSF IIS 2312842, NSF IIS 2402950, a gift grant from KLA, and the MICDE Catalyst Grant. Y.M. acknowledges support from the joint Simons Foundation-NSF DMS grant #2031899, NSF IIS 2402951, the ONR grant N00014-22-1-2102, TBSI, InnoHK, and the University of Hong Kong. The authors also would like to thank Laura Balzano (UMich), Jeff Fessler (UMich), Huikang Liu (SJTU), Dogyoon Song (UC Davis), Liyue Shen (UMich), Rene Vidal (Upenn), and Zhihui Zhu (OSU) for stimulating discussions.

## Appendix A Proofs in Section 2

### A.1 Relation between Score Matching Loss and Denoiser Autoencoder Loss

To estimate $\nabla\log p_{t}(\bm{x})$, one can train a time-dependent score-based model $\bm{s}_{\bm{\theta}}(\bm{x},t)$ via minimizing the following objective []:

$$ $\displaystyle\min_{\bm{\theta}}\int_{0}^{1}\xi_{t}\mathbb{E}_{\bm{x}_{0}\sim p _{\rm data}}\mathbb{E}_{\bm{x}_{t}|\bm{x}_{0}}\left[\left\|\bm{s}_{\bm{\theta} }(\bm{x}_{t},t)-\nabla\log p_{t}(\bm{x}_{t}|\bm{x}_{0})\right\|^{2}\right] \mathrm{d}t,$ (25) $$

where $\xi_{t}:[0,1]\to\mathbb{R}^{+}$ is a positive weighting function. Let $\bm{x}_{\bm{\theta}}(\cdot,t):\mathbb{R}^{d}\times[0,1]\to\mathbb{R}^{d}$ denote a neural network parameterized by parameters $\bm{\theta}$ to approximate $\mathbb{E}[\bm{x}_{0}|\bm{x}_{t}]$.
According to the Tweedie’s formula (5), $\bm{s}_{\bm{\theta}}(\bm{x}_{t},t)=\left(s_{t}\bm{x}_{\bm{\theta}}(\bm{x}_{t},
t)-\bm{x}_{t}\right)/\gamma_{t}^{2}$ can be used to estimate score functions. Substituting this and $\nabla\log p_{t}(\bm{x}_{t}|\bm{x}_{0})=\left(s_{t}\bm{x}_{0}-\bm{x}_{t}\right
)/\gamma_{t}^{2}$ due to (3) yields

$$ $\displaystyle\min_{\bm{\theta}}$ $\displaystyle\int_{0}^{1}\xi_{t}\mathbb{E}_{\bm{x}_{0}\sim p_{\mathrm{data}}} \mathbb{E}_{\bm{x}_{t}|\bm{x}_{0}}\left[\left\|\frac{1}{\gamma_{t}^{2}}\left(s _{t}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)-\bm{x}_{t}\right)-\frac{1}{\gamma_{t}^{ 2}}\left(s_{t}\bm{x}_{0}-\bm{x}_{t}\right)\right\|^{2}\right]\mathrm{d}t$ $\displaystyle=$ $\displaystyle\int_{0}^{1}\frac{\xi_{t}}{s^{2}_{t}\sigma_{t}^{4}}\mathbb{E}_{ \bm{x}_{0}\sim p_{\mathrm{data}}}\mathbb{E}_{\bm{\epsilon}\sim\mathcal{N}(\bm{ 0},\bm{I}_{n})}\left[\left\|\bm{x}_{\bm{\theta}}(s_{t}\bm{x}_{0}+\gamma_{t}\bm {\epsilon},t)-\bm{x}_{0}\right\|^{2}\right]\mathrm{d}t,$ $$

where the equality follows from $\bm{x}_{t}=s_{t}\bm{x}_{0}+\gamma_{t}\bm{\epsilon}$ due to (3). Then, we obtain

$$ $\displaystyle\min_{\bm{\theta}}\int_{0}^{1}\lambda_{t}\mathbb{E}_{\bm{x}_{0} \sim p_{\mathrm{data}}}\mathbb{E}_{\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I}_ {n})}\left[\left\|\bm{x}_{\bm{\theta}}(s_{t}\bm{x}_{0}+\gamma_{t}\bm{\epsilon} ,t)-\bm{x}_{0}\right\|^{2}\right]\mathrm{d}t,$ (26) $$

where $\lambda_{t}={\xi_{t}}/({s^{2}_{t}\sigma_{t}^{4}})$. However, only data points $\{\bm{x}^{(i)}\}_{i=1}^{N}$ sampled from the underlying data distribution $p_{\rm data}$ are available in practice. Therefore, we study the following empirical counterpart of Problem (26) over the training samples, i.e., Problem (6). We refer the reader to [, Section 2.1] for more discussions on the denoising error of this problem.

### A.2 Proof of in Lemma 1

Assuming that the underlying data distribution follows a mixture of low-rank Gaussians as defined in Definition 1, we first compute the ground-truth score function as follows.

###### Proposition 1 .

Suppose that the underlying data distribution $p_{\rm data}$ follows a mixture of low-rank Gaussian distributions in Definition 1. In the forward process of diffusion models, the pdf of $\bm{x}_{t}$ for each $t>0$ is

$$ $\displaystyle p_{t}(\bm{x})=\sum_{k=1}^{K}\pi_{k}\mathcal{N}(\bm{x};\bm{0},s_{ t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\gamma_{t}^{2}\bm{I}_{n}),$ (27) $$

where $\gamma_{t}=s_{t}\sigma_{t}$.
Moreover, the score function of $p_{t}(\bm{x})$ is

$$ $\displaystyle\nabla\log p_{t}(\bm{x})=-\frac{1}{\gamma_{t}^{2}}\left(\bm{x}- \frac{s_{t}^{2}}{s_{t}^{2}+\gamma_{t}^{2}}\frac{\sum_{k=1}^{K}\pi_{k}\mathcal{ N}(\bm{x};\bm{0},s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\gamma_{t}^{2 }\bm{I}_{n})\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\bm{x}}{\sum_{k=1}^{K}\pi_{k }\mathcal{N}(\bm{x};\bm{0},s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+ \gamma_{t}^{2}\bm{I}_{n})}\right).$ (28) $$

###### Proof.

Let $Y\in\{1,\dots,K\}$ be a discrete random variable that denotes the value of components of the mixture model. Note that $\gamma_{t}=s_{t}\sigma_{t}$. It follows from Definition 1 that $\mathbb{P}(Y=k)=\pi_{k}$ for each $k\in[K]$. We first compute

$$ $\displaystyle p_{t}(\bm{x}|Y=k)$ $\displaystyle=\int p_{t}\left(\bm{x}|Y=k,\bm{a}_{k})\mathcal{N}(\bm{a}_{k};\bm {0},\bm{I}_{d_{k}}\right){\rm d}\bm{a}_{k}=\int p_{t}(\bm{x}|\bm{x}_{0}=\bm{U} _{k}^{\star}\bm{a}_{k})\mathcal{N}\left(\bm{a}_{k};\bm{0},\bm{I}_{d_{k}}\right ){\rm d}\bm{a}_{k}$ $\displaystyle=\int\mathcal{N}(\bm{x};s_{t}\bm{U}_{k}^{\star}\bm{a}_{k},\gamma_ {t}^{2}\bm{I}_{n})\mathcal{N}\left(\bm{a}_{k};\bm{0},\bm{I}_{d_{k}}\right){\rm d }\bm{a}_{k}$ $\displaystyle=\frac{1}{(2\pi)^{n/2}(2\pi)^{d_{k}/2}\gamma_{t}^{n}}\int\exp \left(-\frac{1}{2\gamma_{t}^{2}}\|\bm{x}-s_{t}\bm{U}_{k}^{\star}\bm{a}_{k}\|^{ 2}\right)\exp\left(-\frac{1}{2}\|\bm{a}_{k}\|^{2}\right){\rm d}\bm{a}_{k}$ $\displaystyle=\frac{1}{(2\pi)^{n/2}\gamma_{t}^{n}}\left(\frac{s_{t}^{2}+\gamma _{t}^{2}}{\gamma_{t}^{2}}\right)^{-d/2}\exp\left(-\frac{1}{2\gamma_{t}^{2}}\bm {x}^{T}\left(\bm{I}_{n}-\frac{s_{t}^{2}}{s_{t}^{2}+\gamma_{t}^{2}}\bm{U}_{k}^{ \star}\bm{U}_{k}^{\star T}\right)\bm{x}\right)$ $\displaystyle\ \int\frac{1}{(2\pi)^{d_{k}/2}}\left(\frac{\gamma_{t}^{2}}{s_{t} ^{2}+\gamma_{t}^{2}}\right)^{-d/2}\exp\left(-\frac{s_{t}^{2}+\gamma_{t}^{2}}{2 \gamma_{t}^{2}}\left\|\bm{a}_{k}-\frac{s_{t}}{s_{t}^{2}+\gamma_{t}^{2}}\bm{U}_ {k}^{\star T}\bm{x}\right\|^{2}\right){\rm d}\bm{a}_{k}$ $\displaystyle=\frac{1}{(2\pi)^{n/2}}\frac{1}{\left((s_{t}^{2}+\gamma_{t}^{2})^ {d}\gamma_{t}^{2(n-d)}\right)^{1/2}}\exp\left(-\frac{1}{2\gamma_{t}^{2}}\bm{x} ^{T}\left(\bm{I}_{n}-\frac{s_{t}^{2}}{s_{t}^{2}+\gamma_{t}^{2}}\bm{U}^{\star} \bm{U}^{\star T}\right)\bm{x}\right)$ $\displaystyle=\frac{1}{(2\pi)^{n/2}\det^{1/2}(s_{t}^{2}\bm{U}_{k}^{\star}\bm{U }_{k}^{\star T}+\gamma_{t}^{2}\bm{I}_{n})}\exp\left(-\frac{1}{2}\bm{x}^{T} \left(s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\gamma_{t}^{2}\bm{I}_{n} \right)^{-1}\bm{x}\right)$ $\displaystyle=\mathcal{N}(\bm{x};\bm{0},s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^ {\star T}+\gamma_{t}^{2}\bm{I}_{n}),$ $$

where the second equality follows from (3), the third equality uses (22), the fourth equality is due to the fact that $<\bm{x},\bm{U}_{k}^{\star}\bm{a}>$ is an odd function, and the second to last equality uses $\det(s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\gamma_{t}^{2}\bm{I}_{n})
=(s_{t}^{2}+\gamma_{t}^{2})^{d}\gamma_{t}^{2(n-d)}$ and $(s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\gamma_{t}^{2}\bm{I}_{n})^{-1
}=\left(\bm{I}_{n}-s_{t}^{2}/(s_{t}^{2}+\gamma_{t}^{2})\bm{U}_{k}^{\star}\bm{U
}_{k}^{\star T}\right)/\gamma_{t}^{2}$ due to the matrix inversion lemma and $\bm{U}_{k}^{\star T}\bm{U}_{k}^{\star}=\bm{I}_{d_{k}}$. This, together with $\mathbb{P}(Y=k)=\pi_{k}$ for each $k\in[K]$, yields

$$ $\displaystyle p_{t}(\bm{x})$ $\displaystyle=\sum_{k=1}^{K}p_{t}(\bm{x}|Y=k)\mathbb{P}(Y=k)=\sum_{k=1}^{K}\pi _{k}\mathcal{N}(\bm{x};\bm{0},s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+ \gamma_{t}^{2}\bm{I}_{n}).$ $$

Next, we directly compute

$$ $\displaystyle\nabla\log p_{t}(\bm{x})$ $\displaystyle=\frac{\nabla p_{t}(\bm{x})}{p_{t}(\bm{x})}=\frac{\sum_{k=1}^{K} \pi_{k}\mathcal{N}(\bm{x};\bm{0},s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T }+\gamma_{t}^{2}\bm{I}_{n})\left(-\frac{1}{\gamma_{t}^{2}}\bm{x}+\frac{s_{t}^{ 2}}{\gamma_{t}^{2}(s_{t}^{2}+\gamma_{t}^{2})}\bm{U}_{k}^{\star}\bm{U}_{k}^{ \star T}\bm{x}\right)}{\sum_{k=1}^{K}\pi_{k}\mathcal{N}(\bm{x};\bm{0},s_{t}^{2 }\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\gamma_{t}^{2}\bm{I}_{n})}$ $\displaystyle=-\frac{1}{\gamma_{t}^{2}}\left(\bm{x}-\frac{s_{t}^{2}}{s_{t}^{2} +\gamma_{t}^{2}}\frac{\sum_{k=1}^{K}\pi_{k}\mathcal{N}(\bm{x};\bm{0},s_{t}^{2} \bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\gamma_{t}^{2}\bm{I}_{n})\bm{U}_{k}^{ \star}\bm{U}_{k}^{\star T}\bm{x})}{\sum_{k=1}^{K}\pi_{k}\mathcal{N}(\bm{x};\bm {0},s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\gamma_{t}^{2}\bm{I}_{n})} \right).$ $$

∎

###### Proof of Lemma 1 .

According to (5) and Proposition 1, we compute

$$ $\displaystyle\mathbb{E}\left[\bm{x}_{0}|\bm{x}_{t}\right]$ $\displaystyle=\frac{\bm{x}_{t}+\gamma_{t}^{2}\nabla\log p_{t}(\bm{x}_{t})}{s_{ t}}=\frac{s_{t}}{s_{t}^{2}+\gamma_{t}^{2}}\frac{\sum_{k=1}^{K}\pi_{k}\mathcal{ N}(\bm{x};\bm{0},s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\gamma_{t}^{2 }\bm{I}_{n})\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\bm{x}_{t}}{\sum_{k=1}^{K} \pi_{k}\mathcal{N}(\bm{x}_{t};\bm{0},s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{ \star T}+\gamma_{t}^{2}\bm{I}_{n})}$ $\displaystyle=\dfrac{s_{t}}{s_{t}^{2}+\gamma_{t}^{2}}\frac{\sum_{k=1}^{K}\pi_{ k}\exp\left(-\dfrac{1}{2\gamma_{t}^{2}}\left(\|\bm{x}_{t}\|^{2}-\frac{s^{2}_{t }}{s_{t}^{2}+\gamma_{t}^{2}}\|\bm{U}_{k}^{\star T}\bm{x}_{t}\|^{2}\right) \right)\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\bm{x}_{t}}{\sum_{k=1}^{K}\pi_{k} \exp\left(-\dfrac{1}{2\gamma_{t}^{2}}\left(\|\bm{x}_{t}\|^{2}-\frac{s^{2}_{t}} {s_{t}^{2}+\gamma_{t}^{2}}\|\bm{U}_{k}^{\star T}\bm{x}_{t}\|^{2}\right)\right)}$ $\displaystyle=\dfrac{s_{t}}{s_{t}^{2}+\gamma_{t}^{2}}\frac{\sum_{k=1}^{K}\pi_{ k}\exp\left(\dfrac{1}{2\gamma_{t}^{2}}\dfrac{s^{2}_{t}}{s_{t}^{2}+\gamma_{t}^{ 2}}\|\bm{U}_{k}^{\star T}\bm{x}_{t}\|^{2}\right)\bm{U}_{k}^{\star}\bm{U}_{k}^{ \star T}\bm{x}_{t}}{\sum_{k=1}^{K}\pi_{k}\exp\left(\dfrac{1}{2\gamma_{t}^{2}} \dfrac{s^{2}_{t}}{s_{t}^{2}+\gamma_{t}^{2}}\|\bm{U}_{k}^{\star T}\bm{x}_{t}\|^ {2}\right)},$ $$

where the third equality uses (22) and $\left(s_{t}^{2}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\gamma_{t}^{2}\bm{I}_{n}
\right)^{-1}=\left(\bm{I}_{n}-s_{t}^{2}/(s_{t}^{2}+\gamma_{t}^{2})\bm{U}_{k}^{
\star}\bm{U}_{k}^{\star T}\right)/\gamma_{t}^{2}$ due to the matrix inversion lemma. ∎

###### Proposition 1 .

###### Proof.

###### Proof of Lemma 1 .

## Appendix B Proofs in Section 3.1

### B.1 Proof of Theorem 1

###### Proof of Theorem 1 .

Plugging (11) into the integrand of (6) yields

$$ $\displaystyle\ \mathbb{E}_{\bm{\epsilon}}\left[\left\|\frac{s_{t}}{s_{t}^{2}+ \gamma_{t}^{2}}\bm{U}\bm{U}^{T}\left(s_{t}\bm{x}^{(i)}+\gamma_{t}\bm{\epsilon} \right)-\bm{x}^{(i)}\right\|^{2}\right]$ $\displaystyle=$ $\displaystyle\ \left\|\frac{s_{t}^{2}}{s_{t}^{2}+\gamma_{t}^{2}}\bm{U}\bm{U}^{ T}\bm{x}^{(i)}-\bm{x}^{(i)}\right\|^{2}+\frac{(s_{t}\gamma_{t})^{2}}{(s_{t}^{2 }+\gamma_{t})^{2}}\mathbb{E}_{\bm{\epsilon}}\left[\|\bm{U}\bm{U}^{T}\bm{ \epsilon}\|^{2}\right]$ $\displaystyle=$ $\displaystyle\left\|\frac{s_{t}^{2}}{s_{t}^{2}+\gamma_{t}^{2}}\bm{U}\bm{U}^{T} \bm{x}^{(i)}-\bm{x}^{(i)}\right\|^{2}+\frac{(s_{t}\gamma_{t})^{2}d}{(s_{t}^{2} +\gamma_{t})^{2}},$ $$

where the first equality follows from $\mathbb{E}_{\bm{\epsilon}}[\langle\bm{x},\bm{\epsilon}\rangle]=0$ for any given $\bm{x}\in\mathbb{R}^{n}$ due to $\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I}_{n})$, and the second equality uses $\mathbb{E}_{\bm{\epsilon}}\left[\|\bm{U}\bm{U}^{T}\bm{\epsilon}\|^{2}\right]=
\mathbb{E}_{\bm{\epsilon}}\left[\|\bm{U}^{T}\bm{\epsilon}\|^{2}\right]=\sum_{i
=1}^{d}\mathbb{E}_{\bm{\epsilon}}\left[\|\bm{u}_{i}^{T}\bm{\epsilon}\|^{2}
\right]=d$ due to $\bm{U}\in\mathcal{O}^{n\times d}$ and $\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I}_{n})$.
This, together with $\gamma_{t}=s_{t}\sigma_{t}$ and (6), yields

$$ $\displaystyle\ell(\bm{U})=\frac{1}{N}\sum_{i=1}^{N}\int_{0}^{1}\lambda_{t} \left(\|\bm{x}^{(i)}\|^{2}-\frac{1+2\sigma_{t}^{2}}{(1+\sigma_{t}^{2})^{2}}\| \bm{U}^{T}\bm{x}^{(i)}\|^{2}+\frac{\sigma_{t}^{2}d}{(1+\sigma_{t}^{2})^{2}} \right){\rm d}t,$ $$

Obviously, minimizing the above function in terms of $\bm{U}$ amounts to

$$ $\displaystyle\min_{\bm{U}^{T}\bm{U}=\bm{I}_{d}}-\int_{0}^{1}\frac{(1+2\sigma_{ t}^{2})\lambda_{t}}{(1+\sigma_{t}^{2})^{2}}{\rm d}t\frac{1}{N}\sum_{i=1}^{N}\| \bm{U}^{T}\bm{x}^{(i)}\|^{2},$ $$

which is equivalent to Problem (12).
∎

###### Proof of Theorem 1 .

### B.2 Proof of Theorem 2

###### Proof of Theorem 2 .

For ease of exposition, let

$$ $\displaystyle\bm{X}=\begin{bmatrix}\bm{x}^{(1)}&\dots&\bm{x}^{(N)}\end{bmatrix }\in\mathbb{R}^{n\times N},\ \bm{A}=\begin{bmatrix}\bm{a}_{1}&\dots&\bm{a}_{N} \end{bmatrix}\in\mathbb{R}^{d\times N},\ \bm{E}=\begin{bmatrix}\bm{e}_{1}& \dots&\bm{e}_{N}\end{bmatrix}\in\mathbb{R}^{n\times N}.$ $$

Using this and (10), we obtain

$$ $\displaystyle\bm{X}=\bm{U}^{\star}\bm{A}+\bm{E}.$ (29) $$

Let $r_{A}:=\mathrm{rank}(\bm{A})\leq\min\{d,N\}$ and $\bm{A}=\bm{U}_{A}\bm{\Sigma}_{A}\bm{V}_{A}^{T}$ be an singular value decomposition (SVD) of $\bm{A}$, where $\bm{U}_{A}\in\mathcal{O}^{d\times r_{A}}$, $\bm{V}_{A}\in\mathcal{O}^{N\times r_{A}}$, and $\bm{\Sigma}_{A}\in\mathbb{R}^{r_{A}\times r_{A}}$.
It follows from Theorem 1 that Problem (6) with the parameterization (11) is equivalent to Problem (12).

(i) Suppose that $N\geq d$. Applying Lemma 3 with $\varepsilon={1}/{(2c_{1})}$ to $\bm{A}\in\mathbb{R}^{d\times N}$, it holds with probability at least $1-1/2^{N-d+1}-\exp\left(-c_{2}N\right)$ that

$$ $\displaystyle\sigma_{\min}(\bm{A})=\sigma_{d}(\bm{A})\geq\frac{\sqrt{N}-\sqrt{ d-1}}{2c_{1}},$ (30) $$

where $c_{1},c_{2}>0$ are constants depending polynomially only on the Gaussian moment. This implies $r_{A}=d$ and $\bm{U}_{A}\in\mathcal{O}^{d}$. Since Problem (12) is a PCA problem, the columns of any optimal solution $\hat{\bm{U}}\in\mathcal{O}^{n\times d}$ consist of left singular vectors associated with the top $d$ singular values of $\bm{X}$. This, together with Wedin’s Theorem [] and (29), yields

$$ $\displaystyle\left\|\hat{\bm{U}}\hat{\bm{U}}^{T}-\bm{U}^{\star}\bm{U}^{\star T }\right\|_{F}=\left\|\hat{\bm{U}}\hat{\bm{U}}^{T}-(\bm{U}^{\star}\bm{U}_{A})( \bm{U}^{\star}\bm{U}_{A})^{T}\right\|_{F}\leq\frac{2\|\bm{E}\|_{F}}{\sigma_{ \min}(\bm{A})}=\frac{4c_{1}\|\bm{E}\|_{F}}{\sqrt{N}-\sqrt{d-1}}.$ $$

This, together with absorbing $4$ into $c_{1}$, yields (13).

(ii) Suppose that $N<d$. According to Lemma 3 with $\varepsilon={1}/{(2c_{1})}$, it holds with probability at least $1-1/2^{d-N+1}-\exp\left(-c_{2}d\right)$ that

$$ $\displaystyle\sigma_{\min}(\bm{A})=\sigma_{N}(\bm{A})\geq\frac{\sqrt{d}-\sqrt{ N-1}}{2c_{1}},$ (31) $$

where $c_{1},c_{2}>0$ are constants depending polynomially only on the Gaussian moment. This implies $r_{A}=N$ and $\bm{U}_{A}\in\mathcal{O}^{d\times N}$. This, together with the fact that $\bm{A}=\bm{U}_{A}\bm{\Sigma}_{A}\bm{V}_{A}^{T}$ is an SVD of $\bm{A}$, yields that $\bm{U}^{\star}\bm{A}=\left(\bm{U}^{\star}\bm{U}_{A}\right)\bm{\Sigma}_{A}\bm{V
}_{A}^{T}$ is an SVD of $\bm{U}^{\star}\bm{A}$ with $\bm{U}^{\star}\bm{U}_{A}\in\mathcal{O}^{n\times N}$. Note that $\mathrm{rank}(\bm{X})\leq N$. Let $\bm{X}=\bm{U}_{X}\bm{\Sigma}_{X}\bm{V}_{X}^{T}$ be an SVD of $\bm{X}$, where $\bm{U}_{X}\in\mathcal{O}^{n\times N}$, $\bm{V}_{X}\in\mathcal{O}^{N}$, and $\bm{\Sigma}_{X}\in\mathbb{R}^{N\times N}$. This, together with Wedin’s Theorem [] and (31), yields

$$ $\displaystyle\left\|\bm{U}_{X}\bm{U}_{X}^{T}-\bm{U}^{\star}\bm{U}_{A}\bm{U}_{A }^{T}\bm{U}^{\star T}\right\|_{F}\leq\frac{2\|\bm{E}\|_{F}}{\sigma_{\min}(\bm{ A})}=\frac{4c_{1}\|\bm{E}\|_{F}}{\sqrt{d}-\sqrt{N-1}}.$ (32) $$

Note that Problem (12) has infinite optimal solutions when $N<d$, which take the form of

$$ $\displaystyle\hat{\bm{U}}=\begin{bmatrix}\bm{U}_{X}&\bar{\bm{U}}_{X}\end{ bmatrix}\in\mathcal{O}^{n\times d}.$ $$

Now, we consider that $\bar{\bm{U}}_{X}\in\mathcal{O}^{n\times(d-N)}$ is an optimal solution of the following problem:

$$ $\displaystyle\min_{\bm{V}\in\mathcal{O}^{n\times(d-N)},\bm{U}_{X}^{T}\bm{V}= \bm{0}}\|\bm{V}^{T}\bm{U}^{\star}(\bm{I}-\bm{U}_{A}\bm{U}_{A}^{T})\|_{F}^{2}.$ (33) $$

Then, one can verify that the rank of the following matrix is at most $d$:

$$ $\displaystyle\bm{B}:=\begin{bmatrix}\bm{U}_{X}&\bm{U}^{\star}(\bm{I}-\bm{U}_{A }\bm{U}_{A}^{T})\end{bmatrix}$ $$

Then, if $n\geq 2d-N$, it is easy to see that the optimal value of Problem (33) is 0 0. If $n<2d-N$, the optima value is achieved at $\bm{V}^{\star}=[\bm{V}_{1}^{\star}\ \bm{V}_{2}^{\star}]$ with $\bm{V}_{1}^{\star}\in\mathbb{R}^{n\times(n-d)}$ and $\bm{V}_{2}^{\star}\in\mathbb{R}^{n\times(2d-N-n)}$ satisfying $\bm{V}_{1}^{\star T}\bm{B}=\bm{0}$, which implies

$$ $\displaystyle\|\bm{V}^{\star T}\bm{U}^{\star}(\bm{I}-\bm{U}_{A}\bm{U}_{A}^{T}) \|_{F}^{2}=\|\bm{V}_{2}^{\star T}\bm{U}^{\star}(\bm{I}-\bm{U}_{A}\bm{U}_{A}^{T })\|_{F}^{2}\leq 2d-N-n.$ $$

Consequently, the optimal value of Problem (33) is less than

$$ $\displaystyle\max\left\{0,2d-(n+N)\right\}$ (34) $$

Then, we obtain that

$$ $\displaystyle\left\|\hat{\bm{U}}\hat{\bm{U}}^{T}-\bm{U}^{\star}\bm{U}^{\star T }\right\|_{F}$ $\displaystyle=\left\|\bm{U}_{X}\bm{U}_{X}^{T}+\bar{\bm{U}}_{X}\bar{\bm{U}}_{X} ^{T}-\bm{U}^{\star}\bm{U}_{A}\bm{U}_{A}^{T}\bm{U}^{\star T}-\bm{U}^{\star}(\bm {I}-\bm{U}_{A}\bm{U}_{A}^{T})\bm{U}^{\star T}\right\|$ $\displaystyle\geq\|\bar{\bm{U}}_{X}\bar{\bm{U}}_{X}^{T}-\bm{U}^{\star}(\bm{I}- \bm{U}_{A}\bm{U}_{A}^{T})\bm{U}^{\star T}\|_{F}-\left\|\bm{U}_{X}\bm{U}_{X}^{T }-\bm{U}^{\star}\bm{U}_{A}\bm{U}_{A}^{T}\bm{U}^{\star T}\right\|_{F}$ $\displaystyle\geq\sqrt{2(d-N)-2\max\left\{0,2d-(n+N)\right\}}-\frac{4c_{1}\| \bm{E}\|_{F}}{\sqrt{d}-\sqrt{N-1}}$ $\displaystyle\geq\sqrt{2\min\{d-N,n-d\}}-\frac{4c_{1}\|\bm{E}\|_{F}}{\sqrt{d}- \sqrt{N-1}},$ $$

where the second inequality follows from $\bar{\bm{U}}_{X}=\bm{V}^{\star}$ and (34). Then, we complete the proof.

∎

###### Proof of Theorem 2 .

## Appendix C Proofs in Section 3.2

### C.1 Theoretical Justification of the DAE ( 16 )

Since $\bm{x}_{t}=s_{t}\bm{x}_{0}+\gamma_{t}\bm{\epsilon}$, we compute

$$ $\displaystyle\mathbb{E}_{\bm{\epsilon}}\left[\|\bm{U}_{k}^{T}(s_{t}\bm{x}_{0}+ \gamma_{t}\bm{\epsilon})\|^{2}\right]=s_{t}^{2}\|\bm{U}_{k}^{T}\bm{x}_{0}\|^{2 }+\gamma_{t}^{2}\mathbb{E}_{\bm{\epsilon}}[\|\bm{U}_{k}^{T}\bm{\epsilon}\|^{2} ]=s_{t}^{2}\|\bm{U}_{k}^{T}\bm{x}_{0}\|^{2}+\gamma_{t}^{2}d,$ $$

where the first equality is due to $\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I}_{n})$ and $\mathbb{E}_{\bm{\epsilon}}[\langle\bm{U}_{k}^{T}\bm{x}_{0},\bm{U}_{k}^{T}\bm{
\epsilon}\rangle]=\bm{0}$ for each $k\in[K]$. This implies that when $n$ is sufficiently large, we can approximate $w_{k}(\bm{\theta};\bm{x}_{t})$ in (9) well by

$$ $\displaystyle w_{k}(\bm{\theta};\bm{x}_{t})$ $\displaystyle\approx\frac{\exp\left(\phi_{t}\left(s_{t}^{2}\|\bm{U}_{k}^{T}\bm {x}_{0}\|^{2}+\gamma_{t}^{2}d\right)\right)}{\sum_{l=1}^{K}\exp\left(\phi_{t} \left(s_{t}^{2}\|\bm{U}_{l}^{T}\bm{x}_{0}\|^{2}+\gamma_{t}^{2}d\right)\right)}.$ $$

This soft-max function can be further approximated by the hard-max function. Therefore, we directly obtain (17).

### C.2 Proof of Theorem 3

Equipped with the above setup, we are ready to prove Theorem 3.

###### Proof of Theorem 3 .

Plugging (16) into the integrand of (6) yields

$$ $\displaystyle\ \mathbb{E}_{\bm{\epsilon}}\left[\left\|\frac{s_{t}}{s_{t}^{2}+ \gamma_{t}^{2}}\sum_{k=1}^{K}\hat{w}_{k}(\bm{\theta};\bm{x}^{(i)})\bm{U}_{k} \bm{U}_{k}^{T}(s_{t}\bm{x}^{(i)}+\gamma_{t}\bm{\epsilon})-\bm{x}^{(i)}\right\| ^{2}\right]$ $\displaystyle=$ $\displaystyle\ \left\|\frac{s_{t}^{2}}{s_{t}^{2}+\gamma_{t}^{2}}\sum_{k=1}^{K} \hat{w}_{k}(\bm{\theta};\bm{x}^{(i)})\bm{U}_{k}\bm{U}_{k}^{T}\bm{x}^{(i)}-\bm{ x}^{(i)}\right\|^{2}+\frac{(s_{t}\gamma_{t})^{2}}{(s_{t}^{2}+\gamma_{t}^{2})^{ 2}}\mathbb{E}_{\bm{\epsilon}}\left[\left\|\sum_{k=1}^{K}\hat{w}_{k}(\bm{\theta };\bm{x}^{(i)})\bm{U}_{k}\bm{U}_{k}^{T}\bm{\epsilon}\right\|^{2}\right]$ $\displaystyle=$ $\displaystyle\ \frac{s_{t}^{2}}{s_{t}^{2}+\gamma_{t}^{2}}\sum_{k=1}^{K}\left( \frac{s_{t}^{2}}{s_{t}^{2}+\gamma_{t}^{2}}\hat{w}_{k}^{2}(\bm{\theta};\bm{x}^{ (i)})-2\hat{w}_{k}(\bm{\theta};\bm{x}^{(i)})\right)\|\bm{U}_{k}^{T}\bm{x}^{(i) }\|^{2}+\|\bm{x}^{(i)}\|^{2}+\frac{(s_{t}\gamma_{t})^{2}d}{(s_{t}^{2}+\gamma_{ t}^{2})^{2}}\sum_{k=1}^{K}\hat{w}_{k}(\bm{\theta};\bm{x}^{(i)}),$ $$

where the first equality follows from $\mathbb{E}_{\bm{\epsilon}}[\langle\bm{x},\bm{\epsilon}\rangle]=0$ for any fixed $\bm{x}\in\mathbb{R}^{n}$ due to $\bm{\epsilon}\sim\mathcal{N}(\bm{0},\bm{I}_{n})$, and the last equality uses $\bm{U}_{k}\in\mathcal{O}^{n\times d}$ and $\bm{U}_{k}^{T}\bm{U}_{l}=\bm{0}$ for all $k\neq l$. This, together with (6) and $\gamma_{t}=s_{t}\sigma_{t}$, yields

$$ $\displaystyle\ell(\bm{\theta})$ $\displaystyle=\frac{1}{N}\sum_{i=1}^{N}\sum_{k=1}^{K}\int_{0}^{1}\frac{\lambda _{t}}{1+\sigma_{t}^{2}}\left(\frac{1}{1+\sigma_{t}^{2}}\hat{w}_{k}^{2}(\bm{ \theta};\bm{x}^{(i)})-2\hat{w}_{k}(\bm{\theta};\bm{x}^{(i)})\right){\rm d}t\| \bm{U}_{k}^{T}\bm{x}^{(i)}\|^{2}+$ $\displaystyle\frac{1}{N}\int_{0}^{1}\lambda_{t}{\rm d}t\sum_{i=1}^{N}\|\bm{x}^ {(i)}\|^{2}+\left(\int_{0}^{1}\frac{\sigma_{t}^{2}\lambda_{t}}{(1+\sigma_{t}^{ 2})^{2}}{\rm d}t\right)\frac{d}{N}\sum_{i=1}^{N}\sum_{k=1}^{K}\hat{w}_{k}^{2}( \bm{\theta};\bm{x}^{(i)}).$ $$

According to (16), we can partition $[N]$ into $\{C_{k}(\bm{\theta})\}_{k=1}^{K}$, where $C_{k}(\bm{\theta})$ for each $k\in[K]$ is defined as follows:

$$ $\displaystyle C_{k}(\bm{\theta}):=\left\{i\in[N]:\|\bm{U}_{k}^{T}\bm{x}^{(i)} \|\geq\|\bm{U}_{l}^{T}\bm{x}^{(i)}\|,\ \forall l\neq k\right\},\forall k\in[K].$ (35) $$

Then, we obtain

$$ $\displaystyle\sum_{i=1}^{N}\sum_{k=1}^{K}\hat{w}_{k}^{2}(\bm{\theta};\bm{x}^{( i)})=\sum_{k=1}^{K}\sum_{i\in C_{k}(\bm{\theta})}1=N.$ $$

This, together with plugging (35) into the above loss function, yields minimizing $\ell(\bm{\theta})$ is equivalent to minimizing

$$ $\displaystyle\frac{1}{N}\sum_{i=1}^{N}\sum_{k=1}^{K}\int_{0}^{1}\frac{\lambda_ {t}}{1+\sigma_{t}^{2}}\left(\frac{1}{1+\sigma_{t}^{2}}\hat{w}_{k}^{2}(\bm{ \theta};\bm{x}^{(i)})-2\hat{w}_{k}(\bm{\theta};\bm{x}^{(i)})\right){\rm d}t\| \bm{U}_{k}^{T}\bm{x}^{(i)}\|^{2}$ $\displaystyle=$ $\displaystyle\left(\int_{0}^{1}\frac{\lambda_{t}}{1+\sigma_{t}^{2}}\left(\frac {1}{1+\sigma_{t}^{2}}-2\right){\rm d}t\right)\frac{1}{N}\sum_{k=1}^{K}\sum_{i \in C_{k}(\bm{\theta})}\|\bm{U}_{k}^{T}\bm{x}^{(i)}\|^{2}.$ $$

Since $\frac{\lambda_{t}}{1+\sigma_{t}^{2}}\left(\frac{1}{1+\sigma_{t}^{2}}-2\right)<0$ for all $t\in[0,1]$, minimizing the above function is equivalent to

$$ $\displaystyle\max_{\bm{\theta}}\frac{1}{N}\sum_{k=1}^{K}\sum_{i\in C_{k}(\bm{ \theta})}\|\bm{U}_{k}^{T}\bm{x}^{(i)}\|^{2}\qquad\mathrm{s.t.}\ \left[\bm{U}_{ 1}\ \dots\ \bm{U}_{K}\right]\in\mathcal{O}^{n\times dK}.$ $$

Then, we complete the proof.
∎

###### Proof of Theorem 3 .

### C.3 Proof of Theorem 4

###### Proof of Theorem 4 .

For ease of exposition, let $\delta:=\max\{\|\bm{e}_{i}\|:i\in[N]\}$,

$$ $\displaystyle f(\bm{\theta}):=\sum_{k=1}^{K}\sum_{i\in C_{k}(\bm{\theta})}\| \bm{U}_{k}^{T}\bm{x}^{(i)}\|^{2},$ $$

and for each $k\in[K]$,

$$ $\displaystyle C_{k}^{\star}:=\left\{i\in[N]:\bm{x}^{(i)}=\bm{U}_{k}^{\star}\bm {a}_{i}+\bm{e}_{i}\right\}.$ $$

Suppose that (50) and (51) hold with $\bm{V}=\hat{\bm{U}}_{k}$ for all $i\in[N]$ and $k\neq l\in[K]$, which happens with probability $1-2K^{2}N^{-1}$ according to Lemma 5. This implies that for all $i\in[N]$ and $k\neq l\in[K]$,

$$ $\displaystyle\sqrt{d}-(2\sqrt{\log N}+2)\leq\|\bm{a}_{i}\|\leq\sqrt{d}+(2\sqrt {\log N}+2),$ (36) $\displaystyle\|\hat{\bm{U}}_{k}^{T}\bm{U}_{l}^{\star}\|_{F}-(2\sqrt{\log N}+2) \leq\|\hat{\bm{U}}_{k}^{T}\bm{U}_{l}^{\star}\bm{a}_{i}\|\leq\|\hat{\bm{U}}_{k} ^{T}\bm{U}_{l}^{\star}\|_{F}+(2\sqrt{\log N}+2).$ (37) $$

Recall that the underlying basis matrices are denoted by $\bm{\theta}^{\star}=\{\bm{U}_{k}^{\star}\}_{k=1}^{K}$ and the optimal basis matrices are denoted by $\hat{\bm{\theta}}=\{\hat{\bm{U}}_{k}\}_{k=1}^{K}$.

First, we claim that $C_{k}(\bm{\theta}^{\star})=C_{k}^{\star}$ for each $k\in[K]$. Indeed, for each $i\in C_{k}^{\star}$, we compute

$$ $\displaystyle\|\bm{U}_{k}^{\star T}\bm{x}^{(i)}\|=\|\bm{U}_{k}^{\star T}(\bm{U }_{k}^{\star}\bm{a}_{i}+\bm{e}_{i})\|=\|\bm{a}_{i}+\bm{U}_{k}^{\star T}\bm{e}_ {i}\|\geq\|\bm{a}_{i}\|-\|\bm{e}_{i}\|,$ (38) $\displaystyle\|\bm{U}_{l}^{\star T}\bm{x}^{(i)}\|=\|\bm{U}_{l}^{\star^{T}}(\bm {U}_{k}^{\star}\bm{a}_{i}+\bm{e}_{i})\|=\|\bm{U}_{l}^{\star T}\bm{e}_{i}\|\leq \|\bm{e}_{i}\|,\ \forall l\neq k.$ (39) $$

This, together with (36) and $\|\bm{e}_{i}\|<(\sqrt{d}-2\sqrt{\log N})/2$, implies $\|\bm{U}_{k}^{\star T}\bm{x}_{i}\|\geq\|\bm{U}_{l}^{\star T}\bm{x}_{i}\|$ for all $l\neq k$. Therefore, we have $i\in C_{k}(\bm{\theta}^{\star})$ due to (35). Therefore, we have $C_{k}^{\star}\subseteq C_{k}(\bm{\theta}^{\star})$ for each $k\in[K]$. This, together with the fact that they respectively denote a partition of $[N]$, yields $C_{k}(\bm{\theta}^{\star})=C_{k}^{\star}$ for each $k\in[K]$. Now, we compute

$$ $\displaystyle f(\bm{\theta}^{\star})$ $\displaystyle=\sum_{k=1}^{K}\sum_{i\in C_{k}^{\star}}\|\bm{U}_{k}^{\star T}\bm {x}^{(i)}\|^{2}=\sum_{k=1}^{K}\sum_{i\in C_{k}^{\star}}\|\bm{a}_{i}+\bm{U}_{k} ^{\star T}\bm{e}_{i}\|^{2}$ $\displaystyle=\sum_{i=1}^{N}\|\bm{a}_{i}\|^{2}+2\sum_{k=1}^{K}\sum_{i\in C_{k} ^{\star}}\langle\bm{a}_{i},\bm{U}_{k}^{\star T}\bm{e}_{i}\rangle+\sum_{k=1}^{K }\sum_{i\in C_{k}^{\star}}\|\bm{U}_{k}^{\star T}\bm{e}_{i}\|^{2}.$ (40) $$

Next, we compute

$$ $\displaystyle f(\hat{\bm{\theta}})$ $\displaystyle=\sum_{k=1}^{K}\sum_{i\in C_{k}(\hat{\bm{\theta}})}\|\hat{\bm{U}} _{k}^{T}\bm{x}^{(i)}\|^{2}=\sum_{l=1}^{K}\sum_{k=1}^{K}\sum_{i\in C_{k}(\hat{ \bm{\theta}})\cap C_{l}^{\star}}\|\hat{\bm{U}}_{k}^{T}(\bm{U}_{l}^{\star}\bm{a }_{i}+\bm{e}_{i}))\|^{2}$ $\displaystyle=\sum_{l=1}^{K}\sum_{k=1}^{K}\sum_{i\in C_{k}(\hat{\bm{\theta}}) \cap C_{l}^{\star}}\left(\|\hat{\bm{U}}_{k}^{T}\bm{U}_{l}^{\star}\bm{a}_{i}\|^ {2}+2\langle\bm{a}_{i},\bm{U}_{l}^{\star T}\hat{\bm{U}}_{k}\hat{\bm{U}}_{k}^{T }\bm{e}_{i}\rangle\right)+\sum_{k=1}^{K}\sum_{i\in C_{k}(\hat{\bm{\theta}})}\| \hat{\bm{U}}_{k}^{T}\bm{e}_{i}\|^{2}.$ $$

This, together with $f(\hat{\bm{\theta}})\geq f(\bm{\theta}^{\star})$ and (C.3), yields

$$ $\displaystyle\sum_{i=1}^{N}\|\bm{a}_{i}\|^{2}-$ $\displaystyle\sum_{l=1}^{K}\sum_{k=1}^{K}\sum_{i\in C_{k}(\hat{\bm{\theta}}) \cap C_{l}^{\star}}\|\hat{\bm{U}}_{k}^{T}\bm{U}_{l}^{\star}\bm{a}_{i}\|^{2} \leq\sum_{l=1}^{K}\sum_{k=1}^{K}\sum_{i\in C_{k}(\hat{\bm{\theta}})\cap C_{l}^ {\star}}2\langle\bm{a}_{i},\bm{U}_{l}^{\star T}\hat{\bm{U}}_{k}\hat{\bm{U}}_{k }^{T}\bm{e}_{i}\rangle+$ $\displaystyle\sum_{k=1}^{K}\sum_{i\in C_{k}(\hat{\bm{\theta}})}\|\hat{\bm{U}}_ {k}^{T}\bm{e}_{i}\|^{2}-2\sum_{k=1}^{K}\sum_{i\in C_{k}^{\star}}\langle\bm{a}_ {i},\bm{U}_{k}^{\star T}\bm{e}_{i}\rangle-\sum_{k=1}^{K}\sum_{i\in C_{k}^{ \star}}\|\bm{U}_{k}^{\star T}\bm{e}_{i}\|^{2}$ $\displaystyle\leq 4\delta\sum_{i=1}^{N}\|\bm{a}_{i}\|+N\delta^{2}\leq 6\delta N \sqrt{d}+N\delta^{2},$ (41) $$

where the second inequality follows from $\|\bm{e}_{i}\|\leq\delta$ for all $i\in[N]$ and $\bm{U}_{k}^{\star},\hat{\bm{U}}_{k}\in\mathcal{O}^{n\times d}$ for all $k\in[K]$, and the last inequality uses (36).

For ease of exposition, let $N_{kl}:=|C_{k}(\hat{\bm{\theta}})\cap C_{l}^{\star}|$. According to the pigeonhole principle, there exists a permutation $\pi:[K]\to[K]$ such that there exists $k\in[K]$ such that $N_{\pi(k)k}\geq N/K^{2}$. This, together with (C.3), yields

$$ $\displaystyle 6\delta N\sqrt{d}+N\delta^{2}$ $\displaystyle\geq\sum_{i\in C_{\pi(k)}(\hat{\bm{\theta}})\cap C_{k}^{\star}} \left(\|\bm{a}_{i}\|^{2}-\|\hat{\bm{U}}_{\pi(k)}^{T}\bm{U}_{k}^{\star}\bm{a}_{ i}\|^{2}\right)$ $\displaystyle=\langle\bm{I}-\bm{U}_{k}^{\star^{T}}\hat{\bm{U}}_{\pi(k)}\hat{ \bm{U}}_{\pi(k)}^{T}\bm{U}_{k}^{\star},\sum_{i\in C_{\pi(k)}(\hat{\bm{\theta}} )\cap C_{k}^{\star}}\bm{a}_{i}\bm{a}_{i}^{T}\rangle.$ (42) $$

According to Lemma 6 and $N_{\pi(k)k}\geq N/K^{2}$, it holds with probability at least $1-2K^{4}N^{-2}$ that

$$ $\displaystyle\left\|\frac{1}{N_{\pi(k)k}}\sum_{i\in C_{\pi(k)}(\hat{\bm{\theta }})\cap C_{k}^{\star}}\bm{a}_{i}\bm{a}_{i}^{T}-\bm{I}\right\|\leq\frac{9(\sqrt {d}+\sqrt{\log(N_{\pi(k)k}})}{\sqrt{N_{\pi(k)k}}}.$ $$

This, together with the Weyl’s inequality, yields

$$ $\displaystyle\lambda_{\min}\left(\sum_{i\in C_{\pi(k)}(\hat{\bm{\theta}})\cap C _{k}^{\star}}\bm{a}_{i}\bm{a}_{i}^{T}\right)$ $\displaystyle\geq N_{\pi(k)k}-9\sqrt{N_{\pi(k)k}}\left(\sqrt{d}+\sqrt{\log(N_{ \pi(k)k})}\right)$ $\displaystyle\geq\frac{N}{K^{2}}-\frac{9\sqrt{N}}{K}\left(\sqrt{d}+\sqrt{\log N }\right)\geq\frac{N}{2K^{2}},$ $$

where the second inequality follows from $N/K^{2}\leq N_{\pi(k)k}\leq N$ and the last inequality is due to $\sqrt{N}\geq 18K(\sqrt{d}+\sqrt{\log N})$. Using this and Lemma 7, we obtain

$$ $\displaystyle\quad\langle\bm{I}-\bm{U}_{k}^{\star T}\hat{\bm{U}}_{\pi(k)}\hat{ \bm{U}}_{\pi(k)}^{T}\bm{U}_{k}^{\star},\sum_{i\in C_{\pi(k)}(\hat{\bm{\theta}} )\cap C_{k}^{\star}}\bm{a}_{i}\bm{a}_{i}^{T}\rangle$ $\displaystyle\geq\lambda_{\min}\left(\sum_{i\in C_{\pi(k)}(\hat{\bm{\theta}}) \cap C_{k}^{\star}}\bm{a}_{i}\bm{a}_{i}^{T}\right)\mathrm{Tr}\left(\bm{I}-\bm{ U}_{k}^{\star T}\hat{\bm{U}}_{\pi(k)}\hat{\bm{U}}_{\pi(k)}^{T}\bm{U}_{k}^{ \star}\right)$ $\displaystyle\geq\frac{N}{2K^{2}}\mathrm{Tr}\left(\bm{I}-\bm{U}_{k}^{\star T} \hat{\bm{U}}_{\pi(k)}\hat{\bm{U}}_{\pi(k)}^{T}\bm{U}_{k}^{\star}\right).$ $$

This, together with (C.3), implies

$$ $\displaystyle\mathrm{Tr}\left(\bm{I}-\bm{U}_{k}^{\star T}\hat{\bm{U}}_{\pi(k)} \hat{\bm{U}}_{\pi(k)}^{T}\bm{U}_{k}^{\star}\right)\leq 2K^{2}\left(6\delta \sqrt{d}+\delta^{2}\right).$ $$

Using this and $[\bm{U}_{1}^{\star},\dots,\bm{U}_{k}^{\star}]\in\mathcal{O}^{n\times dK}$, we obtain

$$ $\displaystyle\sum_{l\neq k}\|\hat{\bm{U}}_{\pi(k)}^{T}\bm{U}_{l}^{\star}\|_{F} ^{2}$ $\displaystyle=\mathrm{Tr}\left(\sum_{l\neq k}\hat{\bm{U}}_{\pi(k)}^{T}\bm{U}_{ l}^{\star}\bm{U}_{l}^{\star T}\hat{\bm{U}}_{\pi(k)}\right)\leq\mathrm{Tr}\left (\bm{I}-\hat{\bm{U}}_{\pi(k)}^{T}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\hat{ \bm{U}}_{\pi(k)}\right)$ $\displaystyle\leq 2K^{2}\left(6\delta\sqrt{d}+\delta^{2}\right)\leq\frac{3d}{4},$ (43) $$

where the last inequality follows $\delta\leq\sqrt{d}/(24K^{2})$. According to (C.3), we have

$$ $\displaystyle 6\delta N\sqrt{d}+N\delta^{2}$ $\displaystyle\geq\sum_{l\neq k}^{K}\sum_{i\in C_{\pi(k)}(\hat{\bm{\theta}}) \cap C_{l}^{\star}}\left(\|\bm{a}_{i}\|^{2}-\|\hat{\bm{U}}_{\pi(k)}^{T}\bm{U}_ {l}^{\star}\bm{a}_{i}\|^{2}\right)$ $\displaystyle\geq\sum_{l\neq k}^{K}N_{\pi(k)l}\left((\sqrt{d}-\alpha)^{2}- \left(\|\hat{\bm{U}}_{\pi(k)}^{T}\bm{U}_{l}^{\star}\|_{F}+\alpha\right)^{2} \right)\geq\frac{d}{8}\sum_{l\neq k}^{K}N_{\pi(k)l},$ $$

where the second inequality uses (36) and (37), and the last inequality follows from $d\gtrsim\log N$. Therefore, we have for each $k\in[K]$,

$$ $\displaystyle\sum_{l\neq k}^{K}N_{\pi(k)l}\leq\frac{48\delta N\sqrt{d}+8\delta ^{2}N}{d}<1,$ $$

where the last inequality uses $\delta\lesssim\sqrt{d/N}$. This implies $N_{\pi(l)k}=0$ for all $l\neq k$, and thus $C_{\pi(k)}(\hat{\bm{\theta}})\subseteq C_{k}^{\star}$. Using the same argument, we can show that $C_{\pi(l)}(\hat{\bm{\theta}})\subseteq C_{l}^{\star}$ for each $l\neq k$. Therefore, we have $C_{\pi(k)}(\hat{\bm{\theta}})=C_{k}^{\star}$ for each $k\in[K]$. In particular, using the union bound yields event holds with probability at least $1-2K^{2}N^{-1}$. Based on the above optimal assignment, we can further show:

(i) Suppose that $N_{k}\geq d$ for each $k\in[K]$. This, together with (i) in Theorem 2 and $N_{k}\geq d$, yields (19).

(ii) Suppose that there exists $k\in[K]$ such that $N_{k}<d$. This, together with (ii) in Theorem 2 and $N_{k}\geq d$, yields (20).

Finally, applying the union bound yields the probability of these events.
∎

###### Proof of Theorem 4 .

## Appendix D Experimental Setups in Section 2.2

In this section, we provide detailed setups for the experiments in Section 2.2. These experiments aim to validate the assumptions that real-world image data satisfies a mixture of low-rank Gaussians and that the DAE is parameterized according to (9). To begin, we show that $\nabla_{\bm{x}_{t}}\mathbb{E}[\bm{x}_{0}|\bm{x}_{t}]$ is of low rank when $p_{\rm data}$ follows a mixture of low-rank Gaussians and $\sum_{k=1}^{K}d_{k}\leq n$, where $n$ is the ambient dimension of training samples.

### D.1 Verification of Mixture of Low-Rank Gaussian Data Distribution

In this subsection, we demonstrate that a mixture of low-rank Gaussians is a reasonable and insightful model for approximating real-world image data distribution. To begin, we show that $\nabla_{\bm{x}_{t}}\mathbb{E}[\bm{x}_{0}|\bm{x}_{t}]$ is of low rank when $p_{\rm data}$ follows a mixture of low-rank Gaussians with $\sum_{k=1}^{K}d_{k}\leq n$, where $n$ is the dimension of training samples.

###### Lemma 2 .

Suppose that the data distribution $p_{\rm data}$ follows a mixture of low-rank Gaussian distributions as defined in Definition 1. For all $t\in[0,1]$, it holds that

$$ $\displaystyle\min_{k\in[K]}d_{k}\leq\mathrm{rank}\left(\nabla_{\bm{x}_{t}} \mathbb{E}[\bm{x}_{0}|\bm{x}_{t}]\right)\leq\sum_{k=1}^{K}d_{k}.$ (44) $$

###### Proof.

For ease of exposition, let

$$ $\displaystyle h_{k}(\bm{x}_{t}):=\exp\left(\phi_{t}\|\bm{U}_{k}^{\star T}\bm{x }_{t}\|^{2}\right),\ \forall k\in[K].$ $$

Obviously, we have

$$ $\displaystyle\nabla h_{k}(\bm{x}_{t}):=2\phi_{t}\exp\left(\phi_{t}\|\bm{U}_{k} ^{\star T}\bm{x}_{t}\|^{2}\right)\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\bm{x}_ {t}=2\phi_{t}h_{k}(\bm{x}_{t})\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\bm{x}_{t}.$ (45) $$

According to Lemma 1, we have

$$ $\displaystyle\mathbb{E}[\bm{x}_{0}|\bm{x}_{t}]=\frac{s_{t}}{s_{t}^{2}+\gamma_{ t}^{2}}f(\bm{x}_{t}),\ \text{where}\ f(\bm{x}_{t}):=\frac{\sum_{k=1}^{K}\pi_{k }h_{k}(\bm{x}_{t})\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\bm{x}_{t}}{\sum_{k=1} ^{K}\pi_{k}h_{k}(\bm{x}_{t})}.$ $$

Then, we compute

$$ $\displaystyle\nabla_{\bm{x}_{t}}f(\bm{x}_{t})$ $\displaystyle=\frac{1}{\sum_{k=1}^{K}\pi_{k}h_{k}(\bm{x}_{t})}\left(2\phi_{t} \sum_{k=1}^{K}\pi_{k}h_{k}(\bm{x}_{t})\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T} \bm{x}_{t}\bm{x}_{t}^{T}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}+\sum_{k=1}^{K} \pi_{k}h_{k}(\bm{x}_{t})\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\right)$ $\displaystyle\quad-\frac{2\phi_{t}}{\left(\sum_{k=1}^{K}\pi_{k}h_{k}(\bm{x}_{t })\right)^{2}}\left(\sum_{k=1}^{K}\pi_{k}h_{k}(\bm{x}_{t})\bm{U}_{k}^{\star} \bm{U}_{k}^{\star T}\bm{x}_{t}\right)\left(\sum_{k=1}^{K}\pi_{k}h_{k}(\bm{x}_{ t})\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\bm{x}_{t}\right)^{T}$ $\displaystyle=\frac{1}{\sum_{k=1}^{K}\pi_{k}h_{k}(\bm{x}_{t})}\sum_{k=1}^{K} \pi_{k}h_{k}(\bm{x}_{t})\left(2\phi_{t}\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T} \bm{x}_{t}\bm{x}_{t}^{T}+\bm{I}\right)\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}-$ $\displaystyle\quad\frac{2\phi_{t}}{\left(\sum_{k=1}^{K}\pi_{k}h_{k}(\bm{x}_{t} )\right)^{2}}\left(\sum_{k=1}^{K}\pi_{k}h_{k}(\bm{x}_{t})\bm{U}_{k}^{\star}\bm {U}_{k}^{\star T}\right)\bm{x}_{t}\bm{x}_{t}^{T}\left(\sum_{k=1}^{K}\pi_{k}h_{ k}(\bm{x}_{t})\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}\right).$ $$

This directly yields (44) for all $t\in[0,1]$.
∎

Now, we conduct experiments to illustrate that diffusion models trained on real-world image datasets exhibit similar low-rank properties to those described in the above proposition. Provided that the DAE $\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)$ is applied to estimate $\mathbb{E}[\bm{x}_{0}|\bm{x}_{t}]$, we estimate the rank of the Jacobian of the DAE, i.e., $\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)$, on the real-world data distribution, where $\bm{\theta}$ denotes the parameters of U-Net architecture trained on the real dataset. Also, this estimation is based on the findings in [, ] that under the training loss in Equation 6, the DAE $\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)$ converge to $\mathbb{E}[\bm{x}_{0}|\bm{x}_{t}]$ as the number of training samples increases on the real data. We evaluate the numerical rank of the Jacobian of the DAE on four different datasets: CIFAR-10 [], CelebA [], FFHQ [] and AFHQ [], where the ambient dimension $n=3072$ for all datasets.

Given a random initial noise $\bm{x}_{1}\sim\mathcal{N}(\bm{0},\bm{I}_{n})$, diffusion models generate a sequence of images $\{\bm{x}_{t}\}$ according to the reverse SDE in Eq. (4). Along the sampling trajectory $\{\bm{x}_{t}\}$, we calculate the Jacobian $\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)$ and compute its numerical rank via

$$ $\displaystyle\mathrm{rank}\left(\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x} _{t},t)\right):=\operatorname{\arg\min}\left\{r\in[1,n]:\frac{\sum_{i=1}^{r} \sigma_{i}^{2}\left(\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t) \right)}{\sum_{i=1}^{n}\sigma_{i}^{2}\left(\nabla_{\bm{x}_{t}}\bm{x}_{\bm{ \theta}}(\bm{x}_{t},t)\right)}>\eta^{2}\right\}.$ (46) $$

In our experiments, we set $\eta=0.99$. In the implementation, we utilize the Elucidating Diffusion Model (EDM) with the EDM noise scheduler [] and DDPM++ architecture []. Moreover, we employ an 18-step Heun’s solver for sampling and present the results for 12 of these steps. For each dataset, we random sample 15 initial noise $\bm{x}_{1}$, calculate the mean of $\mathrm{rank}(\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t))$ along the trajectory $\{\bm{x}_{t}\}$, and plot ratio of the numerical rank over the ambient dimension against the signal-noise-ratio (SNR) $1/\sigma_{t}$ in Figure 3, where $\sigma_{t}$ is defined in Eq. (3).

Figure: (a) Numerical rank
Refer to caption: extracted/5832094/figure/jacobian_training_iteration.png

###### Lemma 2 .

###### Proof.

### D.2 Verification of Low-Rank Network Parameterization

In this subsection, we empirically investigate the properties of U-Net architectures in diffusion models and validate the simplification of the network architecture to Eq. (9). Based on the results in Section D.1, we use a mixture of low-rank Gaussian distributions for experiments. Here, we set $K=2$, $n=48$, $d_{1}=d_{2}=6$, $\pi_{1}=\pi_{2}=0.5$, and $N=1000$ for the data model Definition 1. Moreover, We use the EDM noise scheduler and 18-step Heun’s solver for both the U-Net and our proposed parameterization (9). To adapt the structure of the U-Net, we reshape each training sample into a 3D tensor with dimensions $4\times 4\times 3$, treating it as an image. Here, we use DDPM++ based diffusion models with a U-Net architecture. In each iteration, we randomly sampled a batch of image $\{\bm{x}^{(j)}\}_{j=1}^{\text{bs}}\subseteq\{\bm{x}^{(i)}\}_{i=1}^{N}$, along with a timestep $t^{(j)}$ and a noise $\bm{\epsilon}^{(j)}$ for each image in the batch to optimize the training loss $\ell(\bm{\theta})$. We define

$$ $\displaystyle\text{kimgs}=\text{bs}\times\frac{\text{training iterations}}{1000}$ (47) $$

to represent the total samples used for training. Here, we pick up the specific model trained under 500 kimgs, 1000 kimgs, 2000 kimgs, and 6000 kimgs for evaluation, as shown in Figure 7(a).

We plot the numerical ranks of $\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)$ for both our proposed parameterization in (9) and for the U-Net architecture in Figure 3(b). According to Lemma 2, it holds that $6\leq\mathrm{rank}(\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t))\leq 12$. This corresponds to the blue curve in Figure 3(b). To supplement our result in Figure 3(b), we further plot the numerical rank against SNR at different training iterations in Figure 7(a) and gradient norm of the objective against training iterations in Figure 7(b). We observe that with the training kimgs increases, the gradient for the U-Net $||\nabla_{\bm{\theta}}\ell||_{F}$ decrease smaller than $10^{-1}$ and the rank ratio of $\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)$ trained from U-Net gradually be close to the rank ratio from the low-rank model in the middle of the SNR $([0.91,10.0])$.

## Appendix E Experimental Setups in Section 4

We use a CPU to optimize Problem (6) for the setting in Section E.1. For the settings in Section E.2 and Section E.3, we employ a single A40 GPU with 48 GB memory to optimize Problem (6).

### E.1 Learning the MoLRG distribution with the theoretical parameterzation

Here, we present the stochastic gradient descent (SGD) algorithm for solving Problem (6) as follows:

Figure: Algorithm 1 SGD for optimizing the training loss (6)

Now, we specify how to choose the parameters of the SGD in our implementation. We divide the time interval $[0,1]$ into $64$ time steps. When $K=1$, we set the learning rate $\eta=10^{-4}$, batch size $M=128N_{k}$, and number of iterations $J=10^{4}$. When $K=2$, we set the learning rate $\eta=2\times 10^{-5}$, batch size $M=1024$, number of iterations $J=10^{5}$. In particular, when $K=2$, we use the following tailor-designed initialization $\bm{\theta}^{0}=\{\bm{U}_{k}^{0}\}$ to improve the convergence of the SGD:

$$ $\displaystyle\bm{U}_{k}^{0}=\bm{U}_{k}^{\star}+0.2\bm{\Delta},\ k\in\{1,2\},$ (48) $$

where $\bm{\Delta}\sim\mathcal{N}(\bm{0},\bm{I}_{n})$. We calculate the success rate as follows. If the returned subspace basis matrices $\{\bm{U}_{k}\}_{k=1}^{K}$ satisfy

$$ $\displaystyle\frac{1}{K}\mathbb{\sum}_{k=1}^{K}||\bm{U}_{\Pi(k)}\bm{U}_{\Pi(k) }^{T}-\bm{U}_{k}^{\star}\bm{U}_{k}^{\star T}||\leq 0.5$ $$

for some permutation $\Pi:[K]\to[K]$, it is considered successful.

### E.2 Learning the MoLRG distribution with U-Net

In our implementation, we set the total dimension of MoLRG as $n=48$ and the number of training samples $N_{\text{eval}}=1000$. To train the U-Net, we use the stochastic gradient descent in Algorithm 1. We use DDPM++ architecture [] for the U-Net and EDM [] noise scheduler. We set the learning rate $10^{-3}$, batch size $64$, and number of iterations $J=10^{4}$.

### E.3 Learning real-world image data distributions with U-Net

To train diffusion models for real-world image datasets, we use the DDPM++ architecture [] for the U-Net and variance preserving (VP) [] noise scheduler. The U-Net is trained using the Adam optimizer [], a variant of SGD in Algorithm 1. We set the learning rate $\eta=10^{-3}$, batch size $M=512$, and the total number of iterations $10^{5}$.

Figure: (a) PCA
Refer to caption: extracted/5832094/figure/PCA_K=3.jpg

### E.4 Correspondence between low-dimensional subspaces and image semantics

We denote the Jacobian of the DAE $\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)$ by $\bm{J}_{t}:=\nabla_{\bm{x}_{t}}\bm{x}_{\bm{\theta}}(\bm{x}_{t},t)\in\mathbb{R}
^{n\times n}$ and let $\bm{J}_{t}=\bm{U}\bm{\Sigma}\bm{V}^{T}$ be an singular value decomposition (SVD) of $\bm{J}_{t}$, where $r=\mathrm{rank}(\bm{J}_{t})$, $\bm{U}=[\bm{u}_{1},\cdots,\bm{u}_{r}]\in\mathcal{O}^{n\times r}$, $\bm{V}=[\bm{v}_{1},\cdots,\bm{v}_{r}]\in\mathcal{O}^{n\times r}$, and $\bm{\Sigma}=\mathrm{diag}(\sigma_{1},\dots,\sigma_{r})$ with $\sigma_{1}\geq\dots\geq\sigma_{r}$ being the singular values. According to the results in Figure 3, it is observed that $\bm{J}_{t}$ is low rank, i.e., $r\ll n$. Now, we compute the first-order approximation of $\bm{x}_{\bm{\theta}}(\bm{x}_{t},\bm{t})$ along the direction of $\bm{v}_{i}\in\mathbb{R}^{n}$, where $\bm{v}_{i}$ is the $i$-th right singular vector of $\bm{J}_{t}$:

$$ $\displaystyle\bm{x}_{\bm{\theta}}(\bm{x}_{t}+\alpha\bm{v}_{i},t)\approx\bm{x}_ {\bm{\theta}}(\bm{x}_{t},t)+\alpha\bm{J}_{t}\bm{v}_{i}=\bm{x}_{\bm{\theta}}( \bm{x}_{t},t)+\alpha\sigma_{i}\bm{u}_{i},$ $$

where the last equality follows from $\bm{J}_{t}\bm{v}_{i}=\bm{U}\bm{\Sigma}\bm{V}^{T}\bm{v}_{i}=\alpha\sigma_{i}\bm
{u}_{i}$. To validate the semantic meaning of the basis $\bm{v}_{i}$, we vary the value of $\alpha$ from negative to positive and visualize the resulting changes in the generated images. Figures 6 and 9(a, c) illustrate some real examples.

In the experiments, we use a pre-trained diffusion denoising probabilistic model (DDPM) [] on the MetFaces dataset []. We randomly select an image $\bm{x}_{0}$ from this dataset and use the reverse process of the diffusion denoising implicit model (DDIM) [] to generate $\bm{x}_{t}$ at $t=0.7T$ (ablation studies for $t=0.1T$ and $0.9T$ are shown in Figure 9(b)), where $T$ denote the total number of time steps. We respectively choose the changed direction as the leading right singular vectors $\bm{v}_{1},\bm{v}_{3},\bm{v}_{4},\bm{v}_{5},\bm{v}_{6}$ and use $\tilde{\bm{x}}_{t}=\bm{x}_{t}+\alpha\bm{v}_{i}$ to generate new images with $\alpha\in[-6,6]$ shown in Figures 6 and 9(a, c).

Figure: (a)
Refer to caption: extracted/5832094/figure/peng_exp2_more_1.jpg

## Appendix F Auxiliary Results

First, we present a probabilistic result to prove Theorem 2, which provides an optimal estimate of the small singular values of a matrix with i.i.d. Gaussian entries. This lemma is proved in [, Theorem 1.1].

###### Lemma 3 .

Let $\bm{A}$ be an $m\times n$ random matrix, where $m\geq n$, whose elements are independent copies of a subgaussian random variable with mean zero and unit variance. It holds for every $\varepsilon>0$ that

$$ $\displaystyle\mathbb{P}\left(\sigma_{\min}(\bm{A})\geq\varepsilon(\sqrt{m}- \sqrt{n-1})\right)\geq 1-\left(c_{1}\varepsilon\right)^{m-n+1}-\exp\left(-c_{2 }m\right),$ $$

where $c_{1},c_{2}>0$ are constants depending polynomially only on the subgaussian moment.

Next, we present a probabilistic bound on the deviation of the norm of weighted sum of squared Gaussian random variables from its mean. This is a direct extension of [, Theorem 5.2.2].

###### Lemma 4 .

Let $\bm{x}\sim\mathcal{N}(\bm{0},\bm{I}_{d})$ be a Gaussian random vector and $\lambda_{1},\dots,\lambda_{d}>0$ be constants. It holds for any $t>0$ that

$$ $\displaystyle\mathbb{P}\left(\left|\sqrt{\sum_{i=1}^{d}\lambda_{i}^{2}x_{i}^{2 }}-\sqrt{\sum_{i=1}^{d}\lambda_{i}^{2}}\right|\geq t+2\lambda_{\max}\right) \leq 2\exp\left(-\frac{t^{2}}{2\lambda_{\max}^{2}}\right),$ (49) $$

where $\lambda_{\max}=\max\{\lambda_{i}:i\in[d]\}$.

Based on the above lemma, we can further show the following concentration inequalities to estimate the norm of the standard norm Gaussian random vector.

###### Lemma 5 .

Suppose that $\bm{a}_{i}\overset{i.i.d.}{\sim}\mathcal{N}(\bm{0},\bm{I}_{d})$ is a Gaussian random vector for each $i\in[N]$. The following statements hold:
(i) It holds for all $i\in[N]$ with probability at least $1-N^{-1}$ that

$$ $\displaystyle\left|\|\bm{a}_{i}\|-\sqrt{d}\right|\leq 2\sqrt{\log N}+2.$ (50) $$

(ii) Let $\bm{V}\in\mathcal{O}^{n\times d}$ be given. For all $i\in C_{k}^{\star}$ and all $k\in[K]$, it holds with probability at least $1-2N^{-1}$ that

$$ $\displaystyle\left|\|\bm{V}^{T}\bm{U}_{k}^{\star}\bm{a}_{i}\|-\|\bm{V}^{T}\bm{ U}_{k}^{\star}\|_{F}\right|\leq 2\sqrt{\log N}+2.$ (51) $$

###### Proof.

(i) Applying Lemma 4 to $\bm{a}_{i}\sim\mathcal{N}(\bm{0},\bm{I}_{d})$, together with setting $t=2\sqrt{\log N}$ and $\lambda_{j}=1$ for all $j\in[d]$, yields

$$ $\displaystyle\mathbb{P}\left(\left|\|\bm{a}_{i}\|-\sqrt{d}\right|\geq 2\sqrt{ \log N}+2\right)\leq 2N^{-2}.$ $$

This, together with the union bound, yields that (50) holds with probability $1-N^{-1}$.

(ii) Let $\bm{V}^{T}\bm{U}^{\star}_{k}=\bm{P}\bm{\Sigma}\bm{Q}^{T}$ be a singular value decomposition of $\bm{V}^{T}\bm{U}^{\star}_{k}$, where $\bm{\Sigma}\in\mathbb{R}^{d\times d}$ with the diagonal elements $0\leq\sigma_{d}\leq\dots\sigma_{1}\leq 1$ being the singular values of $\bm{V}^{T}\bm{U}^{\star}_{k}$ and $\bm{P},\bm{Q}\in\mathcal{O}^{d}$. This, together with the orthogonal invariance of the Gaussian distribution, yields

$$ $\displaystyle\|\bm{V}^{T}\bm{U}^{\star}_{k}\bm{a}_{i}\|=\|\bm{\Sigma}\bm{Q}^{T }\bm{a}_{i}\|\overset{d}{=}\|\bm{\Sigma}\bm{a}_{i}\|=\sqrt{\sum_{j=1}^{d} \sigma_{j}^{2}a_{ij}^{2}}.$ (52) $$

Using Lemma 4 with setting $t=2\sigma_{1}\sqrt{\log N}$ and $\lambda_{j}=\sigma_{j}\leq 1$ for all $j$ yields

$$ $\displaystyle\mathbb{P}\left(\left|\|\bm{V}^{T}\bm{U}^{\star}_{k}\bm{a}_{i}\|- \|\bm{V}^{T}\bm{U}^{\star}_{k}\|_{F}\right|\geq\sigma_{1}\alpha\right)$ $\displaystyle=\mathbb{P}\left(\left|\sqrt{\sum_{j=1}^{d}\sigma_{j}^{2}a_{ij}^{ 2}}-\sqrt{\sum_{j=1}^{d}\sigma_{j}^{2}}\right|\geq\sigma_{1}\alpha\right)\leq 2 N^{-2}.$ $$

This, together with $\sigma_{1}\leq 1$ and the union bound, yields (51).
∎

Next, We present a spectral bound on the covariance estimation for the random vectors generated by the normal distribution.

###### Lemma 6 .

Suppose that $\bm{a}_{1},\dots,\bm{a}_{N}\in\mathbb{R}^{d}$ are i.i.d. standard normal random vectors, i.e., $\bm{a}_{i}\overset{i.i.d.}{\sim}\mathcal{N}(\bm{0},\bm{I}_{d})$ for all $i\in[N]$. Then, it holds with probability at least $1-2N^{-2}$ that

$$ $\displaystyle\left\|\frac{1}{N}\sum_{i=1}^{N}\bm{a}_{i}\bm{a}_{i}^{T}-\bm{I}_{ d}\right\|\leq\frac{9(\sqrt{d}+\sqrt{\log N})}{\sqrt{N}},$ (53) $$

###### Proof.

According to [, Theorem 4.7.1], it holds that

$$ $\displaystyle\mathbb{P}\left(\left\|\frac{1}{N}\sum_{i=1}^{N}\bm{a}_{i}\bm{a}_ {i}^{T}-\bm{I}_{d}\right\|\geq\frac{9(\sqrt{d}+\eta)}{\sqrt{N}}\right)\leq 2 \exp\left(-2\eta^{2}\right),$ $$

where $\eta>0$.
Plugging $\eta=\sqrt{\log N}$ into the above inequality yields

$$ $\displaystyle\mathbb{P}\left(\left\|\frac{1}{N}\sum_{i=1}^{N}\bm{a}_{i}\bm{a}_ {i}^{T}-\bm{I}_{d}\right\|\geq\frac{9(\sqrt{d}+\sqrt{\log N})}{\sqrt{N}}\right )\leq 2N^{-2}.$ $$

This directly implies (53).
∎

###### Lemma 7 .

Let $\bm{A},\bm{B}\in\mathbb{R}^{n\times n}$ be positive semi-definite matrices. Then, it holds that

$$ $\displaystyle\langle\bm{A},\bm{B}\rangle\geq\lambda_{\min}(\bm{A})\mathrm{Tr}( \bm{B}).$ (54) $$

###### Proof.

Let $\bm{U}\bm{\Lambda}\bm{U}^{T}=\bm{A}$ be an eigenvalue decompositon of $\bm{A}$, where $\bm{U}\in\mathcal{O}^{n}$ and $\bm{\Sigma}=\mathrm{diag}(\lambda_{1},\dots,\lambda_{n})$ is a diagonal matrix with diagonal entries $\lambda_{1}\geq\dots\geq\lambda_{n}\geq 0$ being the eigenvalues. Then, we compute

$$ $\displaystyle\langle\bm{A},\bm{B}\rangle=\langle\bm{U}\bm{\Lambda}\bm{U}^{T}, \bm{B}\rangle=\langle\bm{\Lambda},\bm{U}\bm{B}\bm{U}^{T}\rangle\geq\lambda_{ \min}(\bm{A})\mathrm{Tr}(\bm{U}\bm{B}\bm{U}^{T})=\lambda_{\min}(\bm{A})\mathrm {Tr}(\bm{B}),$ $$

where the inequality follows from $\lambda_{i}\geq 0$ for all $i\in[N]$ and $\bm{B}$ is a positive semidefinite matrix.
∎

###### Lemma 3 .

###### Lemma 4 .

###### Lemma 5 .

###### Proof.

###### Lemma 6 .

###### Proof.

###### Lemma 7 .

###### Proof.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: My questions are listed in the weakness section above. I would encourage the authors to analyze more carefully the transferability of claims from MoLRG at the origin toward MoLRG at every basepoint.
soundness: 3
strengths: The paper argues that instead of generating a full ranked diffusion, the diffusion model only need to solve a subproblem corresponding to the intrinsic dimension of the data distribution. This is a very convincing idea. In addition, in the case of MoLRG (mixture of low rank Gaussians). the authors mathematically proves the problem is equivalent to a PCA optimization and established rigorous guarantees of the optimization quality.
summary: The paper suggests that diffusion models can circumvent curse of dimensionality by clustering to fit the intrinsic dimension which is in general much lower than ambient dimension. This is a valuable insight and the authors provided a solid theoretical analysis to support it in a special case. However, the paper suffers from two draw backs: (1) the study of diffusions along subspaces and their study using PCA techniques is not a completely novel ideas; (2) the study focuses narrowly on finite sums of low-rank Gaussians centered at the origin (up to a noise), effectively making the search space finite dimensional.
weaknesses: Weaknesses:
1. It is not a new idea to study behaviors of diffusion models, or generative models in general, along subspaces spanned by leading eigenvectors of the Jacobian. I cannot produce a complete bibliography in this area, but the following papers may be relevant:
- Subspace Diffusion Generative Models, Jing et al., ECCV 2022 , which developed a new diffusion model that restricts the flow vector field to a linear subspace whose dimension shrinks as t->0 
- The Geometry of Deep Generative Image Models and its Applications, Wang & Ponce, ICLR 2021, whose experiments focused on GANs rather than diffusion models, but revealed that generative models identifies top eigenspaces of the Jacobian, which capture important perceptually relevant changes. 
2. The study in this paper assumes the distribution is a sum of finitely many Gaussians centered at 0 with added noises (and even assumed later that they have the same weights). This is an oversimplification as under the current assumption, the task becomes an optimization of finitely matrices. The real-life diffusion is much more complicated as it tries to identify,  instead of linear subspaces,  submanifolds or equilvalently a tangent linear subspace at each point, which is an infinite dimensional task. How would the results change for data lying on low dimensional non-linear manifold? I guess the complexity of the task would depend on how fast the tangent space change among nearby points, which is quantified as curvatures of the manifold. It would be interesting to see discussions addressing this aspect.
3. This contrast is clearly demonstrated by the current paper's own experiments: under the simplified assumption, only a handful of samples (equal to the rank of the matrix) are needed for decent inference quality while for realistic data thousands of samples are needed ( Figure 5a vs Figure 5b ). While it is true the ranking of difficulty among different datasets coincide with that of intrinsic dimension, it should be recognized that a great amount of new details exist in data with higher instrinsic dimensions, and model quality a priori cannot be naively summarized by transferring the analysis on low-rank Gaussians in Theorems 1-4. I think it would be interesting to have more analysis on whether local intrinsic dimension are the only features that need to be handled in real work datasets, in particular, whether the shape and smoothness of the distribution also play a role.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. Figure 2 shows changes in the image as a function of moving in the direction of top singular vectors of the Jacobian of the denoiser. The semantic labels are strange, because there are multiple features changing in each column but only one is chosen as a description. The complex variations of course is a reflection of the fact that the manifolds are not linear and it is not trivial to separate the features. Overall it is not clear that this figure is trying to convey. The actual effect is not consistent with the description. 
2. In figure 3, for real image datasets, there is a jump at the very low SNR for estimated rank from 3 of the dataset. What is causing this strange behavior? Similarly for the toy data, the behavior of UNet seems pretty strange and non monotonic. Is there an intuition or explanation for that? 
3. In line 364, it's stated that the assumption $U_k^TU_l = 0$ follows from the observation of disjoint union of manifold. a) This is another too strong assumption about images. We know that images share features across different classes and images, so it is not natural to assume this orthogonality between the subspaces. b) Even if we assume this orthogonality, why does this follow from the observation that the manifolds are disjoint? That refers to the support to be non-overlapping. In the general case, when the manifolds are not mean zero, they can have many directions in parallel and still be disjoint. c) Finally, if this is an assumption you are making, it is in consistent with the mixture of Gaussian models you show in figure 1, where your orthogonality assumption does not hold. 
4. In figure 4,  it seems like diffusion models are performing worse than the simple PCA models in terms of separating the success and failure cases. Is it clear why? 
5. I suggest eq 48 and 49 be moved to the main text because they are used to generate figures in the main text.
soundness: 3
strengths: 1. The paper raises an important question regarding generalization requirements in diffusion models. So the paper is significant in the sense that it is about a crucial question. 
2. Extending previous theoretical work from mixture of full rank Gaussians to low rank Gaussians is an important and valuable step, since a mixture of low rank Gaussians is a better model of natural images (and generally real world structured data). 
3. Additionally, showing empirical results in both toy and real data is a plus. 
4. Finally, the paper is written clearly and is fairly easy to follow.
summary: This work is about understanding the generalization capabilities of diffusion models from a theoretical perspective. The main assumption is that the true underlying data distribution can be approximated by a mixture of low rank Gaussians. Following this assumption, the authors propose an ideal parametrization for denoising networks of diffusion models. Combined with a few more approximations (e.g. hard max counterpart of weight assignments), they propose: 1) denoising is reduced to a sub-space clustering problem (in theorem 3) , and 2) the error in approximating the true sub-spaces is related to sub-space dimensionality and number of data points in training samples (theorem 4).
weaknesses: 1. The main weakness of the paper is that the main limitation of the analysis is not acknowledged clearly enough. This limitation is due to the main assumption of the work: image data lies on a fixed union of low dimensional sub-spaces.  Although better than a full-rank Gaussian model, this is still a very crude approximation. As cited by the authors, image data lies on a union of **non-linear** manifolds, which cannot necessarily be accurately approximated by a union of linear manifolds (sub-spaces). The nice linear relationships between N and d will not hold as soon as you have non-linear manifolds. For example, for k =1, you would need N=2 to get a perfect estimate of a one dimensional subspace with a linearity assumption. But as soon as you have a non-linear manifold, depending on the degree of non-linearity you would need larger N. 
I suggest the presentation should be modified to reflect the limitation of the analysis. Otherwise, the way the paper is written now, sets up the reader for disappointment, as the results do not extend to the real image data as claimed by the authors. Nevertheless, I think the results are valuable regardless of the simplistic assumption. They just need to be upfront about the assumptions. 
2. Modeling image distribution with a union of sub-spaces is an old idea in signal processing that goes back to wavelet thresholding and later compressive sensing literature. The paper would benefit from citing major papers where these ideas originated. Importantly, the optimal solutions under this kind of assumption has been a very active topic in that area. It would be interesting to see how these results connect to that literature.   
3. Another assumption made is to approximate the weights with a hard-max operation. This is of course a very hard assumption that results in a lot of error for high noise levels (when noisy image is far from the union of subspaces). Importantly, this assumption simplifies the posterior mean too much, which is counter productive when the goal is to explain diffusion models (where you have large levels of noise). 
4. There seems to be a confusion in the experimental results presented in Figure 5. As a reader, I expect the experimental result to support the theory. However, the generalization result rely on a generalization score that is defined in the appendix. The results shows that the generalization as defined in eq 48 is related to N_k/d_k in a sensible way. However, it does not support nor refutes the results presented in the 4 theories (main results). So there is a divide between the theoretical results and the empirical results which is supposed to support the theory.

---

confidence: 2
contribution: 1
initial_rating: 3
presentation: 3
questions: - Is it possible to show some meaningful bounds on the gap between the modeling using MoLRG and the real data distribution?
- Are the findings in the paper relevant specifically for diffusion models or for generative models in general? If the findings are general, it would strengthen the paper to present similar results across different generative methods. Otherwise, It would be interesting to shed some light on the difference in data size requirements of different generative modeling techniques.
soundness: 2
strengths: - The attempt to model and quantify the training dynamics and data size requirements for diffusion models is noteworthy and possibly impactful.
- The motivation behind the MoLRG is intuitive and easy to follow.
- The paper is well written, making both the math and the figures accessible.
summary: Diffusion models are the dominant class of image generation models. They can effectively learn the underlying image distribution during training, despite the high dimensionality of the data. This paper offers a theoretical modeling of the image distribution using a mixture of low-rank Gaussian. With this model, the authors attempt to explain the training dynamics and failure points of training with small datasets, as well as offer several insights.
weaknesses: - I believe the authors should highlight the paper's contribution. Whether the modeling is justified or not, I believe the paper should contain a meaningful takeaway message. For example, some numerical relationship between the denoiser's jacobian's rank and the number of samples required.
- The experiments conducted do not provide sufficient convincing evidence that the chosen modeling is fitting. Moreover, it is unclear what supports the application of the conclusion following Theorem 4 to real data.
- I am uncertain why modeling the DAE as a mixture of zero-meaned Gaussian is justified. Assuming that the data is a union of linear subspaces, using low-rank Gaussians is reasonable yet the subspaces do not necessarily coincide at the origin. Could the authors please shed more light on this choice?
- The use of the principle components of a Denoiser's jacobian for semantic exploration has been explored in previous work, namely [1]. Also, the connection between the proposed MoLRG and the semantic correspondence of the DAE's jacobian's principle component was unclear.

[1] Hila Manor & Tomer Michaeli (2024). On the Posterior Distribution in Denoising: Application to Uncertainty Quantification. In ICLR 2024.