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

Date pendingcs.LG

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

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.

Arman Adibi, Alireza Jafari, Mohammad Ghavamzadeh +1

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

Ask PDF

Oct 5, 2026cs.CV

Learning to Read the Contextual Tokens in Diffusion Transformers

Multimodal Diffusion Transformers (MM-DiTs) jointly process visual and textual representations throughout generation. These models repeatedly update the text tokens through multimodal attention, forming dynamic contextual tokens whose function is not well understood. In this work, we introduce a framework for reading this contextual space through natural-language interrogation. We train a lightweight bottleneck network that maps intermediate contextual tokens into the input space of a frozen Large Language Model (LLM), allowing the LLM to answer questions about the emerging image directly from these hidden representations. Our reader reveals that contextual tokens encode a rich, global representation of the emerging scene: generation-specific semantics, including attributes left underspecified by the prompt, are accessible surprisingly early in denoising, while increasingly fine-grained details become readable over time. Remarkably, this information remains decodable even when the MM-DiT receives an empty prompt, showing that contextual tokens accumulate substantial image-specific information from the evolving visual representation itself. We further find that generations with more readable contextual representations tend to receive higher human-preference scores. Building on these observations, we introduce Contextual Alignment, a training technique that explicitly reinforces the visual-semantic information encoded in the contextual tokens, improving generation quality and distributional coverage. Together, our results establish contextual tokens as both an interpretable view into the internal dynamics of MM-DiTs and an effective target for improving generative models.

Omer Dahary, Etai Sella, Hadar Averbuch-Elor +2

Tel Aviv University · Cornell University

Diffusion Transformers

Ask PDF

Oct 1, 2026cs.CV

Embedding Prediction Helps Image Generation

In diffusion transformers, a class label or a text prompt is embedded once, and the same condition is reused at every denoising step. We ask whether predicted embeddings can serve as this condition instead. Next-Embedding Predictive Autoregression (NEPA) trains a Transformer to predict the next continuous embedding in a sequence. In generation, the clean image follows the noisy image, so its embeddings are the next embeddings after the condition and the noisy image. We train a NEPA model to predict them all at once with Multi-Embedding Prediction, and in Embedding Conditioned Generation, a DiT generator is conditioned on these predictions, recomputed at every denoising step, so the conditioning signal adapts to the current noisy state. Experiments on class-conditional ImageNet

256\times256

study the condition of the generator, the design of Multi-Embedding Prediction, and the scaling of both models. The NEPA model adds a second network to every sampling step; with it, and combined with REPA, our final model, NEPA-DiT-XL, reaches an FID of 1.32 using about a third of the training compute of REPA.

Sihan Xu, Ji Xie, Zilin Wang +2

University of Michigan · Carnegie Mellon University

Diffusion Transformers Image Generation Autoregressive Model

Ask PDF