Transformer-based In-Context Score Operator Learning for Diffusion Models

Wuzhe Xu; Zhaiming Shen; Yulong Lu; Wenjing Liao; Rongjie Lai

cs.LGOct 9, 2026

Transformer-based In-Context Score Operator Learning for Diffusion Models

Authors: Wuzhe Xu, Zhaiming Shen, Yulong Lu, Wenjing Liao, Rongjie Lai

Organizations: South Dakota State University · Georgia Institute of Technology · University of Minnesota · Purdue University

Abstract

Diffusion models typically learn the generative dynamics of a single target distribution or a family of distributions indexed by conditioning variables fixed during training. This limits their ability to adapt to a previously unseen distribution from only a small set of examples. We introduce an in-context solution operator learning framework for diffusion models, in which each target distribution is treated as a task drawn from a meta-distribution and is specified at inference time by an empirical context set. Our method learns an in-context score operator that maps the context samples, a noisy query point, and the diffusion time to the task-dependent score, enabling reverse-diffusion sampling on unseen tasks without retraining or parameter updates. A key observation is that the task-dependent score admits a Gaussian-kernel conditional-mean representation. Its empirical Nadaraya--Watson approximation can be represented exactly by a model in our constructed transformer class, providing a principled foundation for the proposed query-dependent attention architecture. The operator is trained through amortized denoising score matching across task distributions. We establish generalization-error bounds that quantify the effects of the context size, the numbers of training tasks and within-task samples, model complexity, and intrinsic and ambient dimensions. Numerically, we evaluate the method on a range of scientific and visual datasets, including Navier-Stokes, ERA5 climate data, QuickDraw sketches, and Food-101. The results show that the proposed method performs well on relatively simple task families, while its performance degrades as task complexity increases, thereby illustrating both the promise and the limitations of in-context learning for diffusion models.

Figures & tables

\mathfrak{s}_{k}=\{\mathbf{s}_{k,j,l}\}_{j=1,\ldots,J;\,l=1,\ldots,N_{T}},\qquad\mathbf{s}_{k,j,l}=\{\hat{\rho}_{k,j},\,\mathbf{x}_{k,j}(t_{k,j,l}),\,t_{k,j,l}\}.

Algorithm 1 Training of the In-Context Diffusion Model

Panel of Figure 1 — Figure 1: Representative in-distribution inferencesamples for the Navier–Stokes equation. Panels (a)–(c) correspond to viscosities $\nu=1.5\times 10^{-3}$ , $2.2\times 10^{-3}$ , and $3\times 10^{-2}$ , respectively. Within each panel, the rows show prompt samples and samples generated by ICL DM, Cond DM, and SNO, from top to bottom, while the columns show representative realizations. All methods use $n=100$ prompt samples.

Panel of Figure 2 — Figure 2: Scalar-valued empirical PDF of the Navier–Stokes equation in the in-distribution setting. Panels (a)–(c) correspond to $\nu=1.5\times 10^{-3}$ , $2.2\times 10^{-3}$ , and $3\times 10^{-2}$ , respectively. The blue curve denotes the reference distribution, while the red, green, and purple curves represent ICL DM, Cond DM, and SNO, respectively. Greater overlap with the reference curve indicates better agreement with the reference distribution.

Panel of Figure 3 — Figure 3: Out-of-distribution inference samples for Navier–Stokes equation. Panels (a)–(c) correspond to $\nu=1.2\times 10^{-3}$ , $1.3\times 10^{-2}$ , and $6.4\times 10^{-2}$ , respectively. These viscosities lie below, between, and above the two training intervals defined in ( 39 ), respectively. Within each panel, the rows show prompt samples and samples generated by ICL DM, Cond DM, and SNO, from top to bottom.

