Multi-task Neural Diffusion Processes

Joseph Rawson; Domniki Ladopoulou; Petros Dellaportas

cs.LGOct 3, 2025

Multi-task Neural Diffusion Processes

Authors: Joseph Rawson, Domniki Ladopoulou, Petros Dellaportas

Organizations: Department of Statistical Science, University College London, Gower Street, London, WC1E 6BT, United Kingdom. · Department of Statistics, Athens University of Economics and Business, Patission 76, Athens, 104 34, Greece.

Abstract

Neural diffusion processes provide a scalable, non-Gaussian approach to modelling distributions over functions, but existing formulations are limited to single-task inference and do not capture dependencies across related tasks. We propose multi-task neural diffusion processes (MT-NDPs), an extension that incorporates a task encoder to enable task-conditioned probabilistic regression and few-shot adaptation across related functions. The task encoder extracts a low-dimensional representation from context observations and conditions the diffusion model on this representation, allowing information sharing across tasks while accommodating variable-sized context sets and preserving permutation invariance with respect to the ordering of context observations. We evaluate MT-NDPs through synthetic experiments and two real-world applications involving wind turbine data and school examination outcomes. The results demonstrate improved predictive accuracy relative to single-task diffusion baselines when context observations are available, alongside competitive performance against neural-process and Gaussian-process alternatives. In the wind-energy application, MT-NDPs achieve the lowest average coverage error among the evaluated models at larger context sizes, demonstrating improved predictive uncertainty calibration. Overall, our findings highlight the potential of task-conditioned diffusion models for uncertainty-aware multi-task regression and few-shot adaptation to previously unseen tasks.

Figures & tables

Figure 1: Illustration of the forward diffusion process. Gaussian noise is added incrementally at each timestep under a cosine variance schedule, progressively degrading the structure of the data until it converges to white noise. GP: Gaussian process, CI: confidence interval.

Figure 2: Comparison of unconditional and conditional reverse diffusion sampling. We sampled 10 reverse trajectories per input location, plotting the mean of these. Conditional sampling provides a closer fit to the true data across the illustrated timesteps. GP: Gaussian process.

Figure 3: Principal component (PC) projection of encoder embeddings for turbines $1$ – $6$ , trained on turbines $2$ – $6$ . Each panel corresponds to a different range of context set sizes. Colours denote turbines, shaded ellipses show $68$ % coverage regions, and crosses mark cluster means. As context increases, the encoder begins to separate turbines, with turbine $6$ emerging as the most distinct.

Figure 4: Mean absolute error (MAE, left) and root mean squared error (RMSE, right) distributions for the one-dimensional single-task neural diffusion process (1D ST-NDP), five-dimensional single-task neural diffusion process (5D ST-NDP), five-dimensional multi-task neural diffusion process (5D MT-NDP), multi-task Gaussian process (5D MT-GP), attentive neural process (5D ANP), and multi-task neural process (5D MT-NP) on turbine 6 across context sizes $0$ , $25$ , and $50$ , aggregated over $30$ test functions of $100$ observations each. MT-GP results are not reported for context size equal to $0$ , as the implementation requires task-specific observations to condition predictions for an unseen task.

Figure 5: Coverage error (CE) distributions for the one-dimensional single-task neural diffusion process (1D ST-NDP), five-dimensional single-task neural diffusion process (5D ST-NDP), five-dimensional multi-task neural diffusion process (5D MT-NDP), multi-task Gaussian process (5D MT-GP), attentive neural process (5D ANP), and multi-task neural process (5D MT-NP) on turbine 6 across context sizes $0$ , $25$ , and $50$ , aggregated over $30$ test functions of $100$ observations each. MT-GP results are not reported for context size equal to $0$ , as the implementation requires task-specific observations to condition predictions for an unseen task.

	Context = $0$			Context = $25$			Context = $50$
Model	MAE	RMSE	CE	MAE	RMSE	CE	MAE	RMSE	CE
1D ST-NDP	$66.539$	$90.491$	$13.052$	$56.641$	$77.766$	$8.779$	$53.283$	$73.429$	$6.865$
5D ST-NDP	$55.169$	$75.441$	$11.225$	$51.864$	$71.306$	$9.058$	$49.676$	$68.465$	$7.376$
5D MT-GP	–	–	–	$61.120$	$86.372$	$11.184$	$55.906$	$79.639$	$11.480$
5D MT-NP	$62.017$	$79.826$	$15.536$	$48.974$	$64.387$	$9.332$	$47.942$	$63.269$	$8.717$
5D ANP	$66.020$	$85.680$	$16.780$	$49.300$	$64.610$	$11.950$	$\mathbf{46.860}$	$\mathbf{61.800}$	$11.230$

Table 1: Out-of-sample performance of single-task neural diffusion process (ST-NDP), multi-task neural process (MT-NP), multi-task Gaussian process (MT-GP), attentive neural process (ANP) and multi-task neural diffusion process (MT-NDP) models on turbine

