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

1 Introduction

• Standard and conditional diffusion models are not generally designed to infer an unseen target distribution solely from a small unordered context set when conditioning semantics are fixed during training.Adapting may require retraining, fine-tuning, or an additional conditioning representation; direct in-context score-operator learning remains underexplored.Source: §1 · pp. 1–2
• The in-context score operator uses an empirical context set, noisy query, and diffusion time to support reverse-diffusion sampling on unseen tasks without task-specific parameter updates.Tasks are distributions drawn from a meta-distribution, and the operator is trained through amortized denoising score matching across tasks.Source: §1 · p. 2
• The task-dependent score has a Gaussian-kernel conditional-mean representation whose empirical Nadaraya–Watson estimator is exactly representable by the constructed transformer class.This supplies a principled basis for query-dependent interactions among context samples, the noisy query, and diffusion time.Source: §1 · p. 2
• The generalization bound separates finite-context, finite-meta-training-data, and model-complexity effects, with dependence on context size, task and sample counts, and intrinsic and ambient dimensions.The analysis states that the finite-context rate is n^-2/(d+6), with intrinsic dimension governing the exponent and ambient dimension entering through polynomial and logarithmic factors.Source: §1 · pp. 2–3
• Experiments show strong performance in low-intrinsic-dimension regimes, gradual degradation with task complexity, and modest out-of-distribution generalization for Navier–Stokes.The evaluated datasets include Navier–Stokes solution fields, ERA5 climate data, QuickDraw sketches, and Food-101 images.Source: §1 · p. 3

1.1 Related work

• The paper situates its measure-valued task formulation alongside in-context operator learning and distinguishes it from supervised PDE prompting and multimodal operator learning approaches.It also relates its analysis to work on transformer learning algorithms and generalization governed by intrinsic rather than ambient dimension.Source: §1.1 · p. 3
• Prior diffusion methods use contextual examples mainly for paired or query-conditioned image generation, whereas this work learns a score operator over target distributions and interacts directly with context samples.SNO and few-shot diffusion approaches instead condition on a distribution-level embedding or task-dependent feature.Source: §1.1 · p. 3

2.1 Background on Diffusion Models

• Diffusion models learn a score field for a target distribution, while generation reverses the forward noising process through a reverse-time SDE or probability-flow ODE.The forward process progressively corrupts data toward a standard Gaussian; the paper uses Euler–Maruyama to discretize the reverse-time SDE.Source: §2.1 · p. 4

2.2 In-Context Score Function Learning

• Standard denoising score matching learns a time-dependent score for one fixed distribution, whereas the paper amortizes score learning across tasks drawn from a meta-distribution.The task-dependent score operator maps each task distribution to its score function, exploiting shared structure across the family.Source: §2.1 · p. 4, §2.2 · p. 5
• Each task is a probability distribution with a time-dependent score; the population objective trains an operator over task distributions by amortized denoising score matching.The diffusion time is restricted to [δ, 1], with δ > 0, to avoid the zero-noise endpoint where the score can become ill-behaved.Source: §2.1 · p. 4, §2.2 · p. 5
• Because the task distribution is unavailable in practice, a finite sample context is used to infer its score function and evaluate that function at new noisy inputs and diffusion times.The learned operator can adapt to previously unseen distributions through their context samples without task-specific parameter updates.Source: §2.2 · p. 5

3 Methodology

• The neural operator takes an empirical task representation together with a noisy query and diffusion time, approximating the continuous operator from task distributions to score functions.The practical formulation uses a finite context set because the underlying task distribution is not directly observed.Source: §3 · p. 5
• Training discretizes the objective over sampled tasks, context–query episodes, and diffusion-time/noise draws to form a computable empirical denoising score-matching loss.The task- and episode-discretized objective averages over sampled episodes for each training task.Source: §3 · p. 6
• Each training episode pairs an empirical context with a clean query, then constructs a noised query using a sampled diffusion time and Gaussian noise.The training procedure minimizes the empirical loss and updates the operator parameters with Adam.Source: §3 · pp. 6–7
• At inference, the learned score operator is used in the reverse-time SDE to generate samples on a new task with the Euler–Maruyama scheme.The sampling algorithm initializes from Gaussian noise and iterates through reverse steps.Source: §3 · pp. 6–7

4.1 Understanding the Target Score learning: From Kernel Regression to Transformers

• The empirical Nadaraya–Watson score estimator has an exact transformer realization on the specified noisy-query region, with shared parameters independent of task, prompt, query, and time.Task-specific information enters through the context; the shared network implements query–context similarity, kernel-weight normalization, and context aggregation.Source: §4.1 · pp. 9–10

4.2 Generalization Error Analysis

