# Diffusion Transformers are Provably Optimal In-context Generators

arXiv: https://arxiv.org/abs/2610.05333 · HTML: https://mlsanity.com/paper/2610.05333/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 8 of its 74 references that are in the map.
- First posted on arXiv 4 Oct 2026.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

## Builds on (most cited first, 8 of 8)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Convergence of Diffusion Models Under the Manifold Hypothesis in High-Dimensions](https://mlsanity.com/paper/2409.18804/) (2024)
- [Unveil Conditional Diffusion Models with Classifier-free Guidance: A Sharp Statistical Theory](https://mlsanity.com/paper/2403.11968/) (2024)
- [In-Context Learning Is Provably Bayesian Inference: A Generalization Theory for Meta-Learning](https://mlsanity.com/paper/2510.10981/) (2025)
- [Optimal In-context Adaptivity and Distributional Robustness of Transformers](https://mlsanity.com/paper/2510.23254/) (2025)
- [Stochastic control, entropic interpolation and gradient flows on Wasserstein product spaces](https://mlsanity.com/paper/1601.04891/) (2016)
- [Intrinsic Wasserstein Rates for Score-Based Generative Models on Smooth Manifolds](https://mlsanity.com/paper/2605.15822/) (2026)
- [Transformers as In-Context Samplers: From Closed-Form Diffusion to Estimation-Free Sampling](https://mlsanity.com/paper/2609.08981/) (2026)

## Abstract (arXiv)

Generative foundation models are attracting interest for their ability to produce desired outputs from demonstrations given at inference time, without updating parameters. However, since a few demonstrations cannot uniquely identify the intended task, the challenge is how to learn and sample from an output distribution that reflects this task uncertainty. In this work, we theoretically analyze how a Diffusion Transformer (DiT), pretrained across diverse tasks, learns and generates predictive distributions for a new query from demonstrations. We first show that the natural target to generate from finite demonstrations is not an output derived from estimating a single task, but rather a predictive distribution that captures the task uncertainty remaining after observing the demonstrations. We then prove that a DiT can learn this predictive distribution through score estimation, using attention to aggregate information from demonstrations and diffusion to generate samples. Owing to this property, with sufficient pretraining resources and diffusion sampling steps, the resulting DiT achieves the minimax optimal rate over a H"older class of test-time tasks. These results imply that DiT acts as a statistically grounded in-context generator capable of generating distributions adapted to new tasks while retaining the uncertainty inherent in finite demonstrations.