6

, averaged over

30

test functions (

100

points each). Metrics reported: mean absolute error (MAE, kW), root mean squared error (RMSE, kW), and coverage error (CE, %). CE denotes the mean absolute coverage error across the evaluated nominal coverage levels, expressed in percentage points. Bold values indicate the best overall performance across all models and context sizes. MT-GP results are not reported for context size equal to

0

, as the implementation requires task-specific observations to condition predictions for an unseen task. 1D: one-dimensional, 5D: five-dimensional.

Figure 6: Synthetic ablation study across embedding dimensions $\kappa$ and context sizes $C$ . Panels show (a) mean absolute error (MAE), (b) root mean squared error (RMSE), (c) empirical coverage of nominal 95% predictive intervals, and (d) inference time in minutes. Each boxplot represents ten evaluations. The dotted line in panel (c) indicates nominal 95% coverage.

Model	MAE	RMSE	Coverage 95%	CE (%)	Time (s)
MT-NDP	0.564 (0.069)	0.715 (0.080)	95.31 (3.22)	4.940 (2.616)	525.99 (19.23)
ST-NDP	0.634 (0.120)	0.790 (0.133)	91.76 (4.00)	6.377 (3.657)	333.84 (52.31)
MT-NP	0.562 (0.067)	0.707 (0.071)	92.64 (4.01)	4.161 (1.397)	0.17 (0.19)
ANP	0.579 (0.058)	0.721 (0.071)	93.81 (3.20)	4.022 (1.865)	0.80 (1.69)
MT-GP	0.599 (0.047)	0.742 (0.052)	93.26 (4.08)	4.032 (2.094)	1.76 (1.31)

Table 2: Predictive performance on the School Examination dataset using 30 schools for training, five for validation, and six previously unseen schools for testing. Each test school is evaluated using

C=20

context observations. Results are reported as mean (standard deviation) across the six test schools. MAE denotes mean absolute error, RMSE denotes root mean squared error, Coverage 95% refers to the empirical coverage of nominal 95% predictive intervals, and CE denotes calibration error. Time (s) represents the inference time in seconds per test school. Lower values of MAE, RMSE, CE, and inference time indicate better performance, while empirical coverage closer to the nominal 95% level indicates better interval calibration. Bold values indicate the best performance for each metric.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: Forward and reverse trajectories for a single input point under (a) conditional and (b) unconditional sampling. In the forward process (blue), Gaussian noise is gradually added. In the reverse process (red), noise is progressively removed. The shaded area shows one standard deviation.

Panel of Figure 8 — Figure 8: SCADA data for the six turbines from the Kelmarsh wind farm Plumley (2022) after removing standbys and warnings using the operational status and events file. (a) Wind power curves showing mean power output (kW) by mean wind speed (m/s). (b) Joint distribution of mean wind speed (m/s) and mean wind direction (°) for turbine 1 (left), and mean power output (kW) by mean wind speed (m/s) and wind direction (°) for turbine 1 (right).

Panel of Figure 9 — Figure 9: SCADA data for the six turbines from the Kelmarsh wind farm Plumley (2022) after removing standbys and warnings using the operational status and events file. (a) Mean power output (kW) by mean nacelle temperature (°C). (b) Mean power output (kW) by mean transformer temperature (°C).

Figure 10: Fit of Gaussian process (GP) (left) and neural diffusion process (NDP) (right) on a test sample of $100$ points from turbine $1$ , with $95$ % confidence intervals (CI) on normalised wind speed and power output. 1D: One-dimensional.

Panel of Figure 11 — Figure 11: Comparison of Gaussian process (GP) and neural diffusion process (NDP) on turbine $1$ test data. (a) Point prediction accuracy (MAE and RMSE distributions). (b) Predictive uncertainty calibration (coverage probabilities). 1D: One-dimensional, ST: single-task

Figure 12: Distributions of mean absolute error (MAE, left) and root mean squared error (RMSE, right) for the one-dimensional (1D) and five-dimensional (5D) neural diffusion process (NDP) models, evaluated with $0$ , $25$ , and $50$ context points. Results are aggregated over $30$ test functions of $100$ points each from turbine $1$ . ST: single-task.

Figure 13: Observed vs. theoretical coverage probabilities for the one-dimensional (1D) and five-dimensional (5D) neural diffusion process (NDP) models with $0$ , $25$ , and $50$ context points. Results are averaged over $30$ test functions of $100$ points each from turbine $1$ . ST: single-task.

Figure 14: Coverage error for the one-dimensional single-task neural diffusion process (1D ST NDP) and five-dimensional single-task neural diffusion process (5D ST NDP) models with $0$ , $25$ , and $50$ context points. Results are computed over $30$ test functions of $100$ points each from turbine $1$ .

	Context = $0$			Context = $25$			Context = $50$
