Multi-task Neural Diffusion Processes

arXiv:2510.03419v3 Announce Type: replace-cross 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.

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.

Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Applications (stat.AP); Machine Learning (stat.ML)

arXiv:2510.03419 [cs.LG]

(or arXiv:2510.03419v3 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2510.03419 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.

aiqbench analysis

Multi-task Neural Diffusion Processes

Related coverage