• The finite-context error decays as n^-2/(d+6), while ambient dimension D affects the bound only through polynomial and logarithmic factors.As intrinsic dimension d increases, the convergence rate deteriorates.Source: §4.2 · pp. 11–12
• The statistical error includes K^-1/2, (KJ)^-1/2, and (KJNT)^-1/2 terms, so more within-task samples cannot replace diverse training tasks.The K^-1/2 term persists as J and NT increase; J reduces episode- and diffusion-level errors, while NT reduces only diffusion-level error.Source: §4.2 · p. 12

4.3 Proof Outline

• The proof decomposes generalization error into two statistical terms and an approximation term, controlling them with kernel concentration, complexity bounds, and Hoeffding’s inequality.Exact transformer realization reduces approximation error to that of the empirical kernel estimator on the underlying d-dimensional manifold.Source: §4.3 · p. 12
• For statistical error term II, task-, episode-, and diffusion-sample-level deviations are controlled separately using concentration inequalities because the fitted operator is fixed independently of the training set.Source: §4.3 · p. 13
• Combining the statistical-error lemmas with Proposition 1 gives an upper bound comprising three hierarchical statistical error terms and a finite-context approximation error term.Source: §4.3 · p. 13
• Choosing σ(δ) = n^-1/(d+6) balances the dominant approximation terms and yields the rate n^-2/(d+6).Substituting the transformer parameter scaling and loss bound establishes Theorem 1’s bound.Source: §4.3 · p. 13

5 Numerical results

• The evaluation spans Navier–Stokes, ERA5, QuickDraw, and Food-101, using Wasserstein-2, JSD, energy-spectrum metrics, or FID as appropriate.The baselines share a denoising backbone and training protocol, differing in how the prompt set is represented and injected.Source: §5 · pp. 13–14

5.1 2D Navier–Stokes equation

• ICL DM outperforms the baselines in Navier–Stokes in-distribution and out-of-distribution settings, except for MELRw in OOD(high), where Cond DM is lower.The comparison uses weighted MELR and JSD; OOD settings test viscosities outside the training intervals.Source: §5.1 · p. 14, Table 1 · p. 16

5.2 ERA5 dataset

• On ten ERA5 tasks, ICL DM outperforms the baselines on most evaluated metrics and generates samples with patterns similar to reference samples.Evaluation uses W2, weighted and unweighted MELR, and JSD with 100 prompt samples.Source: §5.2 · p. 16, Table 2 · p. 19

5.3 Google sketch dataset

• On QuickDraw, ICL DM outperforms the baselines across most metrics and settings, while performance degrades as K increases and task heterogeneity grows.The comparison covers K = 5, 10, and 100; the authors distinguish this growing task-family complexity from increasing task count under a fixed meta-distribution.Source: §5.3 · p. 19, Table 3 · p. 20

5.4 Food-101 Dataset

• On Food-101, ICL DM outperforms the baselines across most metrics and settings, while all methods generally degrade as K increases.This is an in-distribution evaluation; the authors attribute degradation to a broader task family under fixed per-task data, prompt size, and model capacity.Source: §5.4 · p. 21, Table 4 · p. 24

6 Conclusion

• The in-context score operator uses empirical context samples, a noisy query, and diffusion time to estimate task-dependent scores for reverse-diffusion sampling on unseen tasks without retraining or parameter updates.It is trained through amortized denoising score matching across task distributions.Source: §6 · p. 24
• A Gaussian-kernel conditional-mean representation yields an empirical Nadaraya–Watson estimator exactly representable by the constructed transformer class.The theory separates context-size, meta-training, within-task sampling, model-complexity, and data-geometry effects; finite-context error is governed by intrinsic rather than ambient dimension.Source: §6 · p. 24
• Additional within-task sampling cannot compensate for insufficient meta-training task coverage, making training on a sufficiently large and diverse distribution family important.The finite-context rate is governed by intrinsic task-manifold dimension, while ambient dimension enters through polynomial and logarithmic factors.Source: §6 · p. 24
• Lower-complexity Navier–Stokes tasks show strong in-distribution performance and modest out-of-distribution generalization; ERA5 remains effective in-distribution, while QuickDraw and Food-101 are more challenging.For the more complex families, strong in-distribution performance can also be difficult when prompts and training data are limited.Source: §6 · p. 24

B Proofs for Section 4

• The proof sequence first establishes the exact transformer realization in Lemma 1.The appendix then derives the statistical error bounds in Lemmas 2 and 3 and the approximation error bound in Proposition 1.Source: B Proofs for Section 4 · p. 27
• The statistical and approximation error results are combined to prove Theorem 1.Source: B Proofs for Section 4 · p. 27
• Table 5 provides notation for in-context learning of diffusion-model score operators.Source: Table 5 · p. 28

B.1 Proof of Lemma 1: Exact Transformer Realization

