Diffusion Transformers are Provably Optimal In-context Generators

Guoji Fu; Tomoya Wakayama; Ryotaro Kawata; Atsushi Nitanda; Wee Sun Lee; Taiji Suzuki

cs.LGOct 4, 2026

Diffusion Transformers are Provably Optimal In-context Generators

Authors: Guoji Fu, Tomoya Wakayama, Ryotaro Kawata, Atsushi Nitanda, Wee Sun Lee, Taiji Suzuki

Organizations: National University of Singapore · RIKEN AIP · The University of Tokyo · A*STAR CFAR Nanyang Technological University

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.

Figures & tables

Figure 1: In-context generation from demonstrations. A user specifies a desired behavior through a few examples rather than an explicit rule. Because a finite context may be consistent with multiple underlying rules, different outputs can remain plausible for the same query; collapsing the plausible outputs to a single average can produce an invalid result.

Panel of Figure 2 — Figure 2: Conceptual view of the representability construction. The four attention blocks: B1 pools context log likelihoods; B2 copies the pooled values to each latent-task token; B3 forms task logits and attaches diffusion features; and B4 forms posterior-weighted averages of the density and $\sigma_{t}$ -scaled gradient features. A fixed ReLU readout approximates their ratio to produce $\widetilde{s}$ .

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

$k$	Active query rows and keys	Constructed attention logits	Value and residual destination	State after fixed update $\mathsf{R}_{k,n}$
1	Summary row; keys $I_{\mathrm{ctx}}$	All zero, hence weights $1/n$	$\widehat{\bm{\ell}}(c_{i})\in{\mathbb{R}}^{J}$ to $\mathsf{V}_{\mathrm{pool}}$	Summary pool is $\overline{\widehat{\bm{\ell}}}=n^{-1}\sum_{i}\widehat{\bm{\ell}}(c_{i})\in{\mathbb{R}}^{J}$
2	Every latent row; singleton key $i_{\mathrm{sum}}$	Singleton, hence weight $1$	$\overline{\widehat{\bm{\ell}}}\in{\mathbb{R}}^{J}$ to the latent pool block	Every latent row contains the same pooled vector $\overline{\widehat{\bm{\ell}}}\in{\mathbb{R}}^{J}$
3	Every latent row; singleton key $i_{\mathrm{qry}}$	Singleton, hence weight $1$	$\widehat{\mathbf{q}}\in{\mathbb{R}}^{J(d_{z}+2)}$ to the latent query block	Identifier $e_{j}$ selects $(\widehat{\ell}_{j}^{\mathrm{qry}},\widehat{\mathsf{a}}_{j},\widehat{\mathsf{b}}_{j})$ , and $\mathsf{V}_{\mathrm{lat}}$ becomes $(\widehat{\ell}_{j,n}^{\mathrm{post}},\widehat{\mathsf{a}}_{j},\widehat{\mathsf{b}}_{j})\in{\mathbb{R}}^{d_{z}+2}$
4	Output row; keys $I_{\mathrm{lat}}$	$\widehat{\ell}_{j,n}^{\mathrm{post}}$ for latent key $j$	$(\widehat{\mathsf{a}}_{j},\widehat{\mathsf{b}}_{j})\in{\mathbb{R}}^{d_{z}+1}$ to $\mathsf{V}_{\mathrm{out}}$	Output block is the posterior-weighted pair in 69

Appendix

Table 1: Complete computation implemented by the four masked attention blocks. A singleton mask makes its softmax weight exactly one. Values are projected from the indicated source block and injected into the residual destination in 68 .

Figure 3: Scaling with the number of training episodes. Performance at

0.32

0.64

1.28

1.92

, and

2.56

million episodes. From left to right, the panels show the right-probability error, the

1

-Wasserstein distance between the recovered generated-position distribution and the exact two-component Gaussian posterior predictive, and aligned content IoU. Evaluation uses the same

64

fixed test tasks,

n=64

256

generations per task, and the raw model weights. Error bars indicate the standard error across tasks.

Figure 4: Posterior-predictive contraction with context length. Each column corresponds to

n\in\{1,2,4,8,16,32\}

for the same fixed task and query. Histograms show positions recovered from

256

generated images using the raw model weights, and green curves show the analytically exact posterior-predictive distribution. The dotted line at

c=0.5

separates the left and right modes. The model captures the transition from a broad bimodal predictive law at short contexts to a narrow, essentially unimodal law at long contexts, with visible excess right-mode mass in the intermediate regime.

