Transformers Learn the Optimal DDPM Denoiser for Multi-Token GMMs

Hongkang Li; Hancheng Min; Rene Vidal

cs.LGApr 11, 2026

Transformers Learn the Optimal DDPM Denoiser for Multi-Token GMMs

Authors: Hongkang Li, Hancheng Min, Rene Vidal

Organizations: Department of Electrical and Systems Engineering, University of Pennsylvania, Philadelphia, USA · Institute of Natural Sciences & School of Mathematical Sciences, Shanghai Jiao Tong University, Shanghai, China

Abstract

Transformer-based diffusion models have demonstrated remarkable performance at generating high-quality samples. However, our theoretical understanding of the reasons for this success remains limited. For instance, existing models are typically trained by minimizing a denoising objective, which is equivalent to fitting the score function of the training data. However, we do not know why transformer-based models can match the score function for denoising, or why gradient-based methods converge to the optimal denoising model despite the non-convex loss landscape. To the best of our knowledge, this paper provides the first convergence analysis for training transformer-based diffusion models. More specifically, we consider the population Denoising Diffusion Probabilistic Model (DDPM) objective for denoising data that follow a multi-token Gaussian mixture distribution. We theoretically quantify the required number of tokens per data point and training iterations for the global convergence towards the Bayes optimal risk of the denoising objective, thereby achieving a desired score matching error. A deeper investigation reveals that the self-attention module of the trained transformer implements a mean denoising mechanism that enables the trained model to approximate the oracle Minimum Mean Squared Error (MMSE) estimator of the injected noise in the diffusion steps. Numerical experiments validate these findings.

Figures & tables

Theoretical Works	Network Model	Loss Landscape	Convergence Analysis	Denoising Mechanism
Wang et al. (2024a)	U-Net	Global optimum	✗	PCA
Han et al. (2024)	Fully-connected	N/A	✓	N/A
Wang et al. (2024b)	Fully-connected	N/A	✓	N/A
Han et al. (2025)	Convolutional	Stationary point	✗	Balanced FL
Ours	Transformer	Global optimum	✓	Mean denoising

Table 1: Comparison with existing works about training analysis and the optimality of denoising of diffusion models.

Figure 1: Mean denoising mechanism by the trained Transformer. Attention reduces the noise added to the data. Dark (light) red arrows: attention weights between the query and key that share the same (different) pattern.

Panel of Figure 2 — Figure 2: The convergence performance and the attention behavior of the trained model. (A) The green and red curves are the test loss and score matching error during diffusion model training, respectively. Blue dashed line: Bayes denoising risk. Black dashed line: oracle denoising risk. (B) Excess risk with varying $K$ , the number of Gaussian components per data. (C) Excess risk with varying $\min_{u\in[M]}\tilde{\pi}_{u}$ , i.e., the minimal fraction among all the Gaussian components. A larger $\min_{u\in[M]}\tilde{\pi}_{u}$ indicates a more uniform distribution of all the patterns. (D) Excess risk with different sampling strategies. Red curve: uniform sampling $t\sim\mathrm{Unif}([T])$ . Green curve: sampling from the first $40\%$ time steps, i.e., $t\sim\mathrm{Unif}([1,0.4\cdot T])$ . Blue curve: sampling from the last $40\%$ time steps, i.e., $t\sim\mathrm{Unif}(0.6\cdot T,T)$ . (E) The attention weight summation on keys with the same pattern as the query and on other keys.

Figure 3: FID score of the four generated digits of MNIST. The FID of the minority, “2” decreases more slowly than the others.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Query-key inner products with the same or different patterns, where query patterns are the minimal or maximal of $\tilde{\bm{\pi}}$ . $\min_{u\in[M]}\tilde{\pi}_{u}=0.01$ .

Figure 5: Visualization of the generated digits.

Notations	Annotation
$d$ , $P$	$d$ is the dimension of each token of data. $P$ is the total number of tokens in each data.
$M$ , $K$	$M$ is the total number of mean patterns. $K$ is the number of distinct patterns in each data.
$\tilde{\bm{\pi}}$ , $\{{\bm{\mu}}_{i}\}_{i=1}^{M}$ , $\rho$	$\tilde{\bm{\pi}}$ is the fraction vector of all the Gaussian components of the MTGM distribution. $\{{\bm{\mu}}_{i}\}_{i=1}^{M}$ is the set of all the mean patterns. $\rho^{2}$ is the variance of the Gaussian components.
${\bm{X}}$ , ${\bm{Y}}$ , ${\bm{Z}}$	${\bm{X}}$ denotes the data. ${\bm{Y}}$ ad ${\bm{Z}}$ are the latent variable to define the distribution of ${\bm{X}}$ .
$\Psi$ , ${\bm{W}}$ , $T$ , $\{v_{t}\}_{t=1}^{T}$	$\Psi$ is the set of parameters in the learning model. In our work, $\Psi$ contains ${\bm{W}}$ and $\{v_{t}\}_{t=1}^{T}$ , where ${\bm{W}}$ is the self-attention parameter, and $\{v_{t}\}_{t=1}^{T}$ is the coefficient parameter. $T$ is the total number of diffusion time steps.
${\bm{X}}^{0}$ , $\{\bar{\alpha}_{t}\}_{t=1}^{T}$ , ${\bm{E}}$ , ${\bm{X}}^{t}$	${\bm{X}}^{t}$ is the noisy input at time step $t$ . ${\bm{X}}^{0}$ is the clean input before adding the noise. $\{\bar{\alpha}\}_{t=1}^{T}$ is the noise schedule coefficients. ${\bm{E}}$ is the additive Gaussian noise.

Appendix

Table 2: Summary of Notations

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,136 papers

Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure
Oct 8, 2026Jingda Wu, Changxiao CaiMasked Diffusion Models Diffusion Sampling
Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, USA.
Denoising Diffusion Generative Models Secretly Calculate Attentions
Sep 1, 2026Farzan Haddadi, Leila Monfared, Ebrahim Rezaii +3Attention Control in Diffusion Models Diffusion-Based Image Generation
School of Electrical Engineering, Iran University of Science & Technology, Tehran, IRAN · Independent researcher
Low-dimensional adaptation of diffusion models: Convergence in total variation
Jan 22, 2025Jiadong Liang, Zhihan Huang, Yuxin ChenDiffusion Sampling Diffusion Models
Department of Statistics and Data Science, the Wharton School, University of Pennsylvania

Oct 8, 2026J/K move · Enter open · S save

Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure

Jingda Wu, Changxiao Cai

Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, USA.

Masked Diffusion Models Diffusion Sampling Diffusion Transformer Statistical Learning Theory

Open paper AskPDF arXiv

Masked diffusion models (MDMs) generate discrete data by iteratively unmasking multiple tokens in parallel using a learned mask predictor. Their output carries two errors: parallelization error due to parallel sampling, and estimation error from the mask predictor learned with finite samples. Existing theory treats these errors separately, and the available statistical bounds typically scale with the number

V^L

of length-

L

sequences over a vocabulary of size

V

, suffering from the curse of dimensionality. In this work, we establish end-to-end statistical guarantees showing that transformer-based MDMs learn structured distributions with a sample complexity governed by their latent structure. Specifically, we model this structure through mixtures of hidden Markov trees, which capture latent dependencies within components and heterogeneity across components. When the data distribution is well approximated by such a mixture, a transformer-based MDM with parallel unmasking schedules achieves

\varepsilon

-accurate sampling in KL divergence with

\widetilde{O}(LK_{\mathrm{par}}/\varepsilon)

training samples, where

K_{\mathrm{par}}

is polynomial in

L

V

, the number of hidden states, and the number of mixture components. The transformer network and sampling procedure require no knowledge of the tree topologies or component parameters, thereby adapting to the unknown latent structure automatically. Overall, our results provide theoretical justification for the empirical effectiveness of MDMs on real-world data.

Abstract

Diffusion Models

Explore similar work

Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure

Denoising Diffusion Generative Models Secretly Calculate Attentions

Low-dimensional adaptation of diffusion models: Convergence in total variation