Panel of Figure 4 — Figure 4: Scalar-valued empirical PDFs for the Navier–Stokes equation in the out-of-distribution setting. Panels (a)–(c) correspond to $\nu=1.2\times 10^{-3}$ , $1.3\times 10^{-2}$ , and $6.4\times 10^{-2}$ , respectively. The blue curve denotes the reference distribution, while the red, green, and purple curves represent ICL DM, Cond DM, and SNO, respectively. Greater overlap with the reference curve indicates better recovery of the target marginal distribution. All methods use $n=100$ prompt samples.

	IND		OOD(low)		OOD(mid)		OOD(high)
	MELRw $\downarrow$	JSD $\downarrow$	MELRw $\downarrow$	JSD $\downarrow$	MELRw $\downarrow$	JSD $\downarrow$	MELRw $\downarrow$	JSD $\downarrow$
ICL DM	1.29e-1	5.35e-4	1.17e-1	4.74e-4	9.79e-2	3.82e-4	2.30e-1	5.92e-4
Cond DM	1.48e-1	5.67e-4	1.34e-1	1.44e-3	1.23e-1	4.14e-4	1.84e-1	8.15e-4
SNO	3.97e-1	2.79e-3	6.05e-1	1.02e-2	1.29e-1	7.16e-4	6.49e-1	3.30e-3

Table 1: Navier–Stokes results comparing ICL DM, Cond DM, and SNO in terms of in-distribution (IND) and out-of-distribution (OOD) generalization, as measured by weighted MELR and the JSD between empirical PDFs. The IND results are averaged over viscosity values sampled according to ( 39 ). The OOD(low), OOD(mid), and OOD(high) settings correspond to

\nu=1.2\times 10^{-3}

1.3\times 10^{-2}

, and

6.4\times 10^{-2}

, respectively. Lower values indicate better performance, and the best result in each column is shown in bold. During both training and evaluation, we use

n=100

prompt samples.

Idx	Method	$\mathbf{W_{2}}~\downarrow$	$\mathbf{MELR_{w}}~\downarrow$	$\mathbf{MELR_{u}}~\downarrow$	$\mathbf{JSD}~\downarrow$
0	ICL DM	2.68e1	2.43e-1	2.32e-1	1.33e-2
	Cond DM	4.39e1	1.38e-1	3.18e-1	1.30e-2
	SNO	3.94e1	2.64e-1	1.85e-1	1.40e-2
10	ICL DM	1.85e1	2.27e-2	2.87e-2	1.54e-3
	Cond DM	3.58e1	2.36e-1	3.01e-1	2.70e-2
	SNO	2.60e1	3.10e-2	3.33e-1	3.37e-3

Table 2: In-distribution quantitative comparison of ICL DM, Cond DM, and SNO across ten ERA5 tasks, where Idx denotes the task defined by a specific variable and pressure level. The discrepancy between the generated and reference distributions is evaluated using the Wasserstein-2 distance (

W_{2}

), weighted and unweighted mean energy log ratios (

\mathrm{MELR}_{w}

and

\mathrm{MELR}_{u}

), and the Jensen–Shannon divergence (JSD). Lower values indicate better agreement with the reference distribution, and the best result for each task and metric is highlighted in bold. All methods use

n=100

prompt samples.

Panel of Figure 5 — Figure 5: In-distribution generation for ERA5 temperature at 150 hPa. The upper panel shows 9 prompt samples followed by 9 samples generated by ICL DM, Cond DM, and SNO, respectively. The lower-left panel compares the scalar-valued empirical PDFs, where blue denotes the reference distribution, while red, green, and purple denote ICL DM, Cond DM, and SNO, respectively. The lower right panel reports the spectral discrepancy $\lvert\log(E_{k}/E_{k}^{\mathrm{ref}})\rvert$ as a function of wavenumber $k$ ; values closer to zero indicate better agreement with the reference energy spectrum. All methods use $n=100$ prompt samples.

