Multitask Learning with Stochastic Interpolants

Hugo Negrel1, Florentin Coeurdoux1, Michael Albergo2, Eric Vanden-Eijnden1, 3,
1Capital Fund Management 2Society of Fellows, Harvard University 3Courant Institute of Mathematical Sciences, New York University, New York
Paper arXiv Soon

TL;DR This paper extends stochastic interpolants formalism to solve a very large variety of tasks with one single trained model. It shows strong performance on high-dimensional datasets such as CelebA and \(\varphi^4\). Experimental proofs include exact posterior sampling, inpainting and highly-constrained planning.

Abstract

We propose a framework for learning maps between probability distributions that broadly generalizes the time dynamics of flow and diffusion models.

To enable this, we generalize stochastic interpolants by replacing the scalar time variable with vectors, matrices, or linear operators, allowing us to bridge probability distributions across multiple dimensional spaces.

This approach enables the construction of versatile generative models capable of fulfilling multiple tasks without task-specific training. Our operator-based interpolants not only provide a unifying theoreti cal perspective for existing generative models but also extend their capabilities. Through numerical experiments, we demonstrate the zero-shot efficacy of our method on conditional generation and inpainting, fine-tuning and posterior sampling and multiscale modeling, suggesting its potential as a generic task-agnostic alternative to specialized models.

Inpainting

Classically, interpolation between two multivariates random variables \(x_0\) and \(x_1\) is parametrized by two scalars \(\alpha\) and \(\beta\) via the very simple relation \(I_{\alpha, \beta}(x_0,x_1) = \alpha x_0 + \beta x_1\), where \(x_0 \sim \rho_0\), \(x_1 \sim \rho_1\), and \(\alpha, \beta \in [0, 1]\). By their nature of scalar parameters, they acts uniformally on every entries, letting little flexibility in the diffusion process. Now, define \(I_{\alpha, \beta}(x_0, x_1) = \alpha \odot x_0 + \beta \odot x_1\) with vector \(\alpha \in [0,1]^{n}\) and \(\beta \in [0,1]^n\). While it significantly increases the size of input space, pixel-wise inpainting becomes now a trivial task.

Dataset: AFHQ Cat

Original image

Block masking

Random masking

Dataset: CelebA

Original image

Block masking

Random masking

Dataset: \(\varphi^4\)

Original image

Block masking

Constrained planning

This problem takes root in Reinforcement Learning (RL). Take a maze and two end points. The task is to find the shortest path joining them while abiding by the contraints imposed in the environment. With stochastic interpolants, to generate a path, you simply fix the first and last point, and eventually let diffuse all the other intermediate points.

The constraints do not necessarily have to apply only on the end points, you can also impose to the path to take the detour you want.

Shortest path

Path under a constraint at mid-length

Posterior sampling

With the very same model, you can perform exact posterior sampling. Take here the example of \(\varphi^4\) model, coming from quantum field theory, which famously exhibit a phase transition phenomena. Simply by shifting and scaling interpolation coefficients \(\alpha\) and \(\beta\), you can sample from a posterior without retraining or any kind of approximation.

Dataset: \(\varphi^4\)

BibTeX

@article{negrel2025multitasklearningstochasticinterpolants,
        author    = {Negrel, Hugo and Coeurdoux, Florentin and Albergo, Michael and Vanden-Eijnden, Eric},
        title     = {Multitask Learning with Stochastic Interpolants},
        journal   = {NeurIPS2025},
        year      = {2025},
        url       = {https://arxiv.org/abs/2508.04605}
                }