Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

On Statistical Rates and Provably Efficient Criteria of Latent Diffusion Transformers (DiTs)

Jerry Yao-Chieh Hu, Weimin Wu, Zhao Song, Han Liu

TL;DR

Under the low-dimensional assumption, it is shown that the statistical rates and the computational efficiency are all dominated by the dimension of the subspace, suggesting that latent DiTs have the potential to bypass the challenges associated with the high dimensionality of initial data.

Abstract

We investigate the statistical and computational limits of latent Diffusion Transformers (DiTs) under the low-dimensional linear latent space assumption. Statistically, we study the universal approximation and sample complexity of the DiTs score function, as well as the distribution recovery property of the initial data. Specifically, under mild data assumptions, we derive an approximation error bound for the score network of latent DiTs, which is sub-linear in the latent space dimension. Additionally, we derive the corresponding sample complexity bound and show that the data distribution generated from the estimated score function converges toward a proximate area of the original one. Computationally, we characterize the hardness of both forward inference and backward computation of latent DiTs, assuming the Strong Exponential Time Hypothesis (SETH). For forward inference, we identify efficient criteria for all possible latent DiTs inference algorithms and showcase our theory by pushing the efficiency toward almost-linear time inference. For backward computation, we leverage the low-rank structure within the gradient computation of DiTs training for possible algorithmic speedup. Specifically, we show that such speedup achieves almost-linear time latent DiTs training by casting the DiTs gradient as a series of chained low-rank approximations with bounded error. Under the low-dimensional assumption, we show that the statistical rates and the computational efficiency are all dominated by the dimension of the subspace, suggesting that latent DiTs have the potential to bypass the challenges associated with the high dimensionality of initial data.

On Statistical Rates and Provably Efficient Criteria of Latent Diffusion Transformers (DiTs)

TL;DR

Under the low-dimensional assumption, it is shown that the statistical rates and the computational efficiency are all dominated by the dimension of the subspace, suggesting that latent DiTs have the potential to bypass the challenges associated with the high dimensionality of initial data.

Abstract

We investigate the statistical and computational limits of latent Diffusion Transformers (DiTs) under the low-dimensional linear latent space assumption. Statistically, we study the universal approximation and sample complexity of the DiTs score function, as well as the distribution recovery property of the initial data. Specifically, under mild data assumptions, we derive an approximation error bound for the score network of latent DiTs, which is sub-linear in the latent space dimension. Additionally, we derive the corresponding sample complexity bound and show that the data distribution generated from the estimated score function converges toward a proximate area of the original one. Computationally, we characterize the hardness of both forward inference and backward computation of latent DiTs, assuming the Strong Exponential Time Hypothesis (SETH). For forward inference, we identify efficient criteria for all possible latent DiTs inference algorithms and showcase our theory by pushing the efficiency toward almost-linear time inference. For backward computation, we leverage the low-rank structure within the gradient computation of DiTs training for possible algorithmic speedup. Specifically, we show that such speedup achieves almost-linear time latent DiTs training by casting the DiTs gradient as a series of chained low-rank approximations with bounded error. Under the low-dimensional assumption, we show that the statistical rates and the computational efficiency are all dominated by the dimension of the subspace, suggesting that latent DiTs have the potential to bypass the challenges associated with the high dimensionality of initial data.
Paper Structure (82 sections, 33 theorems, 141 equations, 1 figure, 1 table)

This paper contains 82 sections, 33 theorems, 141 equations, 1 figure, 1 table.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Organization.
  4. Notations.
  5. Background
  6. Score-Matching Denoising Diffusion Models
  7. Forward and Backward Process.
  8. Score Matching.
  9. Score Decomposition in Linear Latent Space
  10. Score Network and Transformers
  11. (Latent) Score Network.
  12. Transformers.
  13. Statistical Rates of Latent DiTs with Subspace Data Assumption
  14. DiT Score Network Class
  15. Score Approximation of DiT
  16. ...and 67 more sections

Key Result

Lemma 2.1

Let data $x = Bh$ follow assumption:1. The decomposition of score function $\nabla \log p_t(\Bar{x})$ is where $p_t^{h}(\Bar{h}) \coloneqq \int \psi_t(\Bar{h}|h)p_h(h) \dd h$, $\psi_t( \cdot | h)$ is the Gaussian density function of $N(\beta(t)h, \sigma(t)I_{d_0})$, $\beta(t) = e^{-t/2}$ and $\sigma(t) = 1 - e^{-t}$. We restate the proof in pf:subspace_score for completeness.

Figures (1)

  • Figure 1: Overview of DiT Score Network Architecture$s_W(\cdot, t)$. $W_B^T$ denotes the linear layer from the input data space to the linear latent space. $f(\cdot) = {R^{-1}\circ f_{\mathcal{T}} \circ R}(\cdot)$ denotes the transformer network $f_{\mathcal{T}}(\cdot)$ with reshaping layer $R(\cdot)$, where $f_{\mathcal{T}}(\cdot) \in \mathcal{T}_p^{r,m,l}$. $W_B$ denotes the linear layer from the linear latent space to the input data space. $\sigma(t)$ denote the variance of the conditional distribution $P(x_t \mid x_0)$.

Theorems & Definitions (100)

  • Remark 2.1
  • Lemma 2.1: Score Decomposition, Lemma 1 of chen2023score
  • Definition 3.1: DiT Reshape Layer $R(\cdot)$
  • Definition 3.2: Transformer Network Class $\mathcal{T}_p^{r,m,l}$
  • Definition 3.3: DiT Score Network Class $\mathcal{S}_{\mathcal{T}_p^{r,m,l}}$ (\ref{['fig:pipeline']})
  • Theorem 3.1: Score Approximation of DiT
  • proof : Proof Sketch
  • Remark 3.1: Comparing with Existing Works
  • Remark 3.2: Latent Dimension Dependency
  • Theorem 3.2: Score Estimation of DiT
  • ...and 90 more