• A masked softmax readout is used to construct K_j(s), which the final decoding layer outputs.Source: §B.1 · p. 32
• The constructed transformer class represents the full vector-valued kernel regression estimator componentwise.Source: §B.1 · p. 32
• The construction uses n prompt tokens and one query token, with diffusion time providing access to α(t) and σ(t).Source: §B.1 · pp. 29, 33

B.2 Statistic Error I

• A covering-number bound is established for the restricted transformer loss class using coefficient perturbation sensitivity.A coefficient perturbation changes each output coordinate by at most Knetε uniformly over restricted inputs and time.Source: §B.2 · p. 34
• The loss decomposition separates task-level statistical error, episode-level statistical error, and time/noise discretization error.Source: §B.2 · p. 35
• For n,D ≥ 1, the bias-inclusive coefficient count satisfies DPT ≤3136nD^3, and the statistical rate is unchanged.Source: §B.2 · p. 36

B.3 Statistical Error II

• Conditional Hoeffding bounds at the episode and time–noise levels use KJ and KJNT, respectively, and combining the three bounds yields (53).The passage does not provide the full expressions for these bounds.Source: §B.3 · p. 36
• For each fixed cutoff, the explicit inequality gives the order stated in (37), with an implied constant that may depend on the fixed loss envelope.When s0 = 1, the passage says the same coarse order follows by taking ϵ = 1/2 and using K, J, NT ≥ 1.Source: §B.3 · p. 36
• When the cutoff varies with n, the proof retains its explicit dependence in (54).The cutoff used in Theorem 1 sets s0 = σ(δ) ∈ (0, 1).Source: §B.3 · p. 36

B.4 Approximation Error III

• The population conditional mean is written as the ratio of Gaussian-kernel expectations N(x(t)) and D(x(t)).The numerator averages the kernel-weighted clean sample, while the denominator averages the kernel weights.Source: §B.4 · p. 36
• The proof bounds the integral without a uniform lower bound on D(x(t)), using cancellation from the noisy-query density and localization to cells on the clean support.The partition argument needs no lower bound on any partition cell’s mass, and the cells are chosen with diameter at most σ(t)/α(t).Source: §B.4 · pp. 37–38
• On the event where the transformer candidate equals K, its error is controlled directly; outside that event, bounded coordinates control the error, and the α(t) = 0 case has zero score error.Averaging the resulting bound over task and time distributions proves (34).Source: §B.4 · p. 39

B.5 Proof of Theorem 1

• With σ(δ) = n^-1/(d+6), the proof obtains an n^-2/(d+6) term in the bound and absorbs fixed factors into its constants.The coefficient combines localization and tail contributions; the candidate uses the n + 1 active-token realization.Source: §B.5 · pp. 39–40

C Auxiliary Lemmas

• The auxiliary interaction lemma constructs attention that selects a token position and applies a ReLU inner-product operation; a decrementing lemma supplies a six-layer residual feed-forward construction.The interaction construction also permits a linear value map.Source: C Auxiliary Lemmas · p. 40

D.1 Network Architecture

• The implementations use a time-conditioned, multi-resolution U-Net, with ICL DM injecting prompt features through cross-attention and Cond DM and SNO using FiLM-style conditioning.Cond DM uses a learned prompt embedding, whereas SNO substitutes a projected kernel mean embedding; their remaining denoising architectures are identical.Source: §D.1 · p. 41

D.2 Training Details

• All models are trained with Adam at a learning rate of 1e-4 and a batch size of 32.Source: §D.2 · p. 41
• Table 6 specifies dataset-level settings, including meta-training tasks K, per-task training samples B, prompt samples n, evaluation generations m, and epochs.The caption does not provide the table’s cell values.Source: Table 6 · p. 42
• Sampling uses δ = 10^-4 and 1000 Euler–Maruyama steps to discretize the reverse-time SDE across all experiments.Source: §D.2 · p. 42

• Scalar-valued empirical PDFs are estimated with kernel density estimation, typically using a Gaussian kernel, while bandwidth h controls estimate smoothness.Source: §D.3 · p. 42
• JSD compares estimated PDFs of two datasets using the Jensen–Shannon divergence, defined with the Kullback–Leibler divergence.Source: §D.3 · p. 42
• Energy spectra summarize how scalar-field energy is distributed across Fourier modes, with E_f(k) giving energy at wave-number magnitude k.Source: §D.3 · p. 42
• MELR measures spectral alignment between empirical datasets, with uniform weights yielding an unweighted ratio and energy-proportional weights yielding a weighted ratio.Source: §D.3 · p. 42
• FID compares means and covariances of pretrained-network image features; FID-QD and FID-Food use dataset-pretrained ResNet-18 features for Quick, Draw! and Food-101.Source: §D.3 · p. 43