Transformers Learn the Optimal DDPM Denoiser for Multi-Token GMMs

TL;DR

This paper provides the first convergence analysis of transformer-based diffusion models, showing how they approximate optimal denoising and the conditions needed for convergence in multi-token Gaussian mixtures.

Contribution

It offers a theoretical understanding of why transformers excel in diffusion models, including convergence conditions and the role of self-attention in denoising.

Findings

Transformer models can converge to the Bayes optimal denoising risk.

Self-attention modules implement a mean denoising mechanism.

Numerical experiments validate the theoretical analysis.

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.