Summary
Diffusion models usually learn how to create data from a fixed type of example, making it hard for them to adapt to new kinds of data with only a few samples. The authors propose a method where a transformer model learns to estimate the necessary steps for generating new data by looking at a small set of examples from an unknown distribution. This lets the model generate new samples from unseen tasks without needing to be retrained. Their method uses attention mechanisms inspired by a known mathematical formula, allowing it to generalize across tasks. They tested it on scientific and image data, showing it performs well on simpler cases but has limitations as tasks get more complex.
What this means in practice
- •For climate data analysts: Generate detailed climate simulations for unseen weather scenarios by conditioning on a small set of recent data samples.
- •For digital artists: Create new sketches or food images by providing a few example drawings or photos, enabling fast style adaptation.
Authors
Wuzhe Xu, Zhaiming Shen, Yulong Lu, Wenjing Liao, Rongjie Lai
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.