arXiv CS open access

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

Wuzhe Xu et al. · arxiv_cs

arXiv CS · Papers

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-

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

T RANSFORMER - BASED I N -C ONTEXT S CORE O PERATOR L EARNING FOR D IFFUSION M ODELS A P REPRINT

arXiv:2610.13690v1 [cs.LG] 9 Oct 2026

Wuzhe Xu South Dakota State University wuzhe.xu@sdstate.edu

Zhaiming Shen Georgia Institute of Technology zshen49@gatech.edu

Wenjing Liao Georgia Institute of Technology wliao@gatech.edu

Yulong Lu University of Minnesota yulonglu@umn.edu Rongjie Lai Purdue University lairj@purdue.edu

October 12, 2026

A BSTRACT 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 taskdependent 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 [12], QuickDraw sketches [20], and Food-101 [3]. 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.

Introduction

Diffusion models have emerged as a powerful class of generative models, achieving strong empirical performance across a broad range of domains. They have been widely adopted in computer vision [42, 13, 43, 21, 2, 36], audio generation [23, 32, 18, 19], and reinforcement learning and control [35, 38, 56]. More recently, their applications have expanded to scientific and engineering problems, including data purification [34, 48, 49, 52], climate-data downscaling [44, 41], and partial differential equation (PDE) problems [17, 15, 51, 29]. Their success has been further enhanced by conditional generation techniques such as classifier guidance [8], classifier-free guidance [14], and conditioned latent diffusion [39], which provide substantial control over the generated samples. Despite this progress, standard diffusion models typically learn the reverse generative dynamics associated with the distribution represented by a particular training dataset [42, 13, 43]. Conditional diffusion models extend this formulation to families of distributions indexed by predefined conditioning variables, such as class labels, classifier outputs, text embeddings, or other task-specific signals [8, 14, 39]. Although these models can generate from multi-

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

A P REPRINT

ple conditional distributions, the form and semantics of the conditioning mechanism are prescribed during training. Consequently, they are not generally designed to infer a previously unseen target distribution solely from a small, unordered set of contextual samples. Adapting to such a distribution may therefore require retraining, fine-tuning, or the construction of an additional conditioning representation. Recent few-shot and score-operator approaches have begun to relax this restriction [11, 27], but learning a general in-context score operator directly from empirical context sets remains underexplored. Many applications, however, require generative models to operate over an entire family of target distributions. In scientific computing, for example, the relevant solution or data distribution may vary with the governing equation, physical parameters, forcing terms, initial or boundary conditions, or available observations. This family-level structure is central to operator-learning methods such as DeepONet and the Fourier neural operator, which learn mappings between problem specifications and their corresponding solutions [28, 26, 24]. Training an independent diffusion model for every possible configuration is computationally prohibitive and fails to exploit the structure shared across related problems. This motivates a solution-operator perspective for diffusion models: rather than learning the score function associated with a single prescribed target distribution, the objective is to learn an operator that maps information specifying a target distribution to its corresponding time-dependent score function [27]. Such an operator should generalize across a family of distributions and, ideally, adapt to previously unseen members of that family without retraining or parameter updates. In this work, we introduce an in-context solution operator learning framework for diffusion models. Each target distribution, referred to as a task, is drawn from a meta-distribution over probability measures. More precisely, we learn an in-context score operator that maps the empirical context distribution, together with a noisy query point and the diffusion time, to the task-dependent score. This formulation casts task adaptation as an operator-evaluation problem: the context set identifies the target distribution, while the shared operator captures common structure across the task family. It thus extends diffusion modeling from learning a score function for a fixed target distribution to learning a reusable score operator over a family of distributions. The solution operator is trained through amortized denoising score matching across task distributions, allowing knowledge to be shared across the entire task family. Once trained, it can perform reverse diffusion for a previously unseen task directly from its prompt samples, without retraining, fine-tuning, or task-specific network parameter updates. This perspective is also closely related to recent work on in-context operator learning over probability measures [7]. Our method operates directly on the empirical context set and constructs score estimates through query-dependent interactions among the context samples, the noisy query point, and the diffusion time. A key observation is that the task-dependent score can be expressed through a Gaussian-kernel conditional-mean representation. Replacing the population conditional mean with its empirical counterpart yields a Nadaraya–Watson estimator based on the context samples. We show that this empirical estimator can be represented exactly by a model in our constructed transformer class. This connection provides a principled motivation for the proposed transformer-based in-context score operator, rather than treating the transformer architecture as an ad hoc design choice. Theoretically, We further establish generalization-error bounds for the learned score operator in Section 4. This bounds quantify the dependence on the context size n, the number of training tasks K, the number of episodes per task J, the number of diffusion samples per episode NT , and the intrinsic and ambient dimensions d and D. The finitecontext error decays as n−2/(d+6) , with the intrinsic dimension governing the exponent and the ambient dimension entering only through polynomial and logarithmic factors. This distinction indicates that in-context score estimation is governed primarily by the intrinsic geometry of the task distributions, exhibiting a curse of intrinsic rather than ambient dimensionality. The statistical error contains terms of order K −1/2 , (KJ)−1/2 , and (KJNT )−1/2 , separating the effects of the numbers of tasks, episodes, and diffusion samples. In particular, the task-level persistent K −1/2 term is unaffected by J and NT . Thus, within the present upper-bound analysis, increasing within-task sampling alone cannot reduce this contribution, highlighting the number and diversity of training tasks as fundamental to generalization.. We summarize our main contributions as follows: • We formulate in-context learning for diffusion models at the level of task distributions. Each task is modeled as a probability distribution drawn from a meta-distribution, and a previously unseen task is specified at inference time by a small context set represented as an empirical measure. • We show that the task-dependent score can be expressed in terms of a Gaussian-kernel conditional mean. We further prove that its empirical Nadaraya–Watson approximation, and hence the corresponding empirical score estimator, can be represented exactly by a model in our constructed transformer class. This connection provides a principled foundation for the proposed transformer-based in-context score operator. Trained through amortized denoising score matching across task distributions, the operator enables reverse-diffusion sampling on unseen tasks without retraining, fine-tuning, or task-specific parameter updates. 2

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

A P REPRINT

• We provide a theoretical analysis of the proposed neural operator and establish a generalization bound under suitable assumptions. The bound separates the contributions of finite context, finite meta-training data, and model complexity, while making explicit the dependence of the in-distribution generalization error on the context size, the numbers of training tasks and samples, and the intrinsic and ambient dimensions of the problem. • We evaluate the proposed method on diverse scientific and visual datasets, including Navier–Stokes solution fields, ERA5 climate data, QuickDraw sketches, and Food-101 images. The results demonstrate strong performance in low-intrinsic-dimension regimes and reveal a gradual degradation as task complexity increases. For low-intrinsic-dimension problems such as Navier–Stokes, the method also exhibits modest generalization to out-of-distribution tasks. The code is available at https://github.com/woodssss/In-context-learning-for-Diffusion-Models. The repository also provides a Lean 4 formalization of Theorem 1. 1.1

Related work

Our framework is closely related to the broader paradigm of in-context learning (ICL), in which a model adapts to a new task through contextual demonstrations without updating its parameters. ICL was initially studied extensively in large language models [4, 31, 55, 30], and related ideas have since been developed for scientific computing and operator learning. ICON [54] studies supervised in-context operator learning for differential equations, using prompted input-output examples to adapt a single network to previously unseen operators. PROSE-PDE [45] takes a complementary multimodal approach, combining numerical solution data with symbolic equation representations for multi-operator learning. Self-supervised in-context operator learning in probability-measure spaces has also been investigated for optimal transport, mean-field games, and Wasserstein gradient flows [16, 9, 7]. This measure-valued task perspective is especially relevant to our formulation, in which each task is itself a target probability distribution. In parallel, several works have incorporated contextual demonstrations into diffusion-based image generation. Prompt Diffusion [47], iPromptDiff [5], and Context Diffusion [33] demonstrate that diffusion models can use visual examples as context for image generation and editing. These methods, however, primarily address visually prompted imageto-image generation in paired or query-conditioned settings. They do not formulate diffusion modeling as in-context learning of a solution operator over a family of target probability distributions, which is the focus of the present work. Other two related directions are Score Neural Operator (SNO) [27] and few-shot diffusion models [11], both of which can be viewed as forms of context-conditioned generation. In SNO, the prompt set is first summarized by a distributionlevel representation, specifically a kernel mean embedding, and the score function is then predicted conditioned on this embedding. In few-shot diffusion models, the support set is encoded into a task-dependent feature, for example by aggregating patch-level information through an encoder, and diffusion sampling is performed conditioned on this feature. In contrast, our method operates directly on the empirical context set and constructs score estimates through query-dependent interactions among the context samples, the noisy query point, and the diffusion time, rather than conditioning solely on a compressed global representation. A growing theoretical literature seeks to explain the mechanisms and generalization properties underlying ICL. Existing studies interpret ICL as implicit Bayesian inference [50], characterize the function classes that transformers can learn from demonstrations [10], and connect transformer forward passes to gradient-based meta-learning [46]. Other works show that transformers can implement statistical learning algorithms in context [1] and study their generalization and stability from an algorithm-learning perspective [25]. The diversity and number of pretraining tasks have also been identified as central factors governing the emergence of ICL and its generalization to unseen tasks [37]. In scientific computing, [6] develops an operator-learning analysis of ICL for linear systems and elliptic PDEs. [7] studies ICL directly on probability-measure spaces and establishes generalization-error bounds for in-context operator learning of optimal transport maps. Particularly relevant to our analysis is [40], which connects attention mechanisms to kernel regression on structured manifolds and establishes generalization guarantees governed by intrinsic rather than ambient dimension. Our theory complements these results by analyzing an in-context score operator and separating the effects of finite context, the number of meta-training tasks, within-task sampling, model complexity, and intrinsic dimension. Organization. The remainder of this paper is organized as follows. Section 2 introduces the problem setup and notation. Section 3 presents the proposed method, including the training objective and algorithm. Section 4 provides a theoretical analysis of the generalization error. Section 5 reports empirical results across a range of datasets. Finally, Section 6 concludes with a discussion of limitations and directions for future work. 3

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

A P REPRINT

Problem Setup

In this section, we first review the background of diffusion models [43, 53] and then formalize the in-context score operator learning problem. 2.1

Background on Diffusion Models

The score-based diffusion models [43, 53] aim to approximate a target data distribution ρ(x) given a dataset {xi }N i=1 , and it operates in two stages: a forward diffusion process and a reverse generative process. Forward Process. In the forward process, data is progressively corrupted with noise through a stochastic differential equation (SDE), transforming the original data into a standard Gaussian distribution. Specifically, for a task ρ ∈ P(RD ), consider a special case of the VP SDE: p (1) dx = − 12 β(t)x dt + β(t) dwt , The perturbed solution of this SDE at time t with initial condition x ∼ ρ is: x(t) := α(t)x + σ(t)ε,

ε ∼ N (0, I).

(2)

where β(t) is a user-specific monotonically increasing function for t ∈ [δ, 1], Z 1 t α(t) = exp − β(s) ds , σ 2 (t) = 1 − α2 (t). 2 0 And we denote ρt by the law of x(t), pt by the density function of ρt . Reverse Process. The reverse process is described by the corresponding reverse time SDE that progressively transforms the standard Gaussian distribution back into the original data distribution. p 1 dx = − β(t)x − β(t)∇x log pt (x) dt + β(t)dw̄, (3) 2 where dw̄ represents a reverse-time Wiener process and log pt (x) is the score function of the marginal distribution of the forward process at time t. Moreover, Song et al. [43] proved the existence of the probability flow (PF) ODE, which shares the same marginal distributions as the reverse-time SDE (3): 1 1 dx = − β(t)x − β(t)∇x log pt (x). dt 2 2

(4)

In this paper, we let the noise scheduling function be β(t) = β0 + (β1 − β0 )t, with β0 = 0.1 and β1 = 20. Score matching. Before introducing our in-context learning objective, we briefly recall the standard denoising score matching formulation. The goal is to learn the time-dependent score field ∇x log pt (x). A standard denoising score matching objective trains a score network sθ by minimizing " # 2 ε Lρ (sθ ) = Ex∼ρ, t∼U[δ,1], ε∼N (0,I) sθ (x(t), t) + (5) σ(t) 2 Here we restrict t to [δ, 1] for some fixed δ > 0 to stay away from the zero-noise endpoint t = 0, where the data distribution may be singular and the score can become ill-behaved as σ(t) → 0. The terminal endpoint t = 1, however, does not pose such a difficulty. Although its law is not exactly standard Gaussian, it is well approximated by R1 a standard Gaussian when the accumulated noise level 0 β(s) ds is sufficiently large, since α(1) ≈ 0, σ(1) ≈ 1 and x(1) = α(1)x + σ(1)ε. Generation. For the generation phase, any black-box ODE solver may be used to integrate (4) initiated at a point sampled from the standard Gaussian, making probability flow (PF) ODE sampling computationally efficient. Alternatively, one may generate samples by integrating the reverse-time SDE (3) with a general SDE solver. In this work, we discretize the reverse-time SDE using the Euler–Maruyama scheme. 4

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

2.2

A P REPRINT

In-Context Score Function Learning

Standard denoising score matching treats a fixed target distribution and learns a single time-dependent score function for it. In many applications, however, generation is required not for a single fixed distribution but for a family of related task distributions, each corresponding to a different task, condition, or physical parameter. Learning an independent score model for each distribution would fail to exploit shared structure across the family and would provide no mechanism for adapting to previously unseen tasks. We therefore amortize score learning across the task family by learning a score operator that maps each task distribution to its corresponding score function. More specifically, we assume that tasks are sampled from a meta-distribution over task distributions. µ ∈ P(P(RD )). Thus, each task ρ ∼ µ is itself a probability distribution on RD . Given a diffusion time interval t ∈ [δ, 1] with δ > 0. for each task ρ, we define the perturbed sample x ∼ ρ,

x(t) = α(t)x + σ(t)ε,

ε ∼ N (0, I),

and let ρt denote the law of x(t). Denote by pρ,t the density of ρt . The corresponding task-dependent score function is then defined as Sρ (x, t) := ∇x log pρ,t (x), (x, t) ∈ RD × [δ, 1]. In this way, each task distribution ρ gives rise to a score function Sρ , which naturally defines the following operator n o T : P(RD ) → S : RD × [δ, 1] → RD , ρ 7→ T (ρ) = Sρ . At the continuous level, we seek to learn an optimal score operator by amortized denoising score matching over the meta-distribution of tasks. This leads to the population loss " # 2 ε x∼ρ, T (ρ) x(t), t + L(T ) := Eρ∼µ [Lρ (T (ρ))] = Eρ∼µ E , (6) σ(t) 2 t∼U[δ,1], ε∼N (0,I) and we denote T ∗ a population-risk minimizer. This formulation yields a distribution-level in-context score-learning problem in which the task distribution itself acts as the context. In practice, however, the learner does not observe the meta distribution ρ directly; it only receives a finite collection of samples drawn from ρ. The corresponding finite-context problem is therefore to infer, from these examples, the score function associated with the underlying task distribution and to evaluate it at new noisy inputs and diffusion times. A learned approximation to T ∗ can consequently adapt to previously unseen task distributions through their observed context samples, without requiring task-specific parameter updates. In the next section, we will provide details on how to compute this model.

