Multitask Learning with Stochastic Interpolants

TL;DR

This paper introduces a generalized stochastic interpolant framework for learning mappings between probability distributions across multiple spaces, enabling versatile, task-agnostic generative models with zero-shot capabilities.

Contribution

It extends stochastic interpolants by incorporating operators, unifying and enhancing generative modeling across various tasks without task-specific training.

Findings

Effective zero-shot conditional generation and inpainting

Versatile in multiscale modeling and posterior sampling

Unifies existing generative models under a common framework

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 theoretical 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.