Model	MAE	RMSE	CE	MAE	RMSE	CE	MAE	RMSE	CE
1D GP	$51.636$	$75.596$	$18.075$	–	–	–	–	–	–
1D ST-NDP	$52.603$	$75.139$	$4.584$	$51.273$	$74.568$	$4.181$	$50.545$	$74.137$	$3.938$
3D ST-NDP	$60.967$	$88.096$	$9.731$	$52.338$	$74.312$	$6.750$	$50.855$	$73.537$	$5.369$
5D ST-NDP	$49.429$	$67.685$	$8.351$	$46.962$	$66.007$	$6.029$	$45.791$	$65.301$	$4.816$

Appendix

Table 3: Out-of-sample performance of Gaussian process (GP) and single-task neural diffusion process (ST-NDP) models on turbine

1

, averaged over

30

test functions (

100

points each). Metrics reported: mean absolute error (MAE, kW), root mean squared error (RMSE, kW), and coverage error (CE, %). 1D: one-dimensional, 3D: three-dimensional, 5D: five-dimensional.

Figure 15: Empirical versus nominal coverage probabilities for the one-dimensional ST-NDP, five-dimensional ST-NDP, and five-dimensional MT-NDP on turbine 6 across context sizes $0$ , $25$ , and $50$ .

Figure 16: Synthetic multi-task regression dataset comprising 20 training tasks (grey) and five previously unseen test tasks (coloured). The functions are generated using a common data-generating process, with task-specific variation to represent differences between related regression tasks. The five held-out tasks are used to evaluate the generalisation performance of MT-NDP.

$\kappa$	$C$	MAE	RMSE	95% Coverage (%)	Time (min)
4	0	0.433 (0.078)	0.497 (0.094)	67.8 (3.8)	2.98 (0.12)
4	5	0.253 (0.078)	0.309 (0.093)	75.9 (4.0)	8.25 (0.19)
4	10	0.194 (0.058)	0.242 (0.074)	82.0 (3.6)	7.80 (0.20)
4	25	0.151 (0.034)	0.190 (0.046)	86.7 (3.2)	8.31 (0.52)
8	0	0.424 (0.093)	0.491 (0.107)	72.1 (3.9)	3.86 (0.11)
8	5	0.259 (0.062)	0.307 (0.069)	79.9 (2.7)	9.71 (0.26)

Appendix

Table 4: Results of the synthetic ablation study evaluating the effect of task embedding dimension (

\kappa

) and context size (

C

) on the predictive performance of MT-NDP. The model is trained on 20 synthetic tasks and evaluated on five previously unseen tasks. Performance is assessed using mean absolute error (MAE), root mean squared error (RMSE), empirical 95% coverage, and inference time in minutes. Results are reported as mean (standard deviation) across five test tasks and two repetitions.

Model	MAE	RMSE	Coverage 95% (%)	Time (min)
MT-NDP	0.805 (0.186)	1.008 (0.190)	95.148 (4.576)	13.124 (0.459)

Appendix

Table 5: Predictive performance of MT-NDP on the School Examination dataset under leave-one-school-out evaluation. Six schools are considered, with five used for training and the remaining school held out for testing in each evaluation. Each held-out school is evaluated using

C=20

context observations. Results are reported as mean (standard deviation) across the six evaluations. MAE denotes mean absolute error, RMSE denotes root mean squared error, Coverage 95% refers to the empirical coverage of nominal 95% predictive intervals, and Time (min) represents inference time in minutes per held-out-school evaluation.

Diffusion Models

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

See all 1,163 papers

Transformer-based In-Context Score Operator Learning for Diffusion Models
Oct 9, 2026Wuzhe Xu, Zhaiming Shen, Yulong Lu +2Diffusion Models In-Context Learning
South Dakota State University · Georgia Institute of Technology · University of Minnesota +1
DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models
May 14, 2026Quanhao Li, Junqiu Yu, Kaixun Jiang +7Diffusion Model Distillation Multi-Task RL
Fudan University, China · Alibaba Group, China
Diffusion Transformers are Provably Optimal In-context Generators
Oct 4, 2026Guoji Fu, Tomoya Wakayama, Ryotaro Kawata +3Diffusion Transformer Diffusion Sampling
National University of Singapore · RIKEN AIP · The University of Tokyo +1

Oct 9, 2026J/K move · Enter open · S save

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

Wuzhe Xu, Zhaiming Shen, Yulong Lu, Wenjing Liao+1

South Dakota State University · Georgia Institute of Technology · University of Minnesota+1

Diffusion Models In-Context Learning Operator Learning Diffusion Transformer

Open paper AskPDF arXiv

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.

Abstract

Diffusion Models

Explore similar work

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

DiffusionOPD: A Unified Perspective of On-Policy Distillation in Diffusion Models

Diffusion Transformers are Provably Optimal In-context Generators