Methodology

In Section 2.2, we formulated the in-context score learning problem at the continuous level, where the optimal map is defined on the space of task distributions and the learning objective is given by the population loss. To approximate this map in practice, we parameterize it by a neural operator Tθ ∈ T , where T is the hypothesis space parameterized by a deep neural network and will be specified in Section 4.1. Since the underlying task distribution is not directly available to the model, we must further replace it by a finite empirical representation. This leads to a sequence of loss functions at four successive levels of discretization: the semi-discretized population loss, the task-discretized loss, the task- and episode-discretized loss, and the fully discretized loss. Semi-discretized Population Loss. In practice, however, the underlying task distribution is not directly observed, and the model cannot operate on it explicitly. Instead, both during training and at test time, each task is specified only through finitely many samples. We therefore now turn to a semi-discretized formulation, in which each task is represented by a finite context set and the population-level objective is replaced by a practical training loss that can be optimized from sampled tasks. Concretely, for a task ρ ∼ µ, we observe Pn a context set (prompt) of n i.i.d. samples x(1) , . . . , x(n) ∼ ρ and represent it by the empirical measure ρ̂ := n1 i=1 δx(i) together with a query point (x(t), t) that serves as the input of the transformer network. Namely, we have a transformer network: Tθ : (RD )n × RD × [δ, 1] → RD ,

(x(1) , . . . , x(n) ; x(t), t) 7→ Tθ (ρ̂)(x(t), t) 5

(7)

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

A P REPRINT

where x(t) is provided by equation (2). We write the corresponding score matching loss for the task ρ as # " 2 ε n , Tθ (ρ̂) x(t), t + Lρ (Tθ ) = E ρ̂∼ρ⊗n , x∼ρ, σ(t) 2 t∼U[δ,1], ε∼N (0,I)

(8)

This leads to the following semi-discretized population operator learning loss: " Ln (Tθ ) := Eρ∼µ [Lnρ (Tθ )] = Eρ∼µ E

ρ̂∼ρ⊗n , x∼ρ, t∼U[δ,1], ε∼N (0,I)

# 2 ε , Tθ (ρ̂) x(t), t + σ(t) 2

(9)

Task-discretized loss. We next discretize the outer expectation over tasks in (9) by replacing it with the empirical average over K training tasks {ρk }K k=1 , sampled i.i.d. from the meta-distribution µ. This gives # " K 2 1 X ε Ln,K (Tθ ) := Tθ (ρ̂k ) xk (t), t + . (10) E ρ̂k ∼ρ⊗n , xk ∼ρk , k K σ(t) 2 t∼U[δ,1], ε∼N (0,I) k=1

Task- and episode-discretized loss. For each sampled task ρk , we further discretize the expectation over contextquery draws by replacing it with a finite average over J sampled episodes indexed by j. Let ρ̂k,j denote the empirical measure of the context set in the j-th episode for task k, and let xk,j ∼ ρk be the corresponding query sample. For each pair (k, j), define xk,j (t) := α(t)xk,j + σ(t)ε, t ∼ U[δ, 1], ε ∼ N (0, I). Then we define " # J K 2 1 XX ε Ln,K,J (Tθ ) := Et∼U[δ,1], ε∼N (0,I) Tθ (ρ̂k,j ) xk,j (t), t + . (11) KJ σ(t) 2 j=1 k=1

Fully discretized empirical loss. Finally, we introduce the fully discretized empirical loss, which serves as the computable training objective in practice. Let {ρk }K k=1 be training tasks drawn i.i.d. from the meta-distribution µ. For each task ρk , we construct an episodic training set and define S := {sk }K k=1 ,

sk := {sk,j,l }j=1,...,J; l=1,...,NT ,

(12)

where j indexes independent context–query pair draws within task k, and l indexes independent draws of diffusion time and noise. More precisely, for each triplet (k, j, l), we first draw n + 1 i.i.d. samples from ρk , (1)

(n)

i.i.d.

xk,j , . . . , xk,j , xk,j ∼ ρk , (1)

(n)

where xk,j , . . . , xk,j are used to form the context and xk,j serves as the clean query sample. We then sample a diffusion time tk,j,l ∼ U[δ, 1] and noise εk,j,l ∼ N (0, I), and construct the corresponding noised query xk,j (tk,j,l ) := α(tk,j,l ) xk,j + σ(tk,j,l ) εk,j,l . Pn The associated empirical context is ρ̂k,j := n1 i=1 δx(i) . Hence each training episode is represented by the tuple k,j

n o sk,j,l = ρ̂k,j , xk,j (tk,j,l ), tk,j,l . By replacing the remaining expectation over the diffusion time and Gaussian noise in (11) with the empirical average over the NT sampled pairs, we obtain the fully discretized objective: K

T XXX 1 εk,j,l LS (θ) = Tθ (ρ̂k,j ) xk,j (tk,j,l ), tk,j,l + . KJNT σ(t k,j,l ) 2 j=1

k=1

(13)

l=1

which serves as the computable training loss used in practice. We then define the empirical loss minimizer by Tθ∗ ,S ∈ arg min LS (θ). Tθ ∈T

In practice, training is performed in an amortized manner across tasks by minimizing the empirical denoising scorematching loss. At inference time, given the empirical context of a previously unseen task, the learned neural operator is 6

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

A P REPRINT

used as a score estimator in the reverse-time SDE, from which new samples are generated using the Euler–Maruyama scheme. The training and sampling procedures are summarized in Algorithms 1 and 2, respectively. Algorithm 1 Training of the In-Context Diffusion Model Input: training set S = {sk }K k=1 , neural operator Tθ , number of epochs Nepoch where sk = {sk,j,l }j=1,...,J; l=1,...,NT , sk,j,l = {ρ̂k,j , xk,j (tk,j,l ), tk,j,l }. 1: Initialize θ 2: for e = 1, . . . , Nepoch do 3: Compute the empirical DSM loss LS (θ) in (13) 4: Update θ using Adam 5: end for 6: return trained operator Tθ∗ ,S

Algorithm 2 Sampling on a New Task (i)

Input: empirical context ρ̂new = {xnew }n i=1 , trained operator Tθ ∗ ,S , number of reverse steps NT 1: Set ∆t = (1 − δ)/NT and define tm = δ + m∆t for m = 0, . . . , NT 2: Sample x̃tNT ∼ N (0, I) 3: for m = NT − 1, NT − 2, . . . , 0 do 4: Sample zm ∼ N (0, I) 5:

x̃tm ← x̃tm+1 +

1 β(tm+1 ) x̃tm+1 2

+ β(tm+1 )Tθ∗ ,S (ρ̂new )(x̃tm+1 , tm+1 ) ∆t p + β(tm+1 ) ∆t zm 6: end for 7: return generated sample x̃t0

Theoretical results

We now present the theoretical results for the proposed in-context score operator. We begin with the observation that the task-dependent population score admits a Gaussian-kernel conditional-mean representation. Replacing the unknown task distribution with its empirical context measure naturally yields a vector-valued Nadaraya–Watson estimator and, consequently, an empirical score estimator. This kernel representation observation motivates both our transformer architecture and its theoretical analysis: we show that the empirical estimator can be represented exactly by a model in our constructed transformer class, revealing a direct correspondence between Gaussian-kernel regression and the attention mechanism. Guided by this connection, we learn an in-context score operator through amortized denoising score matching across task distributions and analyze its generalization to unseen tasks drawn from the metadistribution. Section 4.1 develops the kernel representation and its exact transformer realization. Section 4.2 then presents and interprets the main generalization theorem, followed by Section 4.3, which outlines the corresponding error decomposition and the proof strategy. 4.1

Understanding the Target Score learning: From Kernel Regression to Transformers

In this subsection, we connect the task-dependent score function with a kernel regression estimator that can be represented by a transformer network. For a fixed task distribution ρ, recall that the perturbed law at time t admits density pt . The corresponding score function can be written as i 1 h x(t) − α(t) E(x | x(t)) . (14) T ∗ (ρ)(x(t), t) = S(x(t), t) = ∇x log pt (x) x=x(t) = − 2 σ (t) Here the conditional expectation is given by 2 x exp − ∥x(t)−α(t)x∥ dρ(x) 2 2σ (t) E(x | x(t)) = R . 2 dρ(x) exp − ∥x(t)−α(t)x∥ 2σ 2 (t) R

(15)

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

A P REPRINT

Given context samples ρ̂ = {x(1) , . . . , x(n) }, we approximate (15) by the empirical kernel regression estimator Pn ∥x(t)−α(t)x(i) ∥2 (i) x exp − 2 i=1 2σ (t) , (16) K(s) = P n ∥x(t)−α(t)x(i) ∥2 exp − i=1 2σ 2 (t) where s = {x(1) , . . . , x(n) , x(t), t}. We first observe that (15) admits a kernel representation and that its empirical discretization, given in (16), is a vector-valued Nadaraya–Watson kernel estimator. Each coordinate of this estimator is a scalar-valued kernel regressor and, by Lemma 1 of [40], can be approximated by a transformer network with bounded parameters. Applying this construction coordinatewise yields a transformer representation of the vector-valued estimator. This connection between kernel regression and transformer attention is the key observation motivating our use of transformers to approximate the task-dependent score operator. For completeness, we first define the transformer class considered in our analysis and then state the representation lemma for the empirical kernel estimator. Scalar Transformer Network Class. The scalar transformer network architecture that defined in [40] is as follows. (i) Given an input s = {x(1) , x(2) , · · · , x(n) ; x(t), t}, where xh is the h-th component of vector x(i) , for h = 1, · · · , D and i = 1, · · · , n. We define the output of the transformer network Fθ by  (1) x  0  .  . . H = PE + E(s) =    0 I 1 1

··· ··· .. . ··· ··· ···

x(n) ··· .. . ··· ··· ···

 x(t) 0  ..   .  ∈ Rdembed ×ℓ = R(2D+6)×(n+1) .  0  Iℓ  1

(17)

In matrix H, each column is a token, and each token has dimension dembed = 2D + 6 . The first D rows are data terms and the (D + 1)th row to (2D + 3)th row are 0. The (2D + 4)-th and (2D + 5)-th rows contain the well-known sinusoidal positional encodings ⊤ jπ jπ , sin , Ij = cos 2ℓ 2ℓ which determines how each token will interact with another through the attention mechanism. It is crucial to note that the data terms are dynamic, whereas the positional encoding and constant terms remain static. Furthermore, B1 , · · · , BLT : Rdembed ×ℓ → Rdembed ×ℓ are the transformer blocks which have a residual composition of the form B(θ, t; H) = FFN(MHAθ,t (H) + H) + MHAθ,t (H) + H.

(18)

where FFN is a feed-forward neural network Pmoperating tokenwise on the input, and the multi-head attention (MHA) with m heads is given by MHAθ,t (H) = j=1 Vθ,j (t)Hϕ((Kθ,j (t)H)⊤ Qθ,j (t)H) with ReLU activation from the first to the penultimate layers and Softmax activation for the last layer in the attention module (see definition 2 in [40]). The final attention mask restricts normalization to the n prompt tokens. The attention matrices use prescribed time modulation: for A ∈ {Q, K, V }, each matrix has the form (0)

(α)

Aθ,j (t) = Aθ,j + α(t)Aj

(σ)

+ σ(t)−1 Aj .

(19) (α)

(σ)

The block index is suppressed. Here α and σ are the prescribed diffusion schedules, and Aj , Aj

are fixed archi(0)

tectural templates, selected before training and shared by all networks in the class. Only the base matrices Aθ,j are trainable within these attention maps. A modulation term may be absent by setting its template to zero. The parameter vector θ contains the trainable base matrices and the trainable embedding, decoding, and feed-forward coefficients, all shared across time; the fixed templates are excluded from θ. Both the trainable coefficients and the fixed template entries are bounded in magnitude by κ. The output bound below is imposed uniformly over t ∈ [δ, 1]. Thus “fixed” refers to the schedules and templates, while the input time t still varies. Lastly, DE : Rdembed ×ℓ → R 8

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

A P REPRINT

is the decoding layer which outputs the desired element. We then denote a class of bounded scalar-value transformer networks as following, F(LT , mT , dembed , ℓ, LFFN , wFFN , R, κ) n = Fθ (·) Fθ (·) has the form (17) with LT transformer blocks, at most mT attention heads in each block, embedded dimension dembed , number of hidden tokens ℓ, and LFFN layers of feed-forward networks with hidden width wFFN , with output ∥Fθ (·)∥L∞ (RD ) ≤ R o and weight magnitude ∥θ∥∞ ≤ κ .

(20)

We now extend the above defintion to a vector-valued transformer network class. Vector-Valued Transformer Network Class. For D ∈ N1 , we define the vector-valued transformer network class by G(LT , mT , dembed , ℓ, LFFN , wFFN , R, κ, D), whose components belong to F(LT , mT , dembed , ℓ, LFFN , wFFN , R, κ). That is for Gθ ∈ G(LT , mT , dembed , ℓ, LFFN , wFFN , R, κ, D), we have Gθ (s) = (F1,θ (s), · · · , FD,θ (s)) where Fi,θ ∈ F(LT , mT , dembed , ℓ, LFFN , wFFN , R, κ) for i = 1, · · · , D. Using the coordinatewise construction motivated by Lemma 1 of [40], we obtain a transformer network Fθ ∈ F (LT , mT , dembed , ℓ, LFFN , wFFN , R, κ) that, on the noisy-query region specified below, satisfies Pn (i) ∥x(t)−α(t)x(i) ∥2 i=1 xj exp − 2σ 2 (t) . Fθ (s) = P (21) n ∥x(t)−α(t)x(i) ∥2 exp − 2 i=1 2σ (t) Note that the (21) represents the j-th component of the kernel regression estimator (16). By stacking D such networks componentwise, we can represent the full vector-valued kernel regression estimator in (16), which yields the following lemma indicate relation between kernel regression formula and vector-valued network class. Let n ≥ 1, R ≥ b ≥ 1, and let the fixed schedule satisfy Assumption 2. For C > 0, define XtC := {α(t)x + σ(t)ε : x ∈ M, ∥ε∥2 ≤ C}.

(22)

Lemma 1 (Transformer realization). There exists a network Gθ̂ ∈ G(LT , mT , dembed , ℓ, LFFN , wFFN , R, κ, D) with parameters LT = 5,

mT = n, dembed = 2D + 6, ℓ = n + 1, LFFN = O(1), wFFN = O(D), whose parameters are independent of the task, prompt, query, and time, such that Gθ̂ (s) = K(s) 4

(x(t) ∈ XtC ), 2

∥Gθ̂ (s)∥∞ ≤ b

(x(t) ∈ RD ).

(23)

(24)

A sufficient weight budget is κ = c0 D (n + 1) M , where  !2  √   1 2b D M = max 1, R, +C .   2 σ(δ)

(25)

One may take LFFN = 6 and wFFN = 2D + 6. It is worth emphasizing that the transformer parameters are independent of the task distribution. Rather than encoding a particular distribution, the network implements the generic computations underlying kernel regression: evaluating query–context similarities, normalizing the resulting kernel weights, and aggregating the context samples. All task-specific information therefore enters through the empirical context set, while the shared parameters encode a distribution-agnostic inference algorithm. Consequently, the same transformer can be applied uniformly to any task distribution within the considered class; adapting to a new task requires changing only the prompt, not the network parameters. This observation clarifies the operator-learning nature of the construction and provides a theoretical basis for generalization to previously unseen task distributions. We defer the constructive proof of Lemma 1 to Appendix B.1. 9

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

