On Statistical Rates of Conditional Diffusion Transformers: Approximation, Estimation and Minimax Optimality

TL;DR

It is shown that both conditional DiTs and their latent variants lead to the minimax optimality of unconditional DiTs under identified settings, and that latent conditional DiTs achieve lower bounds than conditional DiTs both in approximation and estimation.

Abstract

We investigate the approximation and estimation rates of conditional diffusion transformers (DiTs) with classifier-free guidance. We present a comprehensive analysis for ``in-context'' conditional DiTs under four common data assumptions. We show that both conditional DiTs and their latent variants lead to the minimax optimality of unconditional DiTs under identified settings. Specifically, we discretize the input domains into infinitesimal grids and then perform a term-by-term Taylor expansion on the conditional diffusion score function under Hölder smooth data assumption. This enables fine-grained use of transformers' universal approximation through a more detailed piecewise constant approximation and hence obtains tighter bounds. Additionally, we extend our analysis to the latent setting under the linear latent subspace assumption. We not only show that latent conditional DiTs achieve lower bounds than conditional DiTs both in approximation and estimation, but also show the minimax optimality of latent unconditional DiTs. Our findings establish statistical limits for conditional and unconditional DiTs, and offer practical guidance toward developing more efficient and accurate DiT models.

Figures (3)

• Figure 1: Conditional DiT Network Architecture. The architecture consists of a reshape layer $R(\cdot)$, a reversed reshape layer $R^{-1}(\cdot)$, and the embedding layers for label $y$ and timestep $t$. The embeddings of $y$ and $t$ are concatenated with input sequences and then processed by a transformer network $f_{\mathcal{T}} \in \mathcal{T}^{\textcolor{blue}{h,s,r}}$.
• Figure 2: Approximate Score Function with Transformer $\mathcal{T}_{\text{score}}$ under \ref{['assumption:conditional_density_function_assumption_1']}. The construction consists of the transformers to approximate local polynomials $f_1$ and $f_2$, and the algebraic operators. We highlight the overall term-by-term approximations and their corresponding lemmas to ensemble the transformers.
• Figure 3: Network Architecture of Latent Conditional DiT. The overall architecture consists of linear layer of encoder and decoder $W_U^{\top}$ and $W_U$ that transform input $x\in \mathbb{R}^{d_x}$ into linear latent space $\mathbb{R}^{d_0}$, reshaping layer $\widetilde{R}(\cdot)$ and $\widetilde{R}^{-1}(\cdot)$, embedding layer for label $y$ and timestep $t$. The embedding concatenates with input sequences and processes by the adapted transformer network ${\mathcal{T}}^{\textcolor{blue}{h,s,r}}_{\widetilde{R}} = \widetilde{R}^{-1} \circ g_{\mathcal{T}} \circ f^{({\rm FF})} \circ \widetilde{R}$.

Theorems & Definitions (157)

• Definition 2.1: Transformer Block
• Definition 2.2: Transformer Network Function Class
• Definition 2.3: DiT Reshape Layer $R(\cdot)$
• Definition 2.4: Transformer Network Function Class with Reshape Layer $\mathcal{T}_R^{\textcolor{blue}{h,s,r}}$
• Definition 3.1: Hölder Space
• Remark 3.1
• Theorem 3.1: Conditional Score Approximation under \ref{['assumption:conditional_density_function_assumption_1']}
• Remark 3.2
• proof : Proof Sketch
• Remark 3.3: Approximation Rate