Research map / Image generation / Fast diffusion sampling
Multitask Learning with Stochastic Interpolants
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Generalizes stochastic interpolants by replacing the scalar time with operators (matrices/linear maps), learning drifts once so that inpainting, posterior sampling, multichannel denoising, and path optimization can be done zero-shot at inference.
Problem
Generative models are typically trained for a single predetermined task (e.g., unconditional generation), so repurposing them for inpainting, denoising, posterior sampling or other tasks needs retraining or approximate guidance.
Approach
Replace the scalar time in stochastic interpolants with a family of bounded linear operators (alpha, beta) and learn operator-conditioned drifts once; any operator path then transports samples, so different tasks correspond to different operator paths at inference without retraining. Includes Hadamard (diagonal) interpolants for masking and a closed-form relation between prior and posterior drifts
Key results
- On CelebA random-mask inpainting, the method reaches 33.76 PSNR / 0.967 SSIM, the best among compared methods.
- On AFHQ-Cat random-mask inpainting, 33.11 PSNR / 0.945 SSIM, the best reported; block-mask 26.96 PSNR / 0.914 SSIM, also best.
- On CelebA block-mask inpainting the method reaches 29.98 PSNR, behind OT-ODE (30.59) but ranks second.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Ours instead of the strongest baseline | CelebA inpainting · CelebA Random | PSNR: 33.54 -> 33.76 | Helped |
| Use Ours instead of the strongest baseline | CelebA inpainting · CelebA Random | SSIM: 0.953 -> 0.967 | Helped |
| Use Ours instead of the strongest baseline | CelebA inpainting · CelebA Block | PSNR: 30.59 -> 29.98 | Hurt |
| Use Ours instead of the strongest baseline | CelebA inpainting · CelebA Block | SSIM: 0.943 -> 0.938 | Hurt |
| Use Ours instead of the strongest baseline | AFHQ-Cat inpainting · AFHQ-Cat Random | PSNR: 32.98 -> 33.11 | Helped |
| Use Ours instead of the strongest baseline | AFHQ-Cat inpainting · AFHQ-Cat Random | SSIM: 0.930 -> 0.945 | Helped |
| Use Ours instead of the strongest baseline | AFHQ-Cat inpainting · AFHQ-Cat Block | PSNR: 26.87 -> 26.96 | Helped |
| Use Ours instead of the strongest baseline | AFHQ-Cat inpainting · AFHQ-Cat Block | SSIM: 0.904 -> 0.914 | Helped |
Limitations
- Learning the larger operator-indexed space increases training complexity and cost.
- Over-smoothing occasionally appears in inpainting results.
- Some closed-form posterior formulas require restrictions (e.g., k<1 for the quadratic reward case).
Experimental setup
Datasets: MNIST CelebA AFHQ-Cat phi4 lattice MCMC samples maze trajectories
Benchmarks: CelebA inpainting AFHQ-Cat inpainting
Metrics: PSNR SSIM magnetization
Baselines: PnP-Flow Flow-Priors D-Flow OT-ODE
Builds on: stochastic interpolants flow matching rectified flow PnP-Flow
Closest papers by meaning
- Stochastic Interpolants with Data-Dependent Couplings (ICML 2024 · similarity 0.89)
- Multimarginal Generative Modeling with Stochastic Interpolants (ICLR 2024 · similarity 0.89)
- Stochastic Interpolants via Conditional Dependent Coupling (NeurIPS 2026 · similarity 0.88)
- A solvable model of learning generative diffusion: theory and insights (NeurIPS 2025 · similarity 0.88)
- Diffusion Generative Flow Samplers: Improving learning signals through partial trajectory optimization (ICLR 2024 · similarity 0.88)
- Denoising Task Difficulty-based Curriculum for Training Diffusion Models (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.