Figure 5: Qualitative posterior-predictive generation. The top half shows a task with a circle query, and the bottom half shows a task with a triangle query; each task is displayed at $n=1$ and $n=8$ . From left to right, each row contains the available context input–output pairs $X_{i}\to Y_{i}$ , the untransformed query input, four independent generations for that query, and the pixel-wise average of $256$ generations. All samples use the raw model weights. For the circle task, the exact/generated right-mode probabilities are $0.550/0.562$ at $n=1$ and $0.999/0.965$ at $n=8$ . For the other task, they are $0.527/0.410$ at $n=1$ and $0.009/0.055$ at $n=8$ . The horizontally blurred averages at $n=1$ reflect bimodality, while the sharp averages at $n=8$ reflect posterior-predictive contraction.

Figure 6: Oracle illustration for the positive sinusoidal family. Left: task uncertainty with nested prompts, compared with the context-free uniform predictor. Right: predictive TV from latent quadrature at

n=256

, compared with the coupled episode-KL upper bound. Means use

5000

independent tasks; bars show

1.96

standard errors across tasks. TV between densities is calculated analytically. This fixed-family diagnostic does not test the worst-case rate, neural training, or reverse sampling.

Observed pattern	Possible bottleneck	Suggested intervention
The task posterior remains spread over tasks with different output laws.	Finite-prompt task uncertainty	Request additional demonstrations that distinguish the plausible task-specific output laws.
The oracle is concentrated, but the model remains diffuse after checking sampler convergence.	Amortization or learning error	Increase independent pretraining episodes; improve model capacity, optimization, or use of context.
Error relative to the oracle varies substantially with sampler steps or schedule.	Sampling error or score–sampler interaction	Refine the sampler steps and schedule; check convergence with the learned score held fixed.
Individual samples are plausible, but their mode frequencies disagree with the oracle.	Probability calibration error	Evaluate mode probabilities, improve episodic distributional training, and check sensitivity to the sampler.
Learning and sampling errors are small relative to finite-prompt task uncertainty.	Prompt information	Prioritize more informative demonstrations over further reductions in learning or sampling error.
A large token budget contains few independent episodes.	Limited independent training data	Track and increase independent episode count; account for within-episode dependence when allocating pretraining data.

Appendix

Table 2: Qualitative diagnostics and resource allocation for in-context generation.

Explore similar work

Transformers as In-Context Samplers: From Closed-Form Diffusion to Estimation-Free Sampling
Date pendingArman Adibi, Alireza Jafari, Mohammad Ghavamzadeh +1Accelerated Sampling Transformer Architectures
School of Computer and Cyber Sciences, Augusta University · Department of Computer Science, University of Virginia · Qualcomm AI Research
Learning to Read the Contextual Tokens in Diffusion Transformers
Oct 5, 2026Omer Dahary, Etai Sella, Hadar Averbuch-Elor +2Diffusion Transformers
Tel Aviv University · Cornell University
Embedding Prediction Helps Image Generation
Oct 1, 2026Sihan Xu, Ji Xie, Zilin Wang +2Diffusion Transformers Image Generation
University of Michigan · Carnegie Mellon University

Date pending, cs.LGJ/K move · Enter open · S save

Transformers as In-Context Samplers: From Closed-Form Diffusion to Estimation-Free Sampling

Arman Adibi, Alireza Jafari, Mohammad Ghavamzadeh, Hadi Daneshmand

School of Computer and Cyber Sciences, Augusta University · Department of Computer Science, University of Virginia · Qualcomm AI Research

Accelerated Sampling Transformer Architectures In-Context Learning In-Context

Open paper AskPDF arXiv

A growing body of work establishes that large language models are not mere statistical memorizers, but are capable of in-context learning: performing inference at test time using only examples provided in the prompt, without any parameter updates. Prior theoretical work has shown that this capability extends to supervised learning tasks such as linear regression. We prove that in-context learning extends further to \emph{data generation}: frozen transformers can simulate iterative generative samplers from in-context samples. We first show that transformers can realize closed-form and smoothed closed-form diffusion samplers. The construction identifies a concrete generative role for softmax attention: it computes responsibility weights and weighted empirical averages, while feedforward layers implement Euler updates. To empirically relate these constructions to pretrained language models, we study \emph{semantic-topic sampling}: prompts consisting of words drawn from a common semantic category, such as animals, foods, or cities. Across transformer layers, the normalized hidden states exhibit a two-stage geometry: they move toward a uniform spherical reference in intermediate layers and then return to structured, topic-dependent representations near the output. We further measure an interacting-particle energy on these hidden-state clouds and observe the same U-shape pattern. We then prove that transformers can approximate an energy-based sampler, constructing the same U-shape energy across the layers.

Abstract

Explore similar work

Transformers as In-Context Samplers: From Closed-Form Diffusion to Estimation-Free Sampling

Learning to Read the Contextual Tokens in Diffusion Transformers

Embedding Prediction Helps Image Generation