Multi-task Neural Diffusion Processes

arXiv:2510.03419v2 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. In many multi-task regression settings, jointly modelling correlated functions and enabling task-aware conditioning is crucial for improving predictive performance and uncertainty calibration, particularly in low-data regimes. We propose multi-task neural diffusion processes, 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 preserving input-size agnosticity and the equivariance properties of neural diffusion processes. The resulting framework retains the expressiveness and scalability of neural diffusion processes while enabling efficient transfer to unseen tasks. Empirical results demonstrate improved point prediction accuracy and better-calibrated predictive uncertainty compared to single-task neural diffusion processes and Gaussian process baselines. We validate the approach on real wind farm data appropriate for wind power prediction. In this high-impact application, reliable uncertainty quantification directly supports operational decision-making in wind farm management, illustrating effective few-shot adaptation in a challenging real-world multi-task regression setting.