Diffusion Transformers are Provably Optimal In-context Generators
Abstract
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ölder 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.
Related
- One-Layer Transformers Provably Learn In-Context K-Nearest Neighbor Prediction with Chain-of-Thought
- Adversarially Trained Linear Transformers Are Optimal Robust In-Context Learners for Gaussian Mixtures
- Transformers as In-Context Samplers: From Closed-Form Diffusion to Estimation-Free Sampling
- One-Layer Transformer Provably Learns Multiclass One-Nearest Neighbor in Context
- Transformer-Based Masked Diffusion Models Provably Learn Distributions with Latent Structure
BibTeX
@misc{fu2026diffusion,
title = {{Diffusion Transformers are Provably Optimal In-context Generators}},
author = {Guoji Fu and Tomoya Wakayama and Ryotaro Kawata and Atsushi Nitanda and Wee Sun Lee and Taiji Suzuki},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.05333},
eprint = {2610.05333},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.05333},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Diffusion Transformers are Provably Optimal In-context Generators
AU - Guoji Fu
AU - Tomoya Wakayama
AU - Ryotaro Kawata
AU - Atsushi Nitanda
AU - Wee Sun Lee
AU - Taiji Suzuki
PY - 2026
DA - 2026/10/04
PB - arXiv
DO - 10.48550/arxiv.2610.05333
UR - https://arxiv.org/abs/2610.05333
AB - 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ölder 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.
ER -