The accompanying repository includes a Lean 4 formalization of the paper's Theorem 1, a theoretical result on the in-context score operator.
Abstract
Diffusion models typically learn the generative dynamics of a single target distribution or a family of distributions indexed by conditioning variables fixed during training. This limits their ability to adapt to a previously unseen distribution from only a small set of examples. We introduce an in-context solution operator learning framework for diffusion models, in which each target distribution is treated as a task drawn from a meta-distribution and is specified at inference time by an empirical context set. Our method learns an in-context score operator that maps the context samples, a noisy query point, and the diffusion time to the task-dependent score, enabling reverse-diffusion sampling on unseen tasks without retraining or parameter updates. A key observation is that the task-dependent score admits a Gaussian-kernel conditional-mean representation. Its empirical Nadaraya–Watson approximation can be represented exactly by a model in our constructed transformer class, providing a principled foundation for the proposed query-dependent attention architecture. The operator is trained through amortized denoising score matching across task distributions. We establish generalization-error bounds that quantify the effects of the context size, the numbers of training tasks and within-task samples, model complexity, and intrinsic and ambient dimensions. Numerically, we evaluate the method on a range of scientific and visual datasets, including Navier-Stokes, ERA5 climate data, QuickDraw sketches, and Food-101. The results show that the proposed method performs well on relatively simple task families, while its performance degrades as task complexity increases, thereby illustrating both the promise and the limitations of in-context learning for diffusion models.
Problem
Standard and conditional diffusion models learn a single target distribution, or a family indexed by conditioning variables fixed at training time. They cannot adapt to an unseen distribution given only a small set of example samples without retraining.
Approach
Each target distribution is treated as a task drawn from a meta-distribution and specified at inference time by an empirical context set. A transformer-based in-context score operator maps the context samples, a noisy query and the diffusion time to the task-dependent score. It is trained by amortized denoising score matching across tasks. The architecture is motivated by a Gaussian-kernel conditional-mean representation of the score, whose Nadaraya–Watson approximation is exactly representable in the constructed transformer class. Generalization-error bounds are derived, and Theorem 1 is formalized in Lean 4 in the repository.
Results
The bounds separate the effects of context size, number of tasks and episodes, model complexity, and intrinsic versus ambient dimension. Experiments on Navier–Stokes, ERA5, QuickDraw and Food-101 show good in-distribution and modest out-of-distribution performance on simpler task families. Performance degrades on more complex visual task families.
(b) \nu=2.2\times 10^{-3}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