Transformers Learn the Optimal DDPM Denoiser for Multi-Token GMMs

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 \textit{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 \emph{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.

Transformer-based diffusion models can now generate realistic images, audio, and other forms of data. Despite their impressive performance, we still do not fully understand why these models work so well. In particular, current models are trained to remove noise from corrupted data examples, but the mathematical reasons why Transformers can successfully learn this task remain unclear. In this work, we provide one of the first theoretical studies explaining how Transformer-based diffusion models learn during training. We analyze a simplified but meaningful setting where the data consist of multiple patterns mixed together. Our results show that, under suitable conditions, standard training methods can reliably guide the model toward an optimal denoising solution. We also identify how the model parameters of the learned model contribute to this success. In addition to the theoretical analysis, we provide numerical experiments that support our findings. Overall, this work helps build a clearer mathematical understanding of why Transformer-based diffusion models are effective by training, which may guide the design of more efficient and trustworthy generative AI systems in the future.