A P REPRINT

Vector-Valued Transformer Network Class for In-Context Score Matching. We now define the transformer network class used to parameterize the task-dependent score function. Recall that the score admits the representation i 1 h x(t) − α(t) E(x | x(t)) . S(x(t), t) = − 2 σ (t) Motivated by this identity, we let Gθ ∈ G (LT , mT , dembed , ℓ, LFFN , wFFN , R, κ, D) be a vector-valued transformer network that approximates the conditional expectation term. We then define the corresponding score network Tθ ∈ T (LT , mT , dembed , ℓ, LFFN , wFFN , R, κ, D) by i 1 h x(t) − α(t) Gθ x(1) , . . . , x(n) , x(t), t . Tθ x(1) , . . . , x(n) , x(t), t := − 2 σ (t)

(26)

For the analysis, take R ≥ b. The componentwise output bound of G implies ∥Gθ (s)∥∞ ≤ R.

sup

(27)

x(1:n) ∈Mn , x(t)∈RD , t∈[δ,1]

The bound applies to the mean network, not to the score in (26) on all of RD . The mean network uses the prescribed time modulation defined in (19). The fixed templates are those of the existing kernel construction; they are independent of its coordinate, localization constant, and clipping threshold. Thus the same construction belongs to this parameter class. Neither an additional trainable time network nor an additional transformer block is introduced. Consequently, the kernel conditional-mean estimator (16) can be represented exactly on XtC by a mean network in G. Applying (26) yields the corresponding empirical score estimator in T on the same region. This connection establishes a direct link between the target score operator, kernel regression, and the attention mechanism. 4.2

Generalization Error Analysis

This subsection analyzes the generalization performance of the empirical risk minimizer Tθ∗ ;S defined in (13). The training set S contains K task distributions, with J context–query episodes sampled from each task and NT diffusion samples generated for each episode. Consider an independent test task ρ ∼ µ, observed through context samples, Pn i.i.d. x(1) , . . . , x(n) ∼ ρ, ρ̂ = n1 i=1 δx(i) . At the task level, we compare the risk of the learned operator, evaluated using the empirical representation ρ̂, with that of the target distribution-level score operator T ∗ (ρ). The corresponding excess risk is Lnρ (Tθ∗ ;S ) − Lρ (T ∗ (ρ)) . Averaging over unseen tasks drawn from the meta-distribution yields the mean excess risk Ln (Tθ∗ ;S ) − L(T ∗ ) = Eρ∼µ Lnρ (Tθ∗ ;S ) − Lρ (T ∗ (ρ)) . This quantity measures the performance gap between the empirical risk minimizer and the target score operator on an unseen task drawn from the same meta-distribution and represented by only n context samples. Because Tθ∗ ;S depends on the randomly sampled training set S, we evaluate its performance by taking the expectation of the excess risk with respect to the randomness of S. Throughout this subsection, we therefore use the term generalization error to denote the expected excess risk ES [Ln (Tθ∗ ;S ) − L(T ∗ )] , where the outer expectation accounts for the randomness in the meta-training set. The accuracy of the empirical kernel estimator depends on how well the context samples cover the support of the task distribution ρ. In a general ambient space RD , achieving sufficient coverage may require an impractically large context set. Fortunately, many high-dimensional datasets possess substantially lower-dimensional intrinsic structure. Natural images provide a familiar example: although they are represented as high-dimensional pixel arrays, semantically meaningful images are often modeled as concentrating near a low-dimensional image manifold. Similar structure arises in scientific applications. For instance, in the Navier–Stokes problem considered in Section 5, task distributions induced by different viscosity parameters retain shared coherent flow structures. Although the probability law varies with viscosity, the resulting solution fields may remain concentrated near a common low-dimensional solution manifold. Motivated by these observations, we assume that all task distributions are supported on a common d-dimensional manifold M ⊂ RD , where d ≪ D. Under this assumption, the complexity of context-based score estimation is governed primarily by the intrinsic dimension d, rather than by the potentially much larger ambient dimension D. For the finite-sample analysis, a uniform loss bound follows from the bounded mean-network output and the positive time cutoff. The precise assumptions are stated below. 10

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

A P REPRINT

Assumption 1 (Regularity of the task distributions). We assume the following: 1. There exists a compact d-dimensional Riemannian manifold M ⊂ RD such that every task distribution ρ under consideration is supported on M. Moreover, M is isometrically embedded in RD , satisfies M ⊂ [−b, b]D for some b > 0, and has positive reach τM > 0. 2. Each task distribution ρ is quasi-uniform on M. That is, if ρ̄ denotes the uniform probability measure on M, then there exist constants Cρ,1 , Cρ,2 > 0 such that Cρ,1 ρ̄(B) ≤ ρ(B) ≤ Cρ,2 ρ̄(B) for every measurable set B ⊆ M. We take the support bound b ≥ 1, enlarging it if necessary. Let N (M, r) ≤ AM r−d ,

AM := sup rd N (M, r) < ∞,

0 < r ≤ 1,

(28)

0<r≤1

where N (M, r) is the Euclidean covering number with centers on M. For a compact smooth embedded manifold, this bound follows by covering finitely many bounded coordinate charts, whose parametrizations are Lipschitz, by ddimensional grids. The constant AM retains the dependence on the geometry of M. The quasi-uniformity condition in Assumption 1 is retained, although the approximation argument below does not use it. Assumption 2 (Gaussian perturbations and diffusion schedule). For t ∈ [δ, 1], the known measurable schedule satisfies 0 ≤ α(t) ≤ 1 and 0 < σ(δ) ≤ σ(t) ≤ 1. The noisy query is x(t) = α(t)x + σ(t)ε,

ε ∼ N (0, ID ),

where ε is independent of the clean query and the prompt. The time variable is independent and uniform on [δ, 1]. All population and empirical losses below use this same Gaussian perturbation law. A finite radius introduced in the proof restricts the region of exact network realization, rather than the distribution of the noise. The denoising score-matching loss is bounded under the full Gaussian law. Indeed, Tθ (ρ̂)(x(t), t) +

ε α(t) = Gθ (s) − x , 2 σ(t) σ(t)

(29)

so that, uniformly over the hypothesis class and all noise realizations, 2

0 ≤ Tθ (ρ̂)(x(t), t) +

ε ≤ RL , σ(t) 2

RL =

D(R + b)2 . σ(δ)4

(30)

This bound uses the output restriction already present in G. Its dependence on δ must be retained when the time cutoff is subsequently chosen as a function of n. Under Assumption 1 and 2, we obtain the following generalization error bound. Theorem 1 (Generalization error). Under Assumption 1 and Assumption 2, let S be the training set defined above. Let n ≥ 2 and choose δ ∈ (0, 1) satisfying σ(δ) = n−1/(d+6) . Consider the class T (LT , mT , dembed , ℓ, LFFN , wFFN , R, κ, D) with parameters LT = 5, LFFN = O(1),

mT = n,

dembed = O(D),

wFFN = O(D),

ℓ = n + 1, 2d+20 κ = O D8 n d+6 b8 R4 ,

Take the concrete choice dembed = wFFN = 2D + 6 and LFFN = 6, with a sufficiently large absolute constant in the weight budget to contain the candidate from Proposition 1. Let Tθ∗ ;S be a measurable empirical minimizer of LS over this class. Then ES [Ln (Tθ∗ ;S ) − L(T ∗ )] 4 p 2 1 1 1 5/2 12 + d+6 2 log(n DKJNT ) √ + √ +√ + C2 Dn− d+6 . ≤ C1 D n KJNT K KJ

(31)

Here C1 depends only on b, R and the fixed depths, and C2 depends only on b, d, AM , R. Both are independent of n, D, K, J, NT and δ. The cutoff choice is required to be attained by the diffusion schedule. The theorem yields the following implications: 11

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

A P REPRINT

1. Context efficiency is governed by intrinsic dimension. The finite-context error decays as n−2/(d+6) . As the intrinsic dimension d increases, the exponent becomes smaller and the convergence rate deteriorates, reflecting a curse of intrinsic dimensionality. The ambient dimension D enters the bound only through polynomial and logarithmic factors. Thus, high-dimensional observations need not be prohibitive when the task distributions concentrate on a low-dimensional manifold. 2. The statistical error separates across sampling levels. The terms K −1/2 , (KJ)−1/2 , and (KJNT )−1/2 quantify the effects of finite sampling at the task, episode, and diffusion-sample levels, respectively. Increasing J reduces the episode-level and diffusion-level errors, while increasing NT reduces only the diffusion-level error. 3. Task diversity is essential for generalization. The persistent K −1/2 term is unaffected by increases in either J or NT . Thus, collecting additional episodes or diffusion samples from existing tasks cannot compensate for limited meta-training task coverage. Generalization to unseen distributions instead requires a sufficiently large and diverse collection of meta-training tasks. From a data-allocation perspective, the theorem therefore suggests prioritizing broader task coverage once the within-task estimation error has been sufficiently reduced. 4.3

Proof Outline

This subsection outlines the proof of Theorem 1. Let Tθ̂ denote the transformer realization of the empirical kernel estimator established in Lemma 1. By adding and subtracting LS (Tθ∗ ;S ), LS (Tθ̂ ), and Ln (Tθ̂ ), we obtain Ln (Tθ∗ ;S ) − L(T ∗ ) = Ln (Tθ∗ ;S ) − LS (Tθ∗ ;S ) + LS (Tθ∗ ;S ) − LS (Tθ̂ ) + LS (Tθ̂ ) − Ln (Tθ̂ ) + Ln (Tθ̂ ) − L(T ∗ ) . Since LS (Tθ∗ ;S ) − LS (Tθ̂ ) ≤ 0, we then have ES [Ln (Tθ∗ ;S ) − L(T ∗ )] ≤ ES [Ln (Tθ∗ ;S ) − LS (Tθ∗ ;S )] + ES LS (Tθ̂ ) − Ln (Tθ̂ ) + Ln (Tθ̂ ) − L(T ∗ ) . {z } | | {z } | {z } statistical error I

statistical error II

approximation error III

(32) In what follows, we control the three terms in (32) separately. We first bound the approximation error by showing that the target score operator T ∗ admits a kernel regression representation and can therefore be approximated by a transformer-based neural operator. We then bound Statistical Error I using covering numbers and Rademacher complexity, followed by Statistical Error II via Hoeffding’s inequality. Combining these estimates establishes the desired generalization bound. Approximation error term III. By the exact transformer realization in Section 4.1, this term reduces to the approximation error of the empirical Nadaraya–Watson estimator. Concentration of the kernel estimator on the underlying d-dimensional manifold gives the following result, whose proof is deferred to Appendix B.4. Proposition 1 (Approximation error). Under Assumption 1 and Assumption 2, T (LT , mT , dembed , ℓ, LFFN , wFFN , R, κ, D) with parameters LT = 5,

mT = n,

dembed = O(D),

LFFN = O(1),

wFFN = O(D),

κ=O

ℓ = n + 1, D 8 n 2 b8 R 4 , σ(δ)8

there exists Tθ̂

∈

(33)

such that Ln (Tθ̂ ) − L(T ∗ ) ≤

Cb,d AM D + 4Db2 σ(δ)2 . nσ(δ)d+4

(34)

The constants hidden in the network budgets are absolute, with R ≥ b ≥ 1; the geometric dependence is explicit in AM . Statistical Error I. Term I measures the deviation between the finite-context population risk and the empirical risk of the learned estimator. Since Tθ∗ ;S depends on the training set, controlling this term requires a uniform generalization analysis over the transformer hypothesis class. Using covering number and Rademacher complexity estimates, we obtain the following bound. 12

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

A P REPRINT

Lemma 2 (Statistic error term I). Use the width convention of Lemma 4. Let Tθ∗ ;S ∈ T be a measurable empirical minimizer of LS (Tθ ). Then ES [Ln (Tθ∗ ;S ) − LS (Tθ∗ ;S )]   s √ 1 1 1 A KJN T 6 DPT log 0,δ +√ + 1  + R L qδ , ≤ 4RL √ + √ RL KJNT K KJ

(35)

where RL is given in (30), qδ in (43), and PT , A0,δ in (47), (50). Furthermore, with parameters LT = 5,

mT = n,

dembed = O(D),

LFFN = O(1),

wFFN = O(D),

κ=O

ℓ = n + 1, D 8 n 2 b8 R 4 , σ(δ)8

we have ES [Ln (Tθ∗ ;S ) − LS (Tθ∗ ;S )] s ! D(R + b)2 nDKJNT 1 1 1 3 √ +√ ≤O nD log +√ + D(R + b)2 σ(δ)2 . σ(δ)4 σ(δ) KJNT K KJ

(36)

Here n ≥ 2, and the notation O(·) hides constants depending only on b, R and the fixed depths. In particular, the dependence of RL on δ is displayed explicitly. The proof of Lemma 2 is provided in Appendix B.2. Statistical error II. Term II measures the deviation between the empirical and finite-context population risks evaluated at Tθ̂ . Since Tθ̂ is fixed independently of the training set, the deviations arising from the task, episode, and diffusion-sample levels can be controlled separately using concentration inequalities. This yields the following bound, whose proof is deferred to Appendix B.3. Lemma 3 (Statistic error term II). . The Statistic error term II can be bounded as p 1 1 1 ES LS (Tθ̂ ) − Ln (Tθ̂ ) ≤ O log(1/σ(δ)) √ + √ +√ + σ 2 (δ) . (37) KJNT K KJ Combining Lemmas 2 and 3 with Proposition 1 in the decomposition (32) yields an upper bound consisting of the three hierarchical statistical error terms and the finite-context approximation error term. Choosing σ(δ) = n−1/(d+6) balances the dominant approximation error terms and yields the rate n−2/(d+6) . Substituting the transformer parameter scaling and the loss bound RL = D(R + b)2 σ(δ)−4 establishes the bound in Theorem 1. The detailed derivation is postponed to Appendix B.5.

Numerical results

In this section, we evaluate our method on four benchmark datasets spanning scientific fields and visual data: solution fields of the 2D Navier–Stokes equation, ERA5 climate data, QuickDraw sketches, and Food-101 images. For Navier– Stokes, we examine both in-distribution and out-of-distribution task generalization. For the remaining datasets, we evaluate in-distribution generation on observed task categories. Data preparation and mini-batch construction Algorithm 1 is presented using an explicit training set S = {sk }K k=1 , where for each task ρk one pre-generates J independent context–query episodes and, for each episode, NT independent draws of diffusion time and noise, yielding tuples sk,j,l = {ρ̂k,j , xk,j (tk,j,l ), tk,j,l } that are stored in full before training begins. In practice, however, pre-generating and storing the full dataset S is prohibitively memoryintensive: since each episode ρ̂k,j requires n fresh samples from ρk , the total number of samples grows as K × J × n, which quickly becomes infeasible for large J. We therefore construct mini-batches on the fly. Specifically, for the i.i.d. NS experiments, we pre-sample K task distributions {ρk }K k=1 ∼ µ and, for each task ρk , draw a fixed bank of B (i)

