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

TL;DR

This paper analyzes the statistical approximation capabilities and computational efficiency of latent Diffusion Transformers (DiTs) under low-dimensional assumptions, providing bounds, criteria, and algorithms for faster inference and training.

Contribution

It offers new theoretical bounds on DiTs score function approximation, distribution recovery, and proposes efficient algorithms for inference and training leveraging low-rank structures.

Findings

Sub-linear approximation error bound in latent space dimension

Almost-linear time inference algorithms for latent DiTs

Low-rank gradient approximations enable faster training

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.