Panel of Figure 6 — Figure 6: In-distribution generation for ERA5 humidity at 150 hPa. The upper panel shows 9 prompt fields followed by 9 samples generated by ICL DM, Cond DM, and SNO, respectively. The lower left panel compares the scalar-valued empirical PDFs, where blue denotes the reference distribution, while red, green, and purple denote ICL DM, Cond DM, and SNO, respectively. The lower right panel reports the spectral discrepancy $\lvert\log(E_{k}/E_{k}^{\mathrm{ref}})\rvert$ as a function of wavenumber $k$ ; values closer to zero indicate better agreement with the reference energy spectrum. All methods use $n=100$ prompt samples.

Panel of Figure 7 — Figure 7: In-distribution generation for ERA5 divergence at 150 hPa. The upper panel shows 9 prompt fields followed by 9 samples generated by ICL DM, Cond DM, and SNO, respectively. The lower left panel compares the scalar-valued empirical PDFs, where blue denotes the reference distribution, while red, green, and purple denote ICL DM, Cond DM, and SNO, respectively. The lower right panel reports the spectral discrepancy $\lvert\log(E_{k}/E_{k}^{\mathrm{ref}})\rvert$ as a function of wavenumber $k$ ; values closer to zero indicate better agreement with the reference energy spectrum. All methods use $n=100$ prompt samples.

Panel of Figure 8 — Figure 8: Representative in-distribution generation results on the Quick, Draw! dataset with $K=5$ task categories. Panels (a)–(c) correspond to Eiffel Tower, Mona Lisa, and Airplane, respectively. Within each panel, the rows show prompt sketches and samples generated by ICL DM, Cond DM, and SNO, from top to bottom. All methods use $n=200$ prompt samples, nine of which are displayed.

	Eiffel Tower			Mona Lisa			Airplane
	FID $\downarrow$	FID-QD $\downarrow$	JSD $\downarrow$	FID $\downarrow$	FID-QD $\downarrow$	JSD $\downarrow$	FID $\downarrow$	FID-QD $\downarrow$	JSD $\downarrow$
ICL DM ( $K=5$ )	122.42	49.72	1.87e-3	102.00	48.96	1.34e-3	146.75	67.16	2.76e-3
Cond DM ( $K=5$ )	253.54	96.68	1.81e-3	150.14	109.17	3.77e-3	182.03	104.82	5.22e-3
SNO ( $K=5$ )	328.88	133.77	5.64e-3	195.21	156.29	2.36e-2	207.13	125.73	3.80e-3
ICL DM ( $K=10$ )	144.59	50.72	8.45e-4	103.04	50.51	3.88e-3	146.90	72.76	1.67e-3
Cond DM ( $K=10$ )	287.01	116.46	3.70e-3	150.31	114.13	6.76e-3	200.89	104.64	2.31e-3

Table 3: In-distribution quantitative comparison on three representative QuickDraw categories as the number of meta-training tasks varies over

K\in\{5,10,100\}

. ICL DM, Cond DM, and SNO are evaluated using the standard Fréchet Inception Distance (FID), the QuickDraw-specific FID (FID-QD), and the Jensen–Shannon divergence (JSD) between empirical pixel-value distributions. FID-QD is computed using features extracted by a ResNet-18 pretrained on the QuickDraw dataset. Lower values indicate better distributional agreement, and the best result for each combination of

K

, category, and metric is highlighted in bold. All methods use

n=200

prompt samples.

Panel of Figure 9 — Figure 9: Representative in-distribution generation results on the Food-101 dataset with $K=5$ task categories. Panels (a)–(c) correspond to Apple Pie, Baby Back Ribs, and Beef Carpaccio, respectively. Within each panel, the rows show prompt images and samples generated by ICL DM, Cond DM, and SNO, from top to bottom. All methods use $n=200$ prompt samples, nine of which are displayed.

	Apple Pie			Ribs			Beef Carpaccio
	FID $\downarrow$	FID-Food $\downarrow$	JSD $\downarrow$	FID $\downarrow$	FID-Food $\downarrow$	JSD $\downarrow$	FID $\downarrow$	FID-Food $\downarrow$	JSD $\downarrow$
ICL DM ( $K=5$ )	173.38	42.74	8.18e-3	185.44	45.57	1.23e-2	153.11	62.51	1.74e-2
Cond DM ( $K=5$ )	232.05	67.75	1.38e-2	246.44	84.22	4.16e-2	238.29	120.05	3.00e-2
SNO ( $K=5$ )	231.43	68.22	1.91e-2	260.38	73.82	5.14e-2	230.45	105.07	1.89e-2
ICL DM ( $K=10$ )	201.67	46.49	2.54e-2	195.81	53.34	2.05e-2	165.45	83.27	2.69e-2
Cond DM ( $K=10$ )	231.73	68.71	2.38e-2	241.53	77.92	3.87e-2	236.25	116.03	2.31e-2

Table 4: In-distribution quantitative comparison of ICL DM, Cond DM, and SNO on three representative Food-101 categories as the number of meta-training tasks varies over

K\in\{5,10,100\}

. Performance is evaluated using the standard Fréchet Inception Distance (FID), the Food-101-specific FID (FID-Food), and the Jensen–Shannon divergence (JSD) between empirical pixel-value distributions. FID-Food is computed using features extracted by a ResNet-18 pretrained on the Food-101 dataset. Lower values indicate better distributional agreement, and the best result for each combination of

K

, category, and metric is highlighted in bold. All methods use

n=200

prompt samples.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Meaning
$\mathbf{x}\in\mathbb{R}^{D}$	Data sample (e.g., an image) with ambient dimension $D$ .
$\mathcal{M}\subset\mathbb{R}^{D}$	Data manifold
$d$	Intrinsic dimension of data manifold
$\mu\in\mathcal{P}(\mathcal{P}(\mathcal{\mathbb{R}^{D}}))$	Meta-distribution over tasks (distribution of distributions).
$\rho\sim\mu$	Task distribution sampled from the meta-distribution.
$t\in[\delta,1]$	Diffusion time.

Appendix

Table 5: Notation for in-context learning of diffusion-model score operators.

Dataset	Resolution	$K$	$B$	$n$	$m$	Epochs
NS	128x128	100	1000	100	100	400
ERA5	128x128	100	1000	100	100	400
QuickDraw	64x64	5, 10, 100	2000	200	200	400
Food-101	128x128	5, 10, 100	1000	200	200	2000

Appendix

Table 6: Dataset specifications and sampling parameters used in the experiments. Here,

K

denotes the number of meta-training tasks,

B

the number of samples used for training per task,

n

the number of prompt samples provided to the model,

m

the number of generated samples used for evaluation, and Epochs the number of training epochs used for ICL DM, Cond DM, and SNO.

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,169 papers

Oct 4, 2026

Diffusion Transformers are Provably Optimal In-context Generators

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.

Guoji Fu, Tomoya Wakayama, Ryotaro Kawata +3

National University of Singapore · RIKEN AIP · The University of Tokyo +1

Diffusion Transformer Diffusion Sampling Conditional Generative Modeling

Ask PDF

Sep 8, 2026

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 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 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

Transformer Diffusion Sampling In-Context Learning Attention Mechanisms

Ask PDF

Apr 11, 2026

Transformers Learn the Optimal DDPM Denoiser for Multi-Token GMMs

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.

Hongkang Li, Hancheng Min, Rene Vidal

Department of Electrical and Systems Engineering, University of Pennsylvania, Philadelphia, USA · Institute of Natural Sciences & School of Mathematical Sciences, Shanghai Jiao Tong University, Shanghai, China

Diffusion Models Diffusion-Based Denoising Diffusion Transformer Score Matching

Ask PDF

Abstract

Diffusion Models

Explore similar work

Diffusion Transformers are Provably Optimal In-context Generators

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

Transformers Learn the Optimal DDPM Denoiser for Multi-Token GMMs