i.i.d.

i.i.d. samples {xk }B i=1 ∼ ρk ; for fixed labeled datasets (ERA5, QuickDraw, and Food-101), we treat each class as a task, select K classes, and randomly sample B datapoints per class. At each training step, we sample a task index 13

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

A P REPRINT

(i)

k ∼ U{1, . . . , K}, randomly draw n + 1 samples from the fixed bank {xk }B i=1 , and split them into a context set (1) (n) (n+1) qry ρ̂k := {xk , . . . , xk } and a clean query xk := xk ; we then sample t ∼ Unif[δ, 1] and ε ∼ N (0, I) independently on the fly to form the noised query xk (t) = α(t)xqry k + σ(t)ε, yielding the training episode s = {ρ̂k , xk (t), t} used in (13). Metrics. To quantify the performance of the in-context diffusion model and the baseline methods, we compare the generated and reference empirical distributions using the Wasserstein-2 distance, the Jensen–Shannon divergence (JSD) between scalar-valued empirical probability density functions, energy-spectrum-based metrics, and Fréchet Inception Distance (FID). Further details are provided in Appendix D.3. Baseline models. We compare three methods: our proposed in-context learning diffusion model (ICL DM), a conditional diffusion model (Cond DM), and the closely related baseline SNO [27]. All three models share the same time-conditioned, multi-resolution U-Net denoising backbone and training protocol, and differ only in how the prompt set is represented and injected. In ICL DM, the prompt set is encoded into multi-scale features that interact with the noisy query through cross-attention. In Cond DM, the prompt set is compressed into a single learned embedding that modulates the U-Net through FiLM-style conditioning. In SNO, this learned prompt representation is replaced by a projected kernel mean embedding (KME), while the rest of the denoising architecture remains the same as in Cond DM. Additional implementation details are given in Appendix D. 5.1

2D Navier–Stokes equation

We consider the vorticity form of the 2D Navier–Stokes equation with Kolmogorov forcing:  2 2  ∂t ω(t, x) + u(t, x) · ∇ω(t, x) − ν∇ ω(t, x) = f (x), x ∈ [0, 2π] , t ∈ [0, T ], ∇ · u(t, x) = 0, x ∈ [0, 2π]2 , t ∈ [0, T ],   ω(0, x) = ω0 (x), x ∈ [0, 2π]2 .

(38)

Here x = (x1 , x2 ), ω(t, x) denotes the vorticity, and the velocity field is given by u(t, x) = ∂x2 ψ(t, x), −∂x1 ψ(t, x) , ω(t, x) = −∇2 ψ(t, x). The Kolmogorov forcing is defined as f (x) = sin(k0 x1 ). In this example, we set the wavenumber k0 = 4. To generate initial conditions, we first draw samples from a log-normal distribution on a uniform 256 × 256 mesh over the spatial domain x ∈ [0, 2π]2 . We parameterize Navier-Stokes tasks by ξ ∈ Ξ, here the ξ := ν is the viscosity. Let π0 be a probability measure over initial conditions ω0 , and for each ξ ∈ Ξ, we let S T,ξ denote the solution operator mapping ω0 to the vorticity field at time T , i.e., S T,ξ : ω0 (x) 7→ ω(T, x; ω0 , ξ); this induces the task distribution ρξ := (S T,ξ )# π0 ∈ P(X ). Let ν be a probability distribution on Ξ, and define the mapping F : Ξ → P(X ) by F (ξ) = ρξ . The meta-distribution over task distributions is then the pushforward µ := F# ν ∈ P(P(X )). In the following experiments, we set T = 1 and use the finite volume method (FVM) implemented in JAX-CFD [22] for generating data. To evaluate both in-distribution and out-of-distribution generalization, we construct 100 tasks by uniformly sampling ν from the union of two intervals, ν ∼ Unif ([1.5e − 3, 1.2e − 2] ∪ [1.5e − 2, 5e − 2]) .

(39)

For the in-distribution setting, we randomly sample five values of ν from the first interval and five values from the second interval, and evaluate the model on these 10 values. The results of the different methods are reported in Table 1. We also present three representative examples with ν = 1.5e − 3, ν = 2.2e − 3, and ν = 3e − 2, as shown in Figures 1 and 2. For the out-of-distribution setting, we evaluate the model at ν1 = 1.2e−3, ν2 = 1.3e−2, and ν3 = 6.4e−2, all of which lie outside the training intervals. The corresponding results are presented in Figures 3 and 4. To assess sample quality, we compare the generated and reference distributions using the weighted mean energy log ratio (MELRw) and the JSD between empirical scalar-valued PDFs. The results are reported in Table 1. As shown in Table 1, our method outperforms the baselines in both in-distribution and out-of-distribution settings, except for MELRw in the OOD(high) setting, where Cond DM achieves a lower value. 14

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

A P REPRINT

(a) ν = 1.5 × 10−3

(b) ν = 2.2 × 10−3

Figure 1: Representative in-distribution inferencesamples for the Navier–Stokes equation. Panels (a)–(c) correspond to viscosities ν = 1.5 × 10−3 , 2.2 × 10−3 , and 3 × 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.

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

(a) ν = 1.5 × 10−3

(b) ν = 2.2 × 10−3

A P REPRINT

Figure 2: Scalar-valued empirical PDF of the Navier–Stokes equation in the in-distribution setting. Panels (a)–(c) correspond to ν = 1.5 × 10−3 , 2.2 × 10−3 , and 3 × 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. IND OOD(low) OOD(mid) OOD(high) MELRw ↓ JSD ↓ MELRw ↓ JSD ↓ MELRw ↓ JSD ↓ MELRw ↓ JSD ↓ 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-ofdistribution (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 ν = 1.2 × 10−3 , 1.3 × 10−2 , and 6.4 × 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.

5.2

ERA5 dataset

ERA5 [12] provides hourly estimates of a wide range of atmospheric, land-surface, and ocean-wave variables on a global grid, with a horizontal resolution of approximately 31 km for the main product and 137 vertical levels. In this study, we prepare the dataset as follows. First, we select nine variables: 0 = temperature, 1 = u-wind, 2 = v-wind, 3 = geopotential, 4 = vertical velocity, 5 = specific humidity, 6 = relative humidity, 7 = divergence, and 8 = vorticity. Second, we select 16 pressure levels: 150, 175, 200, 225, 250, 300, 350, 400, 450, 500, 550, 600, 650, 700, 750, and 800 hPa. For each variable–pressure-level pair, we record data at 8 hour intervals during 2023 and 2024. The resulting pairs define 144 candidate task classes, from which we select K = 100 tasks for the experiments. For each selected task, we randomly sample B = 1000 data points. During training and testing, the number of prompts is set to n = 100 in both cases. Figures 5, 6 and 7 present a comparison of the data generated by the different methods. From the visual comparison, we observe that our method consistently produces samples that shares the similar pattern to the reference samples compared to the baselines. To quantitatively assess generation quality, we compute the Wasserstein-2 distance (W2 ), the weighted and unweighted mean energy log ratios (MELRw and MELRu ), and the JSD between empirical scalar-valued PDFs. The results are reported in Table 2, where our method outperforms the baselines on most of the evaluated metrics, demonstrating its strong ability to capture the underlying data distribution and generate high-quality samples that closely match the reference distribution. 5.3

Google sketch dataset

Google’s Quick, Draw! (QuickDraw) sketch dataset [20] is a large-scale collection of human doodles obtained through the online game “Quick, Draw!,” in which players are asked to draw an object within a short time. The dataset contains approximately 50 million drawings spanning 345 categories. In this paper, we use the first 100 categories and randomly sample 2000 sketches from each category. During both training and testing, we set the number of prompts to n = 200. Figure 8 compares the samples generated by different methods for the case K = 5. From the visual comparison, we observe that our method consistently produces samples that more closely resemble the reference samples than 16

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

A P REPRINT

(a) ν = 1.2 × 10−3

(b) ν = 1.3 × 10−2

Figure 3: Out-of-distribution inference samples for Navier–Stokes equation. Panels (a)–(c) correspond to ν = 1.2 × 10−3 , 1.3 × 10−2 , and 6.4 × 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.

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

(a) ν = 1.2 × 10−3

(b) ν = 1.3 × 10−2

A P REPRINT

Figure 4: Scalar-valued empirical PDFs for the Navier–Stokes equation in the out-of-distribution setting. Panels (a)– (c) correspond to ν = 1.2 × 10−3 , 1.3 × 10−2 , and 6.4 × 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.

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 |log(Ek /Ekref )| 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.

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

Idx

Method

W2 ↓

MELRw ↓

MELRu ↓

JSD ↓

ICL DM Cond DM SNO

2.68e1 4.39e1 3.94e1

2.43e-1 1.38e-1 2.64e-1

2.32e-1 3.18e-1 1.85e-1

1.33e-2 1.30e-2 1.40e-2

ICL DM Cond DM SNO

1.85e1 3.58e1 2.60e1

2.27e-2 2.36e-1 3.10e-2

2.87e-2 3.01e-1 3.33e-1

1.54e-3 2.70e-2 3.37e-3

ICL DM Cond DM SNO

2.56e1 3.04e1 3.06e1

1.42e-1 1.39e-1 2.21e-1

1.20e-1 1.91e-1 2.61e-1

4.59e-3 2.08e-2 2.07e-2

ICL DM Cond DM SNO

2.59e1 3.32e1 3.01e1

1.13e-1 2.00e-1 1.35e-1

2.03e-1 9.53e-2 1.20e-1

4.93e-3 3.20e-2 1.52e-2

ICL DM Cond DM SNO

2.32e1 2.96e1 2.92e1

1.81e-2 3.45e-1 2.40e-1

1.37e-1 1.14e-1 2.58e-1

1.25e-3 6.20e-2 4.37e-2

ICL DM Cond DM SNO

1.51e1 3.55e1 3.03e1

2.98e-2 1.52e-1 1.61e-1

1.38e-1 7.47e-1 8.28e-1

4.55e-3 8.42e-3 3.47e-2

ICL DM Cond DM SNO

1.54e1 2.26e1 2.52e1

5.25e-2 1.86e-2 3.72e-2

1.27e-1 7.04e-1 1.22e0

4.45e-3 8.89e-2 1.09e-1

ICL DM Cond DM SNO

1.25e1 3.77e1 2.33e1

1.33e-1 8.06e-1 1.66e-1

3.78e-1 1.72e0 1.18e0

5.60e-2 2.24e-1 8.55e-2

ICL DM Cond DM SNO

2.46e1 3.84e1 3.16e1

7.60e-2 9.16e-1 6.55e-1

6.87e-2 3.18e-1 2.42e-1

3.06e-3 1.06e-1 5.48e-2

ICL DM Cond DM SNO

1.91e1 3.79e1 3.04e1

2.17e-1 8.63e-1 2.96e-1

1.37e-1 2.29e-1 1.64e-1

7.29e-3 8.84e-2 4.86e-3

A P REPRINT

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 (W2 ), weighted and unweighted mean energy log ratios (MELRw and MELRu ), 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.

those generated by the baselines. For quantitative evaluation, we compute the FID, FID-QD, and JSDPDF metrics to measure the discrepancy between the generated and reference distributions. The results in Table 3 for K = 5, 10, and 100 show that our method outperforms the baselines across most metrics and settings. These results indicate that our method better captures the underlying data distribution and generates higher-quality samples that more closely match the reference distribution. We also observe an overall degradation in performance as K increases. Here, increasing K introduces more distinct categorical task modes into the meta-distribution, increasing task heterogeneity and making generalization more challenging. This differs from the theoretical K −1/2 scaling, which considers increasing the number of tasks under a fixed meta-distribution. Nevertheless, our method maintains a clear performance advantage over the baselines even for larger K, further demonstrating its effectiveness in complex multi-class generation tasks. 19

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

A P REPRINT

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 scalarvalued 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 |log(Ek /Ekref )| 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.

Eiffel Tower Mona Lisa Airplane FID ↓ FID-QD ↓ JSD ↓ FID ↓ FID-QD ↓ JSD ↓ FID ↓ FID-QD ↓ JSD ↓ 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 SNO (K = 10) 333.72 138.12 3.19e-3 182.83 135.98 9.06e-3 215.78 129.79 1.36e-3 ICL DM (K = 100) 217.64 70.31 2.08e-3 142.77 93.19 4.66e-3 143.42 104.52 4.28e-3 Cond DM (K = 100) 316.93 116.63 1.54e-3 171.23 148.58 1.41e-2 197.50 128.99 3.38e-3 SNO (K = 100) 320.16 132.71 4.91e-3 176.47 144.63 1.16e-2 181.37 134.70 3.18e-3 Table 3: In-distribution quantitative comparison on three representative QuickDraw categories as the number of metatraining tasks varies over K ∈ {5, 10, 100}. ICL DM, Cond DM, and SNO are evaluated using the standard Fré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.

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

A P REPRINT

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 scalarvalued 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 |log(Ek /Ekref )| 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.

5.4

Food-101 Dataset

Food-101 [3] is a large-scale food image dataset consisting of 101 food categories, with 1,000 images in each category. The images exhibit substantial variation in viewpoint, scale, illumination, occlusion, and plating style. In this setting, the meta-distribution can be viewed as the overall distribution over food images, while each task distribution corresponds to a specific food category. Following the same procedure as in the QuickDraw experiments, we consider K = 5, 10, and 100 categories and perform the same quantitative evaluation. Note that all 1,000 images in each selected category are used to construct the empirical task distribution from which both training samples and inference prompts and reference samples are drawn. Thus, this experiment evaluates in-distribution generation only and does not assess out-of-distribution generalization. As shown in Table 4, our method outperforms the baselines across most metrics and settings. This suggests that our method more effectively captures the class-conditioned data distribution and generates samples that are closer to the reference distribution. As K increases, the performance of all methods generally degrades. As in the QuickDraw experiment, this reflects the increased complexity of the categorical task family under a fixed per-task data budget, prompt size, and shared model capacity, rather than the K −1/2 scaling under a fixed meta-distribution. Even in this more challenging regime, however, our method retains a substantial advantage over the baselines. These results should not be interpreted as an inherent limitation of the proposed method to categorical task distributions. Rather, they reflect the resource-limited regime considered here, where the per-task data and model capacity are fixed relative to an increasingly broad task family. Whether additional data and model capacity can mitigate this degradation is left for future work. 21

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

A P REPRINT

(a) Eiffel Tower.

(b) Mona Lisa.

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.

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

A P REPRINT

(a) Apple Pie.

(b) Baby Back Ribs.

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.

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

A P REPRINT

Apple Pie Ribs Beef Carpaccio FID ↓ FID-Food ↓ JSD ↓ FID ↓ FID-Food ↓ JSD ↓ FID ↓ FID-Food ↓ JSD ↓ 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 SNO (K = 10) 241.85 72.62 2.13e-2 259.45 83.33 1.10e-1 247.94 131.94 6.52e-2 ICL DM (K = 100) 197.98 55.60 1.98e-2 218.12 60.53 2.81e-2 192.40 102.73 1.75e-2 Cond DM (K = 100) 221.63 69.52 2.64e-2 255.99 87.90 5.79e-2 238.23 119.79 2.08e-2 SNO (K = 100) 235.84 70.97 2.44e-2 250.45 98.45 5.97e-2 250.09 127.23 1.97e-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 ∈ {5, 10, 100}. Performance is evaluated using the standard Fré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.

Conclusion

We introduced an in-context solution operator learning framework for diffusion models. In our formulation, each target distribution is treated as a task drawn from a meta-distribution and is specified at inference time by a finite set of i.i.d. context samples. The proposed in-context score operator acts directly on this empirical measure, constructing taskdependent score estimates through interactions among the context samples, the noisy query point, and the diffusion time. By training the operator through amortized denoising score matching across task distributions, reverse-diffusion sampling can be performed on previously unseen tasks without retraining, fine-tuning, or task-specific parameter updates. A central observation underlying our method is that the task-dependent score admits a Gaussian-kernel conditionalmean representation. Its empirical approximation leads to a Nadaraya–Watson estimator that can be represented exactly by a model in our constructed transformer class. This connection provides a principled explanation for the proposed architecture: the transformer parameters encode a distribution-independent kernel-regression algorithm, while all task-specific information is supplied through the context set. We further established generalization-error bounds that separate the effects of the context size, the number of meta-training tasks, the number of episodes per task, the number of diffusion samples per episode, model complexity, and data geometry. The finite-context rate is governed by the intrinsic dimension of the task manifold, whereas the ambient dimension enters only through polynomial and logarithmic factors. Moreover, the persistent task-level error term shows that additional within-task sampling cannot compensate for insufficient meta-training task coverage, highlighting the importance of training on a sufficiently large and diverse family of distributions. Our empirical results further demonstrate the effectiveness of the proposed framework. For task families of lower complexity, such as Navier-Stokes equation, the method achieves both strong in-distribution performance and modest out-of-distribution generalization. For task families of intermediate complexity, such as ERA5, it remains effective in the in-distribution setting. By contrast, for more complex task families, such as QuickDraw and Food-101, generalization becomes substantially more challenging, and even strong in-distribution performance can be difficult to achieve when both the number of in-context prompts and the amount of training data are limited. Importantly, these findings should not be interpreted as evidence that the method fails on more complex task families. Rather, they indicate that generalization in these settings are intrinsically more challenging under limited context and finite training data.

Acknowledgments

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

A P REPRINT

References [1] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Advances in Neural Information Processing Systems, volume 36, pages 57125–57211, 2023. [2] Arpit Bansal, Hsiao-Ming Chu, Avi Schwarzschild, Sravanti Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 843–852, 2023. [3] Lukas Bossard, Matthieu Guillaumin, and Luc Van Gool. Food-101 – mining discriminative components with random forests. In European Conference on Computer Vision, 2014. [4] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems (NeurIPS), volume 33, pages 1877–1901, 2020. [5] Tianqi Chen, Yongfei Liu, Zhendong Wang, Jianbo Yuan, Quanzeng You, Hongxia Yang, and Mingyuan Zhou. Improving in-context learning in diffusion models with visual context-modulated prompts. arXiv preprint arXiv:2312.01408, 2023. [6] Frank Cole, Yulong Lu, Wuzhe Xu, and Tianhao Zhang. In-context learning of linear systems: Generalization theory and applications to operator learning. arXiv preprint arXiv:2409.12293, 2024. [7] Frank Cole, Dixi Wang, Yineng Chen, Yulong Lu, and Rongjie Lai. In-context operator learning on the space of probability measures. arXiv preprint arXiv:2601.09979, 2026. [8] Prafulla Dhariwal and Alex Nichol. Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems, 2021. [9] Xue Feng, Li Wang, Deanna Needell, and Rongjie Lai. Learn to evolve: Self-supervised neural jko operator for wasserstein gradient flow. arXiv preprint arXiv:2601.05583, 2026. [10] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. arXiv preprint arXiv:2208.01066, 2022. [11] Giorgio Giannone, Didrik Nielsen, and Ole Winther. arXiv:2205.15463, 2022.

Few-shot diffusion models.

arXiv preprint

[12] Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquı́n Muñoz-Sabater, Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, et al. The era5 global reanalysis. Quarterly journal of the royal meteorological society, 146(730):1999–2049, 2020. [13] Jonathan Ho, Ajay N. Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020. [14] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022. [15] Yankun Hong, Harshit Bansal, and Karen Veroy. A score-based generative solver for pde-constrained inverse problems with complex priors. arXiv preprint arXiv:2408.16626, 2024. [16] Han Huang and Rongjie Lai. Unsupervised solution operator learning for mean-field games. Journal of Computational Physics, 537:114057, 2025. [17] Jiahe Huang, Guandao Yang, Zichen Wang, and Jeong Joon Park. Diffusionpde: Generative pde-solving under partial observation. arXiv preprint arXiv:2406.17763, 2024. [18] Rongjie Huang, Zhou Zhao, Huadai Liu, Jinglin Liu, Chenye Cui, and Yi Ren. ProDiff: Progressive fast diffusion model for high-quality text-to-speech. In Proceedings of the 30th ACM International Conference on Multimedia, pages 2595–2605, 2022. [19] Minki Jeong, Hyeon Cho Kim, Seung Jae Cheon, Byeong-Ju Choi, and Nam Soo Kim. Diff-TTS: A denoising diffusion model for text-to-speech. arXiv preprint arXiv:2104.01409, 2021. [20] Jonas Jongejan, Henry Rowley, Takashi Kawashima, Justin Kim, and Nick Fox-Gieg. The Quick, Draw! Dataset, 2017. [21] Gwanghyun Kim, Taeseong Kwon, and Jong Chul Ye. DiffusionCLIP: Text-guided diffusion models for robust image manipulation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2426–2435, 2022. 25

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

A P REPRINT

[22] Dmitrii Kochkov, Jamie A. Smith, Ayya Alieva, Qing Wang, Michael P. Brenner, and Stephan Hoyer. Machine learning–accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 118(21), 2021. [23] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. DiffWave: A versatile diffusion model for audio synthesis. arXiv preprint arXiv:2009.09761, 2020. [24] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces with applications to pdes. Journal of Machine Learning Research, 24(89):1–97, 2023. [25] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 19565–19594. PMLR, 2023. [26] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew M. Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. In International Conference on Learning Representations, 2021. [27] Xinyu Liao, Aoyang Qin, Jacob Seidman, Junqi Wang, Wei Wang, and Paris Perdikaris. Score neural operator: A generative model for learning and generalizing across multiple probability distributions. arXiv preprint arXiv:2410.08549, 2024. [28] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218–229, 2021. [29] Yulong Lu and Wuzhe Xu. Generative downscaling of pde solvers with physics-guided diffusion models. Journal of scientific computing, 101(3):71, 2024. [30] Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. Metaicl: Learning to learn in context. In Proceedings of the 2022 conference of the North American chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2791–2809, 2022. [31] Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 11048–11064, 2022. [32] Gautam Mittal, Jesse Engel, Curtis Hawthorne, and Ian Simon. Symbolic music generation with diffusion models. arXiv preprint arXiv:2103.16091, 2021. [33] Ivona Najdenkoska, Animesh Sinha, Abhimanyu Dubey, Dhruv Mahajan, Vignesh Ramanathan, and Filip Radenovic. Context diffusion: In-context aware image generation. arXiv preprint arXiv:2312.03584, 2023. [34] Weili Nie, Brandon Guo, Yujia Huang, Chaowei Xiao, Arash Vahdat, and Anima Anandkumar. Diffusion models for adversarial purification. arXiv preprint arXiv:2205.07460, 2022. [35] Tim Pearce, Tabish Rashid, Antti Kanervisto, David Bignell, Mingfei Sun, Raluca Georgescu, Sergio V. Macua, Shun Zhang Tan, Ida Momennejad, Katja Hofmann, and Satinder Devlin. Imitating human behaviour with diffusion models. arXiv preprint arXiv:2301.10677, 2023. [36] Ryan Po, Wang Yifan, Volker Golyanik, Kfir Aberman, Jonathan T. Barron, Amit H. Bermano, Eric R. Chan, Tali Dekel, Aleksander Holynski, Angjoo Kanazawa, et al. State of the art on diffusion models for visual computing. arXiv preprint arXiv:2310.07204, 2023. [37] Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. In Advances in Neural Information Processing Systems, volume 36, pages 14228–14246, 2023. [38] Moritz Reuss, Ming Li, Xiaogang Jia, and Rudolf Lioutikov. Goal-conditioned imitation learning using scorebased diffusion policies. arXiv preprint arXiv:2304.02532, 2023. [39] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjorn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022. [40] Zhaiming Shen, Alexander Hsu, Rongjie Lai, and Wenjing Liao. Understanding in-context learning on structured manifolds: Bridging attention to kernel methods. In International Conference on Learning Representations, volume 2026, pages 42067–42103, 2026. 26

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

A P REPRINT

[41] Naufal Shidqi, Chaeyoon Jeong, Sungwon Park, Elke Zeller, Arjun Babu Nellikkattil, and Karandeep Singh. Generating high-resolution regional precipitation using conditional diffusion model. arXiv preprint arXiv:2312.07112, 2023. [42] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, volume 32, 2019. [43] Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Scorebased generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020. [44] Prakhar Srivastava, Ruihan Yang, Gavin Kerrigan, Gideon Dresdner, Jeremy McGibbon, Christopher Bretherton, and Stephan Mandt. Precipitation downscaling with spatiotemporal video diffusion. arXiv preprint arXiv:2312.06071, 2023. [45] Jingmin Sun, Yuxuan Liu, Zecheng Zhang, and Hayden Schaeffer. Towards a foundation model for partial differential equations: Multioperator learning and extrapolation. Physical Review E, 111(3):035304, 2025. [46] Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, 2023. [47] Zhendong Wang, Yifan Jiang, Yadong Lu, Yelong Shen, Pengcheng He, Weizhu Chen, Zhangyang Wang, and Mingyuan Zhou. In-context learning unlocked for diffusion models. In Advances in Neural Information Processing Systems, 2023. [48] Quanlin Wu, Hang Ye, and Yuntian Gu. Guided diffusion model for adversarial purification from random noise. arXiv preprint arXiv:2206.10875, 2022. [49] Chaowei Xiao, Zhongzhu Chen, Kun Jin, Jiongxiao Wang, Weili Nie, Mingyan Liu, Anima Anandkumar, Bo Li, and Dawn Song. Densepure: Understanding diffusion models for adversarial robustness. In The Eleventh International Conference on Learning Representations, 2023. [50] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021. [51] Wuzhe Xu, Yulong Lu, Lian Shen, Anqing Xuan, and Ali Barzegari. Diffusion-based models for unpaired superresolution in fluid dynamics. arXiv preprint arXiv:2504.05443, 2025. [52] Wuzhe Xu, Yulong Lu, Sifan Wang, and Tong-Rui Liu. Improving data fidelity via diffusion model-based correction and super-resolution. arXiv preprint arXiv:2505.08526, 2025. [53] Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications. ACM computing surveys, 56(4):1–39, 2023. [54] Liu Yang, Siting Liu, Tingwei Meng, and Stanley J Osher. In-context operator learning with data prompts for differential equation problems. Proceedings of the National Academy of Sciences, 120(39):e2310142120, 2023. [55] Tony Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In International Conference on Machine Learning (ICML), pages 12697– 12706. PMLR, 2021. [56] Zheyu Zhu, Haoran Zhao, Hongru He, Yiming Zhong, Shangtong Zhang, Yang Yu, and Weinan Zhang. Diffusion models for reinforcement learning: A survey. arXiv preprint arXiv:2311.01223, 2023.

Notation Table

Proofs for Section 4

This section provides the proofs of the theoretical results stated in Section 4. We first establish the exact transformer realization in Lemma 1, then derive the statistical error bounds in Lemmas 2 and 3 and the approximation error bound in Proposition 1. Finally, we combine these results to prove Theorem 1. 27

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

A P REPRINT

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

Meaning

x∈R M ⊂ RD d µ ∈ P(P(RD )) ρ∼µ

Data sample (e.g., an image) with ambient dimension D. Data manifold Intrinsic dimension of data manifold Meta-distribution over tasks (distribution of distributions). Task distribution sampled from the meta-distribution.

t ∈ [δ, 1] β(t), α(t), σ(t) ε ∼ N (0, I) xk ∼ ρ k xk (t) = α(t)xk + σ(t)ε ρk,t pρ,t Sk (xk (t), t)

Diffusion time. Forward SDE coefficients. Standard Gaussian noise. Clean sample from task distribution ρk . Noisy sample from task distribution ρk at time t. Law/distribution of xt under task ρk . Density of the perturbed task distribution ρt ; for task ρk , we write pk,t := pρk ,t . Time-t score function: Sk (xk (t), t) = ∇x log ρk,t (x) .

T (ρk )

In-context solution operator.

n ρ̂ = {x̂(1) , . . . , x̂(n) } {ρ̂k }K k=1 Tθ (ρ̂k )(xk (t), t)

Number of samples in a prompt/context set. Empirical distribution (also used as the in-context prompt). K training tasks sampled from µ. Neural operator predicting score function at xt,k , t given prompt ρ̂k .

K J NT S = {sk }K k=1 Nt sk = {sk,j,l }J, j=1, l=1 sk,j,l = {ρ̂k,j , xk,j (tk,j,l ), tk,j,l } (1) (n) ρ̂k,j = {xk,j , . . . , xk,j }

Number of task distributions in the meta-training set. Number of context-query episodes sampled for each training task. Number of diffusion-time and noise samples for each episode. Total training set. Training set for kth task. One instance containing prompt + a query sample. In-context prompt for the jth episode for kth task.

T (LT , mT , dembed , ℓ, LFFN , wFFN , R, κ, D) LT mT dembed ℓ LFFN wFFN RL κ

Transformer network class (hypothesis class). Number of transformer blocks (depth). Maximum number of attention heads per block. Embedding dimension. Number of hidden tokens (sequence length inside the transformer). Number of layers in each feed-forward subnetwork. Hidden width of each feed-forward subnetwork. Uniform bound on the pointwise DSM loss; see (30). Uniform bound on parameter magnitude: ∥θ∥∞ ≤ κ.

L(T ) Ln (Tθ ) Ln,K (Tθ ) Ln,K,J (Tθ ) LS (Tθ ) T ∗ ∈ arg min L(T ) Tθ̂ ∈ T Tθ∗ ;S ∈ arg min LS (Tθ )

Population DSM loss. Finite-context population DSM risk with n context samples. Task-discretized risk over K sampled task distributions. Task- and episode-discretized risk with J episodes per task. Empirical DSM loss over S. Population-risk minimizer. Transformer approximator of T ∗ obtained by exactly realizing the empirical kernel score estimator. Empirical risk minimizer within hypothesis class T given dataset S.

Tθ ∈T

B.1

x=xk (t)

Proof of Lemma 1: Exact Transformer Realization

For notational convenience, throughout the remainder of the appendix, let ei ∈ R2D+6 denote the ith standard basis vector, whose ith entry is one and whose remaining entries are zero. We also use 0m ∈ Rm to denote the mdimensional zero column vector. Finally, Im×n and 0m×n denote the m × n rectangular identity matrix and zero matrix, respectively. 28

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

A P REPRINT

Proof. It suffices to show that for each coordinate j = 1, · · · , D, there exists a Fθ ∈ F such that for x(t) ∈ XtC , Pn (i) ∥x(t)−α(t)x(i) ∥2 x exp − i=1 j 2σ 2 (t) . (40) Fθ (s) = Kj (s) := P n ∥x(t)−α(t)x(i) ∥2 i=1 exp − 2σ 2 (t) Note that since t is given as an input, both α(t) and σ(t) are accessible. After the embedding layer E and the positional encoding layer PE, we have   (1) x · · · x(n) x(t) 0 · · · · · · 0D+3  ∈ Rdembed ×ℓ = R(2D+6)×(n+1) . (41) H = PE + E(s) =  D+3 I1 ··· ··· Iℓ  1 ··· ··· 1 √ Here Iq = (cos(qπ/(2ℓ)), sin(qπ/(2ℓ)))⊤ and M is defined in (25). Throughout the proof, take u = c M , with an absolute constant c, as a uniform bound for the entries of H, H1 , H2 , H3 on XtC . We write B(m, L, w) for a transformer block with at most m heads and an FFN of depth L and width w. We use hi to denote the i-th column/token in the current embedding matrix. Let us present a step-by-step construction of Kj (s) based on H as the input to the first layer of transformer, guided by Interaction Lemma 6. (i)

• Implementation of xj for i = 1, · · · , n. For each 1 ≤ i ≤ n, we define each attention head Ai , with Vi = eD+1 e⊤ j The data kernels are in the form   0(2D+2)×3  0(2D+3)×(2D+3) Qdata = Kidata =  i 0 0 1

(42)

where Qdata , Kidata ∈ R(2D+3)×(2D+6) . Both the Qdata and Kidata have the last row and last column’s element equal i i to 1 for all i = 1, · · · , n. By Interaction Lemma 6, we can construct Ai , 1 ≤ i ≤ n, such that hi interacts with itself only, i.e.,     0D 0D  =  x(i)  , [Ai (H)]i = ReLU ⟨Qdata hi , Kidata hi ⟩ Vi hi = ReLU(1)  x(i) i j j 0D+5 0D+5 and the output [Ai (H)]q = 0 when q ̸= i. P Let MHA(H) := i Ai (H) denote the multi-head attention, then the residual multi-head attention yields  (1)  x · · · x(n) x(t) (n)  x(1) · · · xj 0  j   0  ··· ··· 0    . . . . . MHA(H) + H =  . .. .. ..  .    0  ··· ··· 0   I1 · · · · · · In+1 1 ··· ··· 1 Let the feed-forward component have all weights and biases equal to zero. Then  (1)  x · · · x(n) x(t) (1) (n)  xj  · · · xj 0    0  ··· ··· 0    . . . . . H1 := B1 (H) =  . .. .. ..  .    0  ··· ··· 0   I1 · · · · · · In+1 1 ··· ··· 1 In total, this process needs B1 ∈ B(n, 1, 2D + 6). The upper bound of the weights parameter in B1 is ∥θB1 ∥∞ ≤ O(D4 ℓ2 u2 b2 ) = O(D4 n2 u2 b2 ). 29

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

A P REPRINT

• Implementation of −α(t)x(i) for i = 1, · · · , n. For each 1 ≤ i ≤ n, we define each attention head Ai , with   0(D+1)×D 0(D+1)×(D+6) 0D×(D+6)  ∈ R(2D+6)×(2D+6) . Vi (t) = −α(t)ID×D 05×D 05×(D+6) , Kidata the same as those used in (42). By Interaction Lemma 6, we can construct Ai , 1 ≤ i ≤ n, We chose Qdata i such that hi interacts with itself only, i.e.,     0D+1 0D+1 [Ai (H1 )]i = ReLU(⟨Qdata hi , Kidata hi ⟩)Vi hi = ReLU(1) −α(t)x(i)  = −α(t)x(i)  , i 05 05 and [Ai (H1 )]t = 0 when t ̸= i. Therefore, the residual multi-head attention yields 

x(1)  x(1) j  −α(t)x(1)  MHA(H1 ) + H1 =  0   0  I1 1

··· ··· ··· ··· ··· ··· ···

x(n) (n) xj −α(t)x(n) ··· ··· ··· ···

 x(t) 0   0   . 0   0   In+1 1

Let the feed-forward component have all weights and biases equal to zero. Then   x(1) ··· x(n) x(t) (n)  x(1) ··· xj 0  j   −α(t)x(1) · · · −α(t)x(n) 0    . H2 := B2 (H1 ) =  0 ··· ··· 0     0 ··· ··· 0    I1 ··· ··· In+1 1 ··· ··· 1 In total, this process needs B2 ∈ B(n, 1, 2D + 6). The upper bound of the weights parameter in B2 is ∥θB2 ∥∞ ≤ O(D4 ℓ2 u2 b2 ) = O(D4 n2 u2 b2 ). • Implementation of x(t) − α(t)x(i) for i = 1, · · · , n. For each 1 ≤ i ≤ n, we define each attention head Ai , with   0(D+1)×D 0(D+1)×(D+6) 0D×(D+6)  ∈ R(2D+6)×(2D+6) , Vi =  ID×D 05×D 05×(D+6) We chose Qdata , Kidata the same as those used in (42). By Interaction Lemma 6, we can construct Ai , 1 ≤ i ≤ n, i such that hi interacts with hn+1 only, i.e., " # " # 0D+1 0D+1 data data [Ai (H2 )]i = ReLU(⟨Qi hi , Ki hn+1 ⟩)Vi hn+1 = ReLU(1) x(t) = x(t) , 05 05 and [Ai (H2 )]t = 0 when t ̸= i. Therefore, the residual multi-head attention yields x(1) (1)  xj  x(t) − α(t)x(1)  MHA(H2 ) + H2 =  0   0  I1 1 

··· ··· ··· ··· ··· ··· ···

x(n) (n) xj x(t) − α(t)x(n) ··· ··· ··· ···

 x(t) 0   0   . 0   0   In+1 1

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

Let the feed-forward component have all weights and biases equal to zero. Then  x(1) ··· x(n) (1) (n)  xj ··· xj  x(t) − α(t)x(1) · · · x(t) − α(t)x(n)  H3 := B3 (H2 ) =  0 ··· ···   0 ··· ···  I1 ··· ··· 1 ··· ···

A P REPRINT

 x(t) 0   0   . 0   0   In+1 1

In total, this process needs B3 ∈ B(n, 1, 2D + 6). The upper bound of the weights parameter in B3 is ∥θB3 ∥∞ ≤ O(D4 ℓ2 u2 b2 ) = O(D4 n2 u2 b2 ). (i) 2

∥ • Implementation of − ∥x(t)−α(t)x + M for i = 1, · · · , n. 2σ 2 (t) Let us define each attention head Ai with Vi = e2D+2 e⊤ 2D+6 , and data kernels in the form

For c ∈ R, define C(c) :=

0 0

0 ∈ R2×3 . c

For s ∈ {−1, 1}, define  0(D+1)×(D+1)  0D×(D+1) W(s) :=  02×(D+1) We then set

0(D+1)×D s √ ID 2σ(t) 02×D

0(D+1)×2 0D×2 02×2

Qdata = W(1), i

 0(D+1)×3 (2D+3)×(2D+6) 0D×3  . ∈R √ C( M )

Kidata = W(−1).

By Interaction Lemma 6, we can construct Ai such that hi interacts with itself only, i.e., [Ai (H3 )]i = ReLU(⟨Qdata hi , Kidata hi ⟩)Vi hi i 1 (i) (i) = ReLU − ⟨(x(t) − α(t)x ), (x(t) − α(t)x )⟩ + M e2D+2 2σ(t)2 ∥(x(t) − α(t)x(i) )∥2 = − + M e2D+2 2σ(t)2 C and [Ai (H3 )]q = 0 for q ̸= i. Here the √ choice of M in (25) ensures that the argument of ReLU is nonnegative on Xt , (i) since ∥x(t) − α(t)x ∥2 /σ(t) ≤ 2b D/σ(δ) + C.

Therefore, the residual multi-head attention yields  x(1) ··· x(n) x(t) (1) (n)  xj ··· xj 0    (1) (n)  x(t) − α(t)x ··· x(t) − α(t)x 0    (1) . )∥2 ∥(x(t)−α(t)x(n) )∥2 MHA(H3 ) + H3 =   − ∥(x(t)−α(t)x + M · · · − + M 0 2 2 2σ(t) 2σ(t)    0 ··· ··· 0    I1 ··· ··· In+1 1 ··· ··· 1 Let the feed-forward component have all weights and biases equal to zero, retaining the common constant M . Then   x(1) ··· x(n) x(t) (1) (n)  xj ··· xj 0    (1) (n)  x(t) − α(t)x ··· x(t) − α(t)x 0    (1) . )∥2 ∥(x(t)−α(t)x(n) )∥2 H4 := B4 (H3 ) =  − ∥(x(t)−α(t)x  + M · · · − + M 0 2σ(t)2 2σ(t)2    0 ··· ··· 0    I1 ··· ··· In+1 1 ··· ··· 1 

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

A P REPRINT

In total, this process needs B4 ∈ B(n, 6, 2D + 6). The upper bound of the weights parameter in B4 is ∥θB4 ∥∞ ≤ O(D4 ℓ2 M 2 · M 2 ) = O(D4 ℓ2 M 4 ). • Implementation of Kj (s) in the (D + 1)-th row and (n + 1)-th column. For the last layer of transformer, we will use softmax activation (with a mask of size n). We can construct V = eD+1 e⊤ D+1 , and data kernels     0 0 0 0 0 0 0 0 .. .. ..  .. .. ..  .. ..   . . . . .  , K data =  . . . , Qdata =    0 0 0 1  1 0 0 0  0

where Qdata ∈ R(2D+3)×(2D+6) has (2D + 2, 2D + 6)-th position equals to 1 and all the other entries are 0, and K data ∈ R(2D+3)×(2D+6) has (2D + 2, 2D + 2)-th position equals to 1 and all the other entries are 0, such that

[A(H4 )]n+1 =

n X

softmax (K data H4 )⊤ Qdata hn+1 i V hi

i=1

(i) ∥x(t)−α(t)x(i) ∥2 x exp M − 2 i=1 j 2σ (t) eD+1 = Kj (s)eD+1 . Pn ∥x(t)−α(t)x(i) ∥2 exp M − 2 i=1 2σ (t)

Pn =

The common factor eM cancels in the ratio. Since Qdata hq is the same for every token q, the attention output is also Kj (s)eD+1 at the prompt tokens. Therefore, the residual attention gives 

x(1) (1) xj + Kj (s) (1)

   x(t) − α(t)x  (1) )∥2 A(H4 ) + H4 =  − ∥(x(t)−α(t)x +M 2σ(t)2   0  I1 1

··· ··· ··· ··· ··· ··· ···

x(n) (n) xj + Kj (s) (n)

x(t) − α(t)x (n) )∥2 − ∥(x(t)−α(t)x +M 2σ(t)2 ··· ··· ···

 x(t) Kj (s)  0    0 .  0   In+1 1

For the final residual FFN, let z denote the (D + 1)-th input coordinate and set f (z) = −b + ReLU(z + b) − ReLU(z − b) − ReLU(z) + ReLU(−z) in row D + 1, with all other output rows zero. Then z + f (z) = max{−b, min{z, b}}. Four ReLU units suffice; identity maps on their nonnegative activations pad the FFN to depth six. Since |Kj (s)| ≤ b, the query output is unchanged on XtC and bounded by b for every noisy query. Therefore, there exists B5 ∈ B(1, 6, 2D + 6) such that, on XtC , 

 x(t) Kj (s)   [H5 ]n+1 := [B5 (H4 )]n+1 =  0D+2  . I  n+1 1 The upper bound of the weights parameter in B5 is ∥θB5 ∥∞ ≤ O(D4 ℓ2 M 2 · 1) = O D4 ℓ2 M 2 . Finally, applying a decoding layer DE to output the element Kj (s), so we get Fθ (s) = Kj (s) on XtC , with |Fθ (s)| ≤ b for all x(t) ∈ RD , as desired. The uniform upper bound for the weight parameters in B5 ◦ B4 ◦ B3 ◦ B2 ◦ B1 is κ ≤ O D4 ℓ2 M 4 . All blocks fit LFFN = 6 and wFFN = 2D+6. By stacking D such networks componentwise, we can represent the full vector-valued kernel regression estimator in (16) on XtC , with ∥Gθ (s)∥∞ ≤ b for all x(t) ∈ RD . 32

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

A P REPRINT

We now fix Cδ2 = 2D log 2 + 24 log(1/σ(δ)),

qδ := P(∥ε∥2 > Cδ ).

(43)

Since E exp(∥ε∥22 /4) = 2D/2 , Markov’s inequality gives 2

qδ ≤ 2D/2 e−Cδ /4 = σ(δ)6 .

(44)

The radius (43) is compatible with the original weight-budget order. Indeed, !2 √ 4b2 D 1 2b D + 2D log 2 + 24 log(1/σ(δ)), + Cδ ≤ 2 σ(δ) σ(δ)2 M ≤ cDb2 Rσ(δ)−2 , κ = c0 D4 (n + 1)2 M 4 ≤ c1

D 8 n 2 b8 R 4 . σ(δ)8

(45)

Here we used D ≥ 1, R ≥ b ≥ 1, and log(1/s) ≤ s−2 for 0 < s ≤ 1. Token budget. The realization uses n + 1 active tokens: the n prompt tokens followed by one query token. This is the token budget used in Lemma 1, Proposition 1, and Theorem 1. B.2

Statistic Error I

Define the pointwise loss 2

fθ (ρ̂, x(t), t, ε) := Tθ (ρ̂)(x(t), t) +

ε . σ(t) 2

Then define the episode-level loss hθ (ρ̂, x) := Et∼U[δ,1], ε∼N (0,I) [fθ (ρ̂, x(t), t, ε)] , and the task-level loss gθ (ρ) := Eρ̂∼ρ⊗n , x∼ρ [hθ (ρ̂, x)] . With these abbreviations, the four losses are as follows: • The semi-discretized population loss is Ln (Tθ ) := Eρ∼µ gθ (ρ) . • The task-discretized loss is K

1 X gθ (ρk ). Ln,K (Tθ ) := K k=1

• The task- and episode-discretized loss is K

Ln,K,J (Tθ ) :=

1 XX hθ (ρ̂k,j , xk,j ) . KJ j=1 k=1

• The fully discretized empirical loss is K

LS (Tθ ) :=

T XXX 1 fθ (ρ̂k,j , xk,j (tk,j,l ), tk,j,l , εk,j,l ) . KJNT j=1

k=1

l=1

For the statistical analysis only, set fθC = fθ 1{∥ε∥2 ≤Cδ } ,

Uδ = max{1, b + Cδ }.

(46)

C Use the same averages to define hC θ , gθ , and the corresponding population and empirical losses. On this event, ∥H∥∞,∞ ≤ Uδ . This restriction is used only in the proof; the noise law and the training objective remain unchanged.

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

A P REPRINT

For the statistical bounds below, use the concrete width choice already available in Lemma 1: w = dembed = wFFN , with w = 2D + 6, LT = 5, LFFN = 6, and mT = n for the main theorem. Each scalar network has at most PT := w(D + 1) + LT 3w2 mT + LFFN (w2 + w) (47) free coefficients. This counts the embedding, decoder, one trainable base matrix for each Q/K/V map, and all affine FFN matrices and biases. The prescribed time templates in (19) are not free coefficients. Tied weights and unused heads only reduce this count. We give an explicit covering constant by a layer recurrence. Set s0 = σ(δ), τ = 2 + s−1 0 , and initialize u0 = 1 + DκUδ ,

v0 = DUδ .

For k = 0, . . . , LT − 1, define qk = wτ κuk ,

ek = w(τ uk + τ κvk ),

ak = mT (qk + ℓwqk3 ),

ck = mT (1 + 4wqk2 + 3ℓwqk2 )ek , dk = vk + ck .

zk = uk + ak ,

For the FFN in that block, set rk,0 = zk , tk,0 = dk , and, for j = 0, . . . , LFFN − 1, let rk,j+1 = κ + wκrk,j ,

tk,j+1 = wrk,j + 1 + wκtk,j .

Complete the block recursion with uk+1 = zk + rk,LFFN ,

vk+1 = dk + tk,LFFN .

The resulting mean-network sensitivity is Knet := w(uLT + κvLT ).

(48)

In particular, a maximum coefficient perturbation of ε changes every output coordinate by at most Knet ε, uniformly over the restricted inputs and all t ∈ [δ, 1]. Lemma 4 (Covering number estimation). For the transformer class with the width and fixed time-operation convention above, the restricted loss class satisfies, for 0 < η ≤ RL , A0,δ , log N η, {fθC : Tθ ∈ T }, ∥ · ∥L∞ ≤ DPT log η

(49)

where a sufficient loss-covering constant is A0,δ := 4κ(Kloss + 1) + 2RL ,

Kloss :=

2D(R + b) Knet . s40

(50)

C The same bound holds for hC θ and gθ .

Proof. We prove the network estimate directly. Each scheduled matrix has entry bound τ κ; its parameter perturbation is at most τ ε. Fixed time templates have zero parameter difference. A matrix-vector product therefore has magnitude at most qk and perturbation at most ek ε. The ReLU attention head has magnitude at most ℓwqk3 and perturbation at most 3ℓwqk2 ek ε. For the normalized softmax readout, if the values have magnitude at most B, then | Readout(u, v) − Readout(u′ , v ′ )| ≤ 2B∥u − u′ ∥∞ + ∥v − v ′ ∥∞ . Indeed, the sum of the absolute score derivatives is at most 2B, and the value derivatives are the softmax weights, whose sum is one. Interpolating scores and values proves the inequality. Since the attention-score perturbation is at most 2wqk ek ε, the softmax head has perturbation at most (1+4wqk2 )ek ε. Summing these sufficient ReLU and softmax bounds over the heads gives ak , ck . Residual addition and the affine/ReLU FFN estimates give the stated recursions. The embedding and decoder give the initial values and (48). These estimates use the original blocks and are uniform in time. By (29), the loss conversion is |fθ − fθ′ | ≤

2D(R + b) ∥Gθ − Gθ′ ∥∞ . s40 34

(51)

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

A P REPRINT

Thus the restricted loss is Kloss -Lipschitz in the coefficient maximum norm. Both restricted losses vanish outside the same noise event. Partition the DPT -dimensional coefficient cube into cells of side r = η/[2(Kloss + 1)] and select one admissible representative from each occupied cell. Allowing endpoint cells, at most DPT DPT 4κ(Kloss + 1) + 2η A0,δ DPT = (2κ/r + 2) ≤ η η representatives suffice. Their coefficient distance from any member of the same cell is at most 2r, so their loss distance is strictly below η. The centers belong to the legal bounded-output class. This proves (49). All network operations and the restriction event are measurable, and the common loss envelope gives integrability. Probability averaging cannot increase uniform distances, so the same parameter grid covers hC and g C . We are now ready to prove Lemma 2. Proof. Note that ES [Ln (Tθ∗ ;S ) − LS (Tθ∗ ;S )] ≤ ES sup (Ln (Tθ ) − LS (Tθ )). Tθ ∈T

Since the tail contribution to the empirical loss is nonnegative and the population tail is at most RL qδ , for every training set, C L sup (Ln − LS )(Tθ ) ≤ sup (LC (52) n − LS )(Tθ ) + R qδ . Tθ ∈T

Tθ ∈T

It remains to apply the hierarchical symmetrization argument to these auxiliary losses. For each fixed Tθ , decompose C C C LC n (Tθ ) − LS (Tθ ) = [Ln (Tθ ) − Ln,K (Tθ )] C + [LC n,K (Tθ ) − Ln,K,J (Tθ )] C + [LC n,K,J (Tθ ) − LS (Tθ )].

The first bracket is the task-level statistical error, the second is the episode-level statistical error, and the third is the time/noise discretization error. For the first term, standard symmetrization gives C C d Eρ1:K sup (LC n − Ln,K )(Tθ ) ≤ 2Eρ1:K RadK ({gθ : Tθ ∈ T }; ρ1:K ). Tθ ∈T

For the second term, condition on all tasks; for the third, condition on all tasks and episodes. Under these conditional laws the samples are independent, though they need not be identically distributed. Independent ghost copies and Jensen’s inequality give the same symmetrization bound with KJ and KJNT signs, respectively. For a function class H evaluated on N observations, write " d N (H; Z1:N ) := Eζ Rad

# N 1 X ζi h(Zi ) , sup h∈H N i=1

where the ζi are independent Rademacher variables. By Dudley’s entropy integral, ! Z RL p 12 d N (H) ≤ inf ERad 2ξ + √ log N (η, H, ∥ · ∥L∞ ) dη 0<ξ≤RL N ξ q 12RL ≤ inf 2ξ + √ DPT log(A0,δ /ξ) 0<ξ≤RL N  s  √ RL  A0,δ N ≤√ 12 DPT log + 2 . RL N √ The last inequality is obtained by letting ξ = RL / N . Apply this estimate at N = K, KJ, KJNT and sum the three bounds. Adding (52) proves (35). For the stated network parameters, DPT = O(nD3 ). Equations (43), (48), and (50) imply √ A0,δ KJNT nDKJNT log ≤ c log , RL σ(δ) 35

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

A P REPRINT

where c depends only on b, R and the fixed depths. Indeed, Cδ2 = 2D log 2 + 24 log(1/σ(δ)). For fixed LT = 5 and LFFN = 6, the recursions defining Knet are finite polynomials with nonnegative coefficients in D, w, mT , ℓ, κ, Uδ and τ = 2 + σ(δ)−1 . Hence (50) has polynomial growth in n, D, σ(δ)−1 under the stated weight budget, and its logarithm has the displayed order. In particular, the bias-inclusive count satisfies DPT ≤ 3136nD3 when n, D ≥ 1; the statistical rate is unchanged. Substituting these bounds, RL = D(R + b)2 σ(δ)−4 , and qδ ≤ σ(δ)6 into (35) proves (36).

B.3

Statistical Error II

Proof. Note that ES [LS (Tθ̂ ) − Ln (Tθ̂ )] = ES [Ln,K (Tθ̂ ) − Ln (Tθ̂ )] + ES [Ln,K,J (Tθ̂ ) − Ln,K (Tθ̂ )] + ES [LS (Tθ̂ ) − Ln,K,J (Tθ̂ )]. iid

Since the variables gθ̂ (ρk ) with ρk ∼ µ are bounded by RL , Hoeffding’s inequality gives, for any u > 0, 2Ku2 P(Ln,K (Tθ̂ ) − Ln (Tθ̂ ) ≥ u) ≤ exp − L 2 . (R ) Thus with probability at least 1 − ϵ, r Ln,K (Tθ̂ ) − Ln (Tθ̂ ) ≤ R and therefore

log(1/ϵ) , 2K

log(1/ϵ) + 2ϵRL . 2K Conditional Hoeffding bounds at the episode and time–noise levels give the same expressions with KJ and KJNT . Combining these three bounds yields (53). ES [Ln,K (Tθ̂ ) − Ln (Tθ̂ )] ≤ R

For every 0 < ϵ < 1, the same argument gives r ES [LS (Tθ̂ ) − Ln (Tθ̂ )] ≤ R

log(1/ϵ) 2

1 1 1 √ +√ +√ KJNT K KJ

+ 6ϵRL .

(53)

For the cutoff used in Theorem 1, set s0 = σ(δ) ∈ (0, 1) and choose ϵ = s60 . The explicit inequality (53) becomes ES [LS (Tθ̂ ) − Ln (Tθ̂ )] 1 1 1 D(R + b)2 p √ √ √ 3 log(1/s0 ) + + + 6D(R + b)2 s20 . ≤ s40 KJN K KJ T

(54)

Here log(1/s60 ) = 6 log(1/s0 ) and RL = D(R + b)2 s−4 0 . For each fixed cutoff, this gives the order in (37); its implied constant can depend on the fixed loss envelope. If s0 = 1, taking ϵ = 1/2 and using K, J, NT ≥ 1 gives the same coarse order. When the cutoff varies with n, the theorem proof retains the explicit dependence in (54).

B.4

Approximation Error III

Fix a task distribution ρ and t ∈ [δ, 1]. Define ∥z∥22 Kσ(t) (z) = exp − , 2σ(t)2 n 1X N̂n (x(t)) = Kσ(t) x(t) − α(t)x(i) x(i) , n i=1

n 1X D̂n (x(t)) = Kσ(t) x(t) − α(t)x(i) , n i=1

N (x(t)) = Ex∼ρ [Kσ(t) (x(t) − α(t)x) x],

D(x(t)) = Ex∼ρ [Kσ(t) (x(t) − α(t)x)]. 36

(55)

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

A P REPRINT

Here D denotes the ambient dimension, whereas D(x(t)) denotes the population kernel denominator. Since the Gaussian kernel is strictly positive, Bayes’ formula gives K(s) =

N̂n (x(t)) D̂n (x(t))

E(x | x(t)) =

N (x(t)) . D(x(t))

(56)

The noisy-query density and its score satisfy 1 D(x(t)), (2πσ(t)2 )D/2 α(t) x(t) + E(x | x(t)). T ∗ (ρ)(x(t), t) = − 2 σ(t) σ(t)2 pρ,t (x(t)) =

(57) (58)

Differentiation under the integral in (57) is justified by the compact support of ρ and the smooth Gaussian kernel. Lemma 5 (Finite prompt approximation error). Under Assumptions 1 and 2, let Tθ̂ be the candidate from Lemma 1 with C = Cδ and weight budget (45). For every task distribution ρ and every t ∈ [δ, 1], Eρ̂,x,ε ∥Tθ̂ (ρ̂)(x(t), t) − T ∗ (ρ)(x(t), t)∥22 ≤

Cb,d AM D + 4Db2 σ(δ)2 , nσ(δ)d+4

(59)

where Cb,d = 32e 2d b2 , the prompt is sampled from ρ⊗n , and the clean query and noise are independent of the prompt. Proof. Assume first that α(t) > 0. Fix x(t), and suppress the arguments of N, D, N̂n , D̂n within this paragraph. The n 1X N N N̂n − D̂n = Kσ(t) x(t) − α(t)x(i) x(i) − . D n i=1 D

Control at a fixed noisy query. centered numerator is

Each summand has mean zero, since Ex∼ρ [Kσ(t) (x(t) − α(t)x) (x − N/D)] = 0. Moreover, N/D and every clean sample lie in [−b, b]D . Independence therefore yields # " 2 2 N 1 N 2 Eρ̂ N̂n − D̂n = Ex∼ρ Kσ(t) (x(t) − α(t)x) x − D n D 2 2 ≤

4Db2 2 Ex∼ρ [Kσ(t) (x(t) − α(t)x) ]. n

(60)

On the event D̂n ≥ D/2, divide (60) by D2 /4 to obtain 2 h i 16Db2 E x∼ρ [Kσ(t) (x(t) − α(t)x) ] . Eρ̂ ∥K(s) − E(x | x(t))∥22 1{D̂n ≥D/2} ≤ n D(x(t))2 On the complementary event, the squared error is at most 4Db2 . Chebyshev’s inequality gives ! 2 4 Var(D̂n ) 4 Ex∼ρ [Kσ(t) (x(t) − α(t)x) ] Pρ̂ (D̂n < D/2) ≤ = −1 . D2 n D(x(t))2 Combining both events proves 2

Eρ̂ ∥K(s) − E(x | x(t))∥22 ≤

32Db2 Ex∼ρ [Kσ(t) (x(t) − α(t)x) ] . n D(x(t))2

Integration against the noisy-query law. Using (57), the expectation of the ratio in (61) is " # 2 Ex∼ρ [Kσ(t) (x(t) − α(t)x) ] Ex(t) D(x(t))2 Z R 2 Kσ(t) (x(t) − α(t)x) dρ(x) 1 M dx(t). = D(x(t)) (2πσ(t)2 )D/2 RD

(61)

(62)

The factor D(x(t)) in the query density cancels one denominator. We bound the resulting integral directly, without a uniform lower bound on D(x(t)). 37

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

A P REPRINT

Localization on the clean support. Let A1 , . . . , AJ ′ be a finite measurable partition of M; the symbol J ′ here is only a partition size and is unrelated to the number of training episodes. Discard cells of zero ρ-mass. For nonnegative aj and positive bj , P X aj X aj aj Pj P = ≤ . bj j bj k bk j j Consequently, R

R 2 X Aj Kσ(t) (x(t) − α(t)x)2 dρ(x) K (x(t) − α(t)x) dρ(x) σ(t) M R ≤ . D(x(t)) Kσ(t) (x(t) − α(t)x) dρ(x) Aj j

(63)

For every positive-mass cell, define its conditional probability law by ρ(B | Aj ) = ρ(B ∩ Aj )/ρ(Aj ) for measurable B. The Pcommon mass ρ(Aj ) cancels from the numerator and denominator of the corresponding term. Thus (62) is at most j Ij (t), where Z 2 Ex∼ρ(·|Aj ) [Kσ(t) (x(t) − α(t)x) ] 1 Ij (t) := dx(t). (64) (2πσ(t)2 )D/2 RD Ex∼ρ(·|Aj ) [Kσ(t) (x(t) − α(t)x)] No lower bound on the mass of a partition cell is needed. Gaussian integral on each cell. fixed query,

Draw x, x′ independently from ρ(· | Aj ). Convexity of u 7→ 1/u implies, at each 2

Ex [Kσ(t) (x(t) − α(t)x) ] 2 ≤ Ex [Kσ(t) (x(t) − α(t)x) ] Ex′ [1/Kσ(t) (x(t) − α(t)x′ )] Ex′ [Kσ(t) (x(t) − α(t)x′ )] # " 2 Kσ(t) (x(t) − α(t)x) . = Ex,x′ Kσ(t) (x(t) − α(t)x′ ) All integrands are nonnegative. Tonelli’s theorem therefore yields " # Z 2 Kσ(t) (x(t) − α(t)x) 1 Ij (t) ≤ Ex,x′ dx(t) . (2πσ(t)2 )D/2 RD Kσ(t) (x(t) − α(t)x′ )

(65)

Completing the square gives 2∥x(t) − α(t)x∥22 − ∥x(t) − α(t)x′ ∥22 = ∥x(t) − α(t)(2x − x′ )∥22 − 2α(t)2 ∥x − x′ ∥22 . Hence Z 2 Kσ(t) (x(t) − α(t)x) 1 dx(t) (2πσ(t)2 )D/2 RD Kσ(t) (x(t) − α(t)x′ ) Z α(t)2 ∥x − x′ ∥22 1 ∥x(t) − α(t)(2x − x′ )∥22 = exp exp − dx(t) σ(t)2 2σ(t)2 (2πσ(t)2 )D/2 RD α(t)2 ∥x − x′ ∥22 . = exp σ(t)2 The normalized Gaussian integral is exactly one in ambient dimension D.

(66)

Choice of the partition. Take a cover of M by balls Bj of radius rt = σ(t)/(2α(t)) with centers on M and cardinality N (M, rt ). Turn the cover into a measurable partition by setting [ A1 = M ∩ B 1 , Aj = (M ∩ Bj ) \ Bi . i<j

Every cell has diameter at most σ(t)/α(t). Thus (65)–(66) imply Ij (t) ≤ e. Applying (61)–(62) and (28), σ(t) 32eDb2 Ex(t),ρ̂ ∥K(s) − E(x | x(t))∥22 ≤ N M, n 2α(t) 32eDb2 Cb,d AM D ≤ N (M, σ(t)/2) ≤ . (67) n nσ(t)d The second inequality uses α(t) ≤ 1 and monotonicity of covering numbers; the last evaluates (28) at σ(t)/2 ≤ 1. 38

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

A P REPRINT

Transfer to the transformer candidate. On Eδ := {∥ε∥2 ≤ Cδ }, Lemma 1 gives Gθ̂ = K. On the complement, both Gθ̂ and E(x | x(t)) have coordinates in [−b, b]. Therefore α(t)2 E∥Gθ̂ − E(x | x(t))∥22 σ(t)4 α(t)2 4Db2 α(t)2 ≤ E∥K − E(x | x(t))∥22 + qδ 4 σ(t) σ(t)4 Cb,d AM D + 4Db2 σ(δ)2 . ≤ nσ(δ)d+4

E∥Tθ̂ − T ∗ ∥22 =

(68)

This proves (59). If α(t) = 0, both score functions equal −x(t)/σ(t)2 , so the score error is zero and the same bound holds. Under the stated assumptions, we are now ready to prove Proposition 1. Proof. The candidate satisfies the budgets in (33) by (45) and the token-budget discussion. The prompt is conditionally independent of the clean query and its perturbation, given the task. Therefore (58) gives ε ∗ E T (ρ)(x(t), t) + ρ, t, x(t), ρ̂ = 0. σ(t) Expanding the difference of the two squared losses, the cross term has conditional expectation zero. Consequently, Ln (Tθ ) − L(T ∗ ) = Eρ,ρ̂,x,t,ε ∥Tθ (ρ̂)(x(t), t) − T ∗ (ρ)(x(t), t)∥22 .

(69)

All quantities are integrable: the two mean terms are bounded and σ(t) ≥ σ(δ) > 0. Averaging (59) over the task and time distributions proves (34).

B.5

Proof of Theorem 1

Proof. Set s0 = σ(δ) and 1 1 1 ΦK,J,NT := √ + √ +√ . KJNT K KJ We use the inequality (54) for Statistical Error II. Since n ≥ 2 and D, K, J, NT ≥ 1, r p √ nDKJNT 3 log(1/s0 ) ≤ 3 nD3 log . s0 Thus its statistical contribution is absorbed into Statistical Error I, while its tail contribution is 6D(R + b)2 s20 . By combining (36), (54), and (34) in (32), we obtain ES [Ln (Tθ∗ ;S ) − L(T ∗ )] ≤ (C1′ + +

√

D(R + b)2 3) σ(δ)4

s nD3 log

nDKJNT σ(δ)

1 1 1 √ +√ +√ KJNT K KJ

(70)

Cb,d AM D + D[4b2 + 7(R + b)2 ]σ(δ)2 . nσ(δ)d+4

The coefficient 7(R+b)2 combines the localization remainder (R+b)2 from Statistical Error I and the tail contribution 6(R + b)2 from Statistical Error II. Choose δ such that σ(δ) = n−1/(d+6) , as in the statement of the theorem. Then 1 nDKJNT log = log n1+ d+6 DKJNT ≤ log(n2 DKJNT ), σ(δ) and

2 1 = n− d+6 , nσ(δ)d+4

σ(δ)2 = n− d+6 .

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

The loss bound becomes RL =

A P REPRINT

D(R + b)2 = D(R + b)2 n4/(d+6) , σ(δ)4

so its dependence on n must be included in the statistical term. The candidate from Proposition 1 uses the n + 1 active-token realization of Lemma 1. Its weight budget satisfies 8 2 8 4 D n b R 8 2d+20 d+6 b8 R4 κ=O = O D n . σ(δ)8 Therefore ES [Ln (Tθ∗ ;S ) − L(T ∗ )] √ 4 p 1 1 1 1 ≤ (C1′ + 3)(R + b)2 D5/2 n 2 + d+6 log(n2 DKJNT ) √ + √ +√ KJNT K KJ

(71)

+ D[Cb,d AM + 4b2 + 7(R + b)2 ]n− d+6 . √ With C1 = (C1′ + 3)(R + b)2 and C2 = Cb,d AM + 4b2 + 7(R + b)2 , absorbing the fixed factors proves (31).

Auxiliary Lemmas

For completeness, we state the following lemmas; Their proof can be found in [40]. Lemma 6 (Interaction Lemma). Let H = [ht ]1≤t≤ℓ ∈ Rdembed ×ℓ be an embedding matrix such that (d −2):(dembed −1) ht embed = It and hdt embed = 1. Fix 1 ≤ t1 , t2 ≤ ℓ, 1 ≤ i ≤ dembed , and ℓ ∈ N. Suppose dembed ≥ 5 and ∥H∥∞,∞ < U for some U > 0, and the data kernels Qdata ∈ R(dembed −3)×dembed (the first (dembed − 3) rows in the query matrix Q) and K data ∈ R(dembed −3)×dembed (the first (dembed − 3) rows in the key matrix K) satisfy max{∥Qdata ∥∞,∞ , ∥K data ∥∞,∞ } ≤ κ. Then one can construct an attention head A with ReLU activation such that ReLU(⟨Qdata ht , K data ht2 ⟩)ei if t = t1 , [A(H)]t = 0 otherwise. The weight parameters of this attention head satisfies ∥θA ∥∞ = O(d4embed κ2 ℓ2 U 2 ). The same positional selection applies to any linear value map V , replacing ei by V ht2 . The weight bound becomes the maximum of the original bound and ∥V ∥∞,∞ . Lemma 7 (Decrementing Lemma). Let dembed ≥ 5 and H = [ht ]1≤t≤ℓ ∈ Rdembed ×ℓ , be an embedding matrix such (d −2):(dembed −1) = (It1 , It2 ) = It and hdt embed = 1. Then for any r1 , r2 with 1 ≤ r1 ≤ r2 ≤ dembed − 3 that ht embed and any k1 , k2 with 1 ≤ k1 , k2 ≤ ℓ and any M > 0, there exists a six-layer residual feed-forward network (FFN) such that   h t  if t ∈ {1, · · · , k1 } ∪ {k2 , · · · , ℓ}    (h ) t 1      ..      .         (h )    t r −1 1         (h ) − M    t r1   ..     . FFN(ht ) + ht =    (ht )r − M  otherwise  2       (h )   t r +1    2      .    .  .          (ht )dembed −3          It    1 Additionally, we have ∥θFFN ∥∞ ≤ O(ℓM ). 40

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

Additional Experimental Details

D.1

Network Architecture

A P REPRINT

We provide additional implementation details for the architectures used in our experiments. While the models employed in practice are more expressive than the simplified versions presented in the main text for theoretical analysis, they retain the same core design principles. All three methods use a time-conditioned, multi-resolution U-Net as the denoising backbone, and differ primarily in how the prompt set is encoded and injected. In ICL DM, the prompt set is represented by multi-scale features that interact with the noisy query through cross-attention. In Cond DM, the prompt set is compressed into a single learned embedding that modulates the U-Net through FiLM-style conditioning. In SNO, the learned prompt embedding is replaced by a projected kernel mean embedding (KME), while the denoising U-Net remains otherwise the same as in Cond DM. ICL DM. Our implementation of ICL DM uses a time-conditioned, multi-resolution U-Net with prompt-dependent cross-attention. Given an unordered prompt set ρ̂ = {x(i) }ni=1 , a noisy query x(t) and diffusion time t, we first encode the diffusion time using a sinusoidal embedding followed by a multilayer perceptron. The prompt set is then processed by a prompt encoder composed of an input convolution, an initial downsampling layer, and a hierarchy of residual blocks with time modulation. At each resolution, a set-transformer block applies self-attention across the prompt elements, allowing information to be exchanged before the features are aggregated by permutation-invariant mean pooling. This produces multi-scale key-value feature maps that summarize the prompt set. In parallel, the noisy query x(t) is passed through a U-Net encoder-decoder with residual blocks at corresponding resolutions, while crossattention modules at every scale in both the encoder and decoder inject the pooled prompt features into the query pathway. After a final upsampling stage to the original spatial resolution, the network outputs the predicted noise field, which serves as the model prediction for denoising score matching. Specifically, we use a U-Net architecture with 4 resolution levels, where the number of channels at each level is [64, 128, 256, 512]. The timestep embedding is produced by a two-layer MLP with hidden dimension 256, each residual block contains two convolutional layers, cross-attention uses 4 heads at every scale, and the prompt set transformer uses 2 attention heads. Cond DM. Our conditional diffusion model uses a time-conditioned, multi-resolution U-Net in which the prompt set is compressed into a single global conditioning vector rather than injected through cross-attention. Given an unordered prompt set ρ̂ = {x(i) }ni=1 , a noisy query x(t), and diffusion time t, we encode t using a sinusoidal embedding followed by a multilayer perceptron. The prompt set is passed through a prompt encoder consisting of an input convolution, an initial downsampling layer, and a hierarchy of residual blocks, and only the deepest prompt feature is retained. This feature is averaged over the prompt dimension to ensure permutation invariance, aggregated over spatial locations by a learned attention-based pooling module, and mapped to a prompt embedding vector. The resulting prompt embedding is added to the time embedding and used to modulate all residual blocks of the denoising U-Net through FiLM-style conditioning. The denoising network follows an encoder-decoder architecture with skip connections and includes a self-attention block at the bottleneck to enhance global spatial interactions. A final upsampling stage restores the original spatial resolution, after which the network predicts the noise field used for denoising score matching. Specifically, the U-Net has 4 resolution levels with channel dimensions [64,128,256,512]; the timestep embedding and prompt embedding both have dimension 256; the spatial attention pooling module uses a hidden dimension of 128; and both the deepest prompt self-attention block and the bottleneck self-attention block use 2 attention heads. SNO. Our implementation of SNO uses a time-conditioned, multi-resolution U-Net in which the conditioning information is provided by a kernel mean embedding (KME) of the prompt set rather than by cross-attention over prompt features. Given a KME vector associated with the prompt set, a noisy query x(t), and diffusion time t, we first encode the diffusion time using a sinusoidal embedding followed by a multilayer perceptron. The KME vector is then mapped by a two-layer projection network to the same latent dimension as the timestep embedding, and the resulting conditioning vector is added to the time embedding. This combined embedding is used to modulate all residual blocks in the denoising U-Net through FiLM-style conditioning. Other than replacing the prompt encoder and attention-based prompt pooling in Cond DM by the KME projection module, the rest of the denoising architecture is identical to the conditional diffusion model baseline. Specifically, we use a U-Net architecture with 4 resolution levels, where the number of channels at each level is [64, 128, 256, 512]. The timestep embedding and projected KME embedding both have dimension 256, and the bottleneck self-attention block uses 2 attention heads. D.2

Training Details

All models are trained using the Adam optimizer with a learning rate of 1e-4 and a batch size of 32. 41

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

A P REPRINT

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

For sampling, we set δ = 10−4 and discretize the reverse-time SDE using Nrev = 1000 Euler–Maruyama steps for all experiments. D.3

Evaluation Metrics

The metrics used to evaluate the performance of the models are as follows: • Scalar-Valued Empirical Probability Density Functions (PDFs). For each ensemble of scalar field values {ui }N i=1 := {ui,p,q | i = 1, . . . , N ; p, q = 1, . . . , Nx }, we estimate the empirical probability density function (PDF) of the scalar variable u using Kernel Density Estimation (KDE): Ns 1 X u − uk PDF(u) := K , Ns = N × Nx2 , (72) Ns h h k=1

where K(·) is a smooth kernel function, typically chosen as the Gaussian kernel, 2 1 z K(z) = √ exp − , 2 2π

(73)

and h > 0 denotes the bandwidth parameter controlling the smoothness of the resulting density estimate. • Jensen–Shannon divergence (JSD) of PDFs. Given two empirical datasets {xi }ni=1 and {yj }m j=1 , we first estimate their corresponding probability density functions p and q. The Jensen–Shannon divergence between p and q is defined by JSD(p∥q) :=

1 1 KL(p∥m) + KL(q∥m), 2 2

m :=

1 (p + q), 2

(74)

where KL denotes the Kullback–Leibler divergence. • Energy spectra. For a scalar field f (x) defined on a uniform grid {xn }, the energy spectrum describes how the energy of f is distributed across Fourier modes. It is defined by 2

Ef (k) =

X X |k|=k

f (xn ) exp (−j 2π k · xn /L) ,

(75)

where xn denotes the grid points, j is the imaginary unit, k is the wave vector, and k = |k| is its magnitude. Thus, Ef (k) represents the total energy contained in all Fourier modes with wave number magnitude k. • Mean Energy Log Ratio (MELR). To measure the spectral alignment between two empirical datasets {xi }ni=1 and {yj }m j=1 , we consider the Mean Energy Log Ratio (MELR), defined as X Ex (k) MELR(x, y) = wk log , (76) Ey (k) k

where Ex (k) and Ey (k) denote the average energy spectra of the two datasets at wave number k, and wk is a user-specifiedPweight function. If wk = 1 for all k, we obtain the unweighted mean energy log ratio. If wk = Ey (k)/ j Ey (j), we obtain the weighted mean energy log ratio. 42

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

A P REPRINT

• Fréchet Inception Distance (FID). The Fréchet Inception Distance (FID) quantifies the discrepancy between the distributions of real and generated images by comparing the means and covariance matrices of features extracted from a pretrained Inception-v3 network; smaller values indicate better agreement. • Problem-specific FID. We use standard Fréchet-distance formula, but evaluate it in a problem-specific feature space rather than the default feature space induced by an ImageNet-pretrained Inception-v3 network. Specifically, we replace the standard Inception encoder with a feature extractor trained for the dataset of interest. For both the Quick, Draw! dataset [20] and the Food-101 dataset [3], we use a ResNet-18 model pretrained on the corresponding dataset to extract features from real and generated samples, and then compute the FID score from the resulting empirical feature distributions. The corresponding metrics are denoted by FID-QD and FID-Food, respectively.

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

Analyze & explore related research in Conceptio

Related documents