Multitask Learning  with Stochastic Interpolants
Authors: Hugo Negrel
Florentin Coeurdoux
Michael Samuel Albergo
Eric Vanden-Eijnden
Venue: NeurIPS 2025 spotlight ()
Keywords: generative models
flows
diffusions
conditional generation
pinpointing
fine-tuning
multiscale generation
zero-shot
multitask
Primary area: deep_learning
TL;DR: We introduce operator-based interpolants—a novel framework that replaces time variables with linear operators in generative models—enabling a single model to perform multiple tasks without task-specific training.

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.

Multitask Learning with Stochastic Interpolants
Hugo Negrel
Florentin Coeurdoux
Capital Fund Management 23 Rue de l'Université, 75007 Paris hugonegrel13dgman fr
Capital Fund Management 23 Rue de l'Université, 75007 Paris f florentin. cow.rdouxxcc...mmmmmmmm com
Michael S Albergo Society of Fellows, Harvard University malbergo@fsn as h.w.ardd edu
Eric Vanden- Eijnden Machine Learning Lab Capital Fund Management 23 Rue de I'Université, 75007 Paris eric. vanden- jwwennecciideeemmmmmmmmm com
Abstract
We propose 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 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 sam- pling, and multiscale modeling, suggesting its potential as a generic agnostic alternative to specialized models.
only provide unifying theoreti-
Introduction
Recent years have witnessed remarkable advances in generative modeling, with transport-basee
approaches such normalizing flows (Lipman et al., 2022 Albergo and Vanden-Eijnden 2022
Liu et 2022) and diffusion models (Ho et al.) 2020 Song and Ermon [2020] De Bortoli al.
2021: Albergo et al.) (023a) emerging as the techniques across various application domains (Rombach et al. 2022 Mazé and Ahmed, 2023 Alverson et al., 2024). These methods have revolutionized our ability to generate high-quality images, text, and other complex data types by viewing these data as samples from an unknown target distribution and learning transform simple (e.g., noise) distributions into this target. This transformation is effectively achieved via transport of the samples by a flow or diffusion process with a drift (or score) parameterized by neural networks and estimated via simulation-free quadratic regression, enabling highly efficient training.
Despite their impressive performance, these generative frameworks face fundamental limitation: they are typically designed and trained for specific, predetermined tasks, with the generative objective specified before training. For example, diffusion model trained generate images cannot easily be repurposed to perform inpainting or other editing tasks without substantial modification or retraining.
also at: Courant Institute of Mathematical Sciences, New York University, New York, NY 10012, USA,
eve20cims. nyu. edu
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

While some flexibility can be achieved through conditioning variables or prompting, these approaches remain constrained within narrowly defined operational boundaries established beforehand. Recent attempts at multitask generation using approximated guidance strategies (Chung et al. 2023 Song
jtt al. 2022. Wang et al. 2024) have shown promise, but rely on uncontrolled approximations that limit their theoretical guarantees and can lead to unpredictable results. These methods typically operate within predefined space of capabilities and lack the flexibility to adapt to novel tasks without retraining, often requiring domain-specific architectural modifications specialized training procedures that further limit their versatility.
In this paper, we introduce novel framework for training truly multi- tasti-task generative models based on a generalized formulation of stochastic interpolants. Our key insight is replace the scalar time variable traditionally used in transport-based models with a linear operator. These operator-based interpolants enable interpolation between random variables across multiple dimensional planes or setups, providing a unified mathematical formulation that treats various generative tasks as different ways of traversing the same underlying space, rather than as separate problems requiring distinct models. This dramatically expands the space of possible tasks that single model can perform.
Our main contributions include theoretical advances that establish framework for multiple genera- tive applications:
We extend traditional scalar interpolation in dynamical generative models to higher-dimensional structures, developing a unified mathematical formulation of operator-based interpolants that treats various generative tasks as different ways of traversing the same underlying space. This opens up fundamentally new ways of seeing how generative models can be structured to handle multiple objectives simultaneously.
We show how this framework enables generative models with continual self-supervision over wide purview of generative tasks, making possible: universal inpainting models that work with arbitrary masks, multichannel data denoisers with operators in the Fourier domain, posterior sampling with quadratic rewards, and test-time dynamical optimization with rewards and interactive user feedback, all with one pretrained model.
We demonstrate these various tools on image-infiling, data de-corruption, statistical physics simulation, and dynamical robotic planning tasks across a number of datasets, showing that our method matches or surpasses existing approaches without being specifically tied to any single generative objective.
All common augmentations like text-conditioning and guidance can still be used just as before in our setup. While our approach increases the complexity of the initial learning problem, we provide arguments that this additional complexity can be addressed through scale. In essence, the "pretraining" phase becomes more challenging, but the resulting model gains substantial flexibility and versatility that compensates for the pretraining costs. That is, our approach can be seen as way of amortizing learning over a variety of tasks. This points toward more general paradigm of universal generative models that can be trained once and then applied to variety of objectives, potentially reducing computational and environmental costs associated with training separate models for each task.
1.1 Related works
Flow matching and diffusion models. Our approach extends the theoretical groundwork estab- lished in flow matching with stochastic interpolant and rectified flows (Lipman et al., 2022 Albergo and Vanden Vanden-Eijnden 2022 Albergo et al.) 2023a Liu et al., 2022) as well as the probability flow formulations in diffusion models (Ho et al. 2020: Song al. 2020). which established techniques for handling multiple target distributions simultaneously. Unlike data-dependent coupling approaches (Albergo et al., 2023b) that require constructing specific couplings for tasks like inpainting, our method learns general operator space that naturally accommodates such tasks without additional cou- pling design. By introducing operator-valued interpolants, we enable a richer space of transformations between distributions, unlocking a flexible framework for multiple generative tasks.
Inverse problems and inpainting Our framework offers unified approach to inverse problems, contrasting with traditional methods that require problem-specine variational optimization procedures Pereyra al. (2015). Recent diffusion-based approaches Chung et al. (2023); Song et al. (2022):
Kawar et al. (022); Martin et al. (2025) typically need guided sampling trajectories tailored to

---

each task. Similarly, methods using MCMC/SMC sampling Coeurdoux et al. (2024): Sun et al. (2024), variational approximations Mardani et (023); Alkan et al. (2023), or optimization-diiv techniques Wang et al. 2 Daras et al. (2022) remain fundamentally task-specific. Our approach encodes solution paths within the interpolant operator structure itself, enabling multiple inverse problems to be addressed through appropriate operator path selection during seerenncewiihhuut additional training. Our work can also be seen as a way to formalize methods that seek to clean data corrupted in various ways Bansal et al. (2022)
Multiscale and any-order generation Recent approaches to generative modeling have explored hicrarchical strategies through progressive refinement. Visual Autoregressive Modeling (Tian et al., 2024) employs next-scale prediction, while Fractal Generative Models (Li et al. 2025) utilize self- similar structures for multiscale representations. These methods typically constrain generation to fixed paths established during training. In contrast, our framework decouples the training process from generation trajectories, allowing flexible path selection at inference time. This is conceptually related to optimizing generation order in discrete diffusion (Shi et al.) 2025) and token ordering studies (Kim et al. 2023), but provides greater flexibility by enabling post- training optimization and dynamic, self-guided generation strategies that adapt based on intermediate results.
2 Theoretical framework
Imagine we want to create a single generative model capable of multiple tasks sampling new data, inpainting, denoising, and more. To achieve this, we need to expand beyond the traditional "single path" between noise and data. In this section, we develop the theoretical foundations of operator-based interpolants, which allow flexible navigation through richer multidimensional space and will enable multitask generative capabilities discussed in Section 3
2.1 Operator-based interpolants
Suppose that we are given a couple of random variables (No, x1) both taking values in a Hilbert space H (for example Re4) and drawn from a joint distribution y p(dro, dx1). Our aim is to design a transport between broad class of distributions supported on H involving mixtures of zo and æ1. We will do so by generalizing the framework of stochastic interpolant.
Definition 2.1 (Operator- based interpolants). Let B(H) be connected of bounded linear operators on H and let S C B(H) x BHH), also connected. Given any pair of linear operators (a, B) E S, the operator-basek interpolant 1 [(a, () is the stochastic process given by
I(a, 3) azg + 821,
(a,8)
(1)
where (Io, x1) ~ μ.. We will denote by Pa.3 the probability distribution of I [(a, 3).
If we picked e.g. = (1-t) f) and standard stochastic interpolant, but we stress that the operator- based interpolants from Definition 2. are much more general objects. For example, H Rdd we could take B(H) to set of d x matrices with real entries choices of a. 3 tailored to several multitask generation will be discussed in Section 3 The main objective of this paper is to show how to exploit this flexibility of design. Specifically we will show that we can learn model that can be used to transport samples of [(xx, 3) along any paths (a, in S without having to choose any such path during training. This will mean that transport problems in broad class associated to variety of tasks, including inpainting and block generation of any order, finettunig, etc. will be pretrained into this model.
ld with E [0, 1], respectively, we would go back to a
2.2 Multipurpose drifts and score
To proceed, we introduce two drifts, which are functions of S x H. taking value in Ht: Definition 2.2 (Multipurpose drift). The drifts no. 77 H - are given by
(a,a, 3,x) E[zo][((e B) z). m1 8 x) E[x; [(a, 3) z).
(2)
3

---

where EE[ [T 3) x denotes expectation over the coupling (20, ~ conditioned on the event (a, 5) z.
Using the L characterization of the conditional expectation, these drifts can be estimated via solution of tractable optimization problem with an objective function involving an expectation:
Lemma 2.3 (Drift objective). Let v da, d3) be probability distribution whose support is S. Then the drifts 1Jo,1 (ax, 3, x) in Definition 2.2 can be characterized globally for all (a, 3) E S and all
E supp(Ma.,)) via solution of the optimization problems
T/o argmin (a,8)>w [lloo(. 3,, 9)) -rool? -rol?
Yo
(10
711 argmin E (a,
(1, 3, 8)) xx1l²)
(4)
"1 (20:21)~
where || denotes the norm in H.
This lemma is proven in Appendix[A] Below will use (3 and (4) to learn 1o, over rich parametric class made of deep neural networks. Note that the drifts and 771 are not linearly independent since the definition of (a, 33 in (9) together with the equality x E[ ()) = x] imply that x = =07o 3.x)) + 3m 3,x).
(5) Therefore can obtain no from 771 if is invertible and 7/1 from no if 3 is invertible. Note also that, when Zo is Gaussian and øy 1 Z1. 7oo is related to the score of the distribution of the stochastic interpolant:
Lemma 2.4 (Score), Assume that H R& and that the probability distribution pa.3 of the stochastic interpolant M(, 9) is absolutely continuous with respect to the Lebesgue measure with density Pa. Assume also that zω NO, Id) and xo I 21. Then the score sa,p(z)
log (Z) is related to the drift no(a, 3,z) via
no(a, 3, x) a (Ix).
(6)
This lemma is proven in Appendix A
2.3 Transport with flows and diffusions
We can state the main theoretical results of this paper: if we are able to sample the stochastic interpolant I [(o,, 3o) at a specific value (oo, 38) E S. then we can produce sample of [(,,, 3e) along any curve (or, 93)) E S by solving either probability flow ODE estimated the drifts 17.. (ax, 3,x)) x) from Definition 2.2 along this curve:
SDE, assuming have
Proposition 2.5 (Probability flow). Let (as,, 3e (0. be any one-parameter family of operators (at, E Assume that an Br are differentiable for all E [0, 1]. Then, for all E [0, 1), the law of I (a B,) is the same as the law of the solution X, to
X, = (0r, Bc. X.) + Bem (on, Br. Xt),
Xo d I(oo, Bo).
This proposition proven in Appendix A Similarly. for generation with an SDE, we have:
Proposition 2.6 (Diffusion). Assume that H RRd and the probability distribution pa, of the stochastic interpolant [(a, 3) is absolutely continuous with respect to the Lebesgue measure. Assume also that, in I [a, B), To Gaussian and Fo - ø1. Then, under the same conditions as in Proposition 2.5 if qu is invertible, for all t [0, and any 2 O, the law of I I(a,, Be) is the same as the law of the solution X; to
dx;
Br, X;)dt + (on, Br, X()dt+ + /2GdW, Xi d I(ao, 90). (8)
where Wr is Wiener process in IRd.
This proposition is also proven in Appendix A Note that the SDE g reduces to the ODE 7 if we set E, 0. Note that if a, is positive-definite which set the noise term in to 22 /WW, this allows to extend this SDE to paths along which we can have at 0.
use &, c as new diffusion coefficient,

---

Figure 1: Multi-task, self-supervised sampling: Schematic representation of various sub-tasks that are captured by the minimizer of learning objective using the Hadamard- producct interpolant in (2). generative task is chosen in a zcro-shot manner by specifying a function of time after training. This serves as continual self- supervision of what has been unmasked what remains. Top: dy chosen to generate pixels an autoregressive fashion. Mid- dle: du is chosen to sample along fractal morton order. Bottom: can be chosen to do zero-shot in- painting.
AR
Morton
In-paint
3
Multitask generation
In this section we discuss how to use the theoretical framework introduced in Section 2 to perform multiple generative tasks without having doing any retraining.
3.1 Self-supervised generation and inpainting
Inpainting- the task of filling in missing parts of an tradittoiai requires specialized training for each possible mask configuration. Our operator-based framework enables fundamentally different approach: a single model that can perform inpainting with any arbitrary mask, chosen at inference time, or can generate samples from scratch in an arbitrary ordering of the generation. This includes standard generation of all dimensions once, autoregressive generation dimension by dimension, blockwise fractal generation, and so forth. In particular, we may want to construct generative model that fills in missing entries from a sample 21 E Fd drawn from a data distribution Mr. We would like this model to be universal, in the sense that it can be used regardless of which entries are missing; their number and position can be arbitrary and changed post- post-training. allowing for flexible inpainting and editing (see Figure for an illustration). This approach creates a natural self- supervisoin mechanism, as the model continuously tracks which parts have been generated and which remain to be filled.
To perform this task, assume that in is drawn from the data distribution fa of interest and ann drawn independently from (0,, Id), so that N(0, Id) #1, set 3 = 1 a in the operator interpolant (.. and assume that diagonal matrix. With a slight abuse notations we can then identify the diagonal elements of the matrix a with a vector E Id and write (1) as
((a) = a c + (1 a) + I1,
(9)
where @ denotes the Hadamard (i.e. entrywise) product. The drift to learn in this case is
na, x) E no æ |ax @ + (1 a)) 0z1 x]]
(10)
for a E [0, 1]d, and this learning can be done via solution of
min E a [Il I(a (10 (0,1d)
20)11
(11)
2
Denoting 41 (x) of.
suppose that we observe of for the entries with é C {l..... .......... .)
and would like infer the missing entries with of {l...., , d} o. To perform this inpainting we can use the probability flow ODE (7) with a path a such that af if E o and aj = t i cf Note that this can be done for any choice of without retraining.
3.2 Multichannel denoising
Suppose that B1, Bz.... B, are deterministic corruption operators that can be applied to the data. For example B, could be a high-pass filter, B2 a motion blur, etc. and they could defined primarily in the Fourier representation of the data. Similarly, if zo N0, Id), let A1, Ag.... Am be operators
5

---

Gaussian
Figure 2: Multichannel denoising: Possible in- terpolations fulfilled by various choices opera- tors in (1).. We present two such examples in the form of Gaussian and motion blurring, realized interpolations defined in the Fourier domain.
Motion
that give some structure to this noise (e.g. some spatial correlation over the domain of the data). Set Au Bg Id and take
0;
by B,
(12)
i=0
where d;, b, are nonegative scalars each taking values in some range that includes and bn 1. With this choice we can find path (04,,8)ue(0.11 that bridges data 1 corrupted in any possible channel as >i im A,A0 + >i B, for some choice of (a1, . ,,. b2....) back to the clean data
via a path that bridges this choice of parameters to bg and ag a1 ... bb = 0. See Figure for an illustration of two possible corruption schemes.
3.3 Fine-tuming and posterior sampling
Suppose that we are given data from a distribution 141 (dz), and that we would like to generate samples
from p (dx) = - (x(()) pz (dx) where r R is reward function and Z = JN en(a) #r (dx) is normalization function, which is unknown to us but we assume finite in the context of Bayesian inference, μi plays the role prior distribution, r the likelihood, and ju[ the posterior distribution.
We will assume that the reward r is quadratic function, i.e.
rxx = 1 (x, Ax) + (6,x)
(13)
where a definite negative bilinear operator on H. H. and (., denotes the inner product on H. For simplicity we also assume that H R.d: the general case can treated similarly.
In this context, assume that we have learned the drifts 1/,11 associated with the interpolant
Ma, B) = arro + Bri,
To ~ (0, Id);
31 Mi,
To 1 x1
(14)
so that we can generate samples from the prior distribution. Our next result shows that this gives us access to the drifts n6,1 associated with the interpolant
(a, (3) ozo + 3x..
To
Id), of ~ pl. To 1 x.
(15)
involving data xi from the posterior distribution. Proposition 3.1. Let
no 8,x) = |[ (ax B) = x),
m (a, 2,x) = () =x],
(16)
be the drifts associated with the interpolant
14 and
vola, 3,x) ) = E 4, B) = x (a x, = 3) = x),
(17)
be the drifts associated with the interpolant 115)
If a and 3 are invertible, then
(a B, - 63 a, Br, xr) + a - -3,,' xx)
(18)
(a, B, x) = n (1 3,, as long as we can find pair (ar, 3r) that satisfies
(19)
8.. a, -T, - B, = 3T -T -T A
(20)

---

and 2y is given by
a,a 3. a
(21)
This proposition is proven in Appendix as corollary of Proposition that relates the probability distribution of I, to that of I. Proposition 3.1 offers way to sample the posterior distribution without retraining, by using the drifts [88 and (19) the ODE 7 or the SDE (8).
3.4 Inference adaption
Suppose that we have learned the drifts 70.1 in Definition 2.2 and wish to transport samples along a path (t) with fixed end points. We can leverage the flexibility of our formulation to perform inference adaptation, that is, optimize the path (ax, (5)) used during generation to achieve specific objectives, such as minimizing computational cost, maximizing sample quality, or satisfying user constraints. This can be done in two ways: (I) offline optimization, where pre-compute optimal paths for different scenarios using objectives like Wasserstein length minimization, and (2) online adaptation, where paths are dynamically adjusted during generation based intermediate results or user feedback.
In the case of offline optimization, we could for example optimize the Wasserstein length of the path. That is, if we want to bridge the distributions pa of [(ao 90) and Pas, of [ar, 31) via Haa. with (a,, (e) E S for all t E [0, 1] then the path that minimizes the Wasserstein length of the bridge distribution Pa solves
/ E Bllláenaaa,, Be, B 3() + Bem (a,, Be, Man, B()) )
min (c),
(22)
where the minimization is performed over paths (ar, B:)r = (ar, (e) such that (att, 3r) E S for all t [0, 1] with their end points (ao, Su) and (a1, 31) prescribed and fixed.
Algorithm 1: Multitask learner
Samples (xg, 21) ~ μú; choice of distribution v'daa, d3) and associated sampler. repeat
input:
Draw batch (ar) ,r, 04 B:) ~
Compute I, 0,220 + B:r
Compute i, = x7 Σ [ 3r, It) xfl12 I| (ac, 3i, It) - 111
Take gradient step L to update iuu and in
until converged;
output: Drifts îo and f/
Algorithm 2: Multitask generator
input: Drifts î.. m1: choice of path (ayB)) ke)oo.1 tailored to the generation task; data
I(00, Bo) = ag + f0r1; diffusion coefficient e, > O; time step 1/K K with E
initialize: X6- = I(ao. So):
for k- = O..... K do
set i nolaun Buhr Xgn): T malau BEn, Xin), and z4 N(0,
update X = X + h(dak - cxnara n + h32 + /2EunhZN.
end
output: Xi d I(a11 31) (approximately)
4 Algorithmic aspects
The algorithmic aspects our framework can be summarized few key steps. First, we define connected set S of (a, B) such that the ensemble of different tasks will want to perform correspond
7

---

to getting samples of I(a, 3) at some value of (av, Pb) E S and generating from them new data at another value of (a1. 31) S. Second, we specify measure v on S for the learning of the drifts Mo and ?/ defined in 2 Third, we learn these drifts via minimization of the objectives in (3) and (4). using the procedure outlined Algorithm Note that we can possibly simplify this algorithm, learning only one of the two drifts and obtaining the other through the relation ( Finally, given any pairs (axo Bo). (a1,31) E S, we use a path (at, (e) with æt 3r for all t [o, 1/ and integrate the SDE (8) (or possibly the ODE (7 if we set 6, 0) to perform the generation, as outlined in Algorithm| 2 Note that this path could also be adapted on-the- during inference, using some
feedback about the solution of the SDE.
5 Numerical experiments
Below we provide numerical realization of some of the various objectives that can be fulfilled with the multitask objective. Details of the experimental setup can be found in Appendix
5.1 Multitask inpainting and sequential generation
We evaluate our method on three datasets: MNIST, with images of size 28 x 28, CelebA, resized to 128 x 128, and of Animal FacesHQ focused cat class, with images resized to 256 256. Details of the experimental setup are standard and can be found in Appendix B In these experiments, we use the Hadamard interpolant (
MNIST. We demonstrate the versatility of our operator-based interpolant framework through in- painting and sequential generation tasks MNIST. The results are shown Figure where all the generated images come from the same model without any retraining.
For inpainting (left panels), we replace masked regions with Gaussian noise (shown as pink for clarity), then generate only these regions while preserving unmasked pixels. This is achieved by setting the entries of a to 1 - t for masked pixels and 0 for unmasked ones. To preserve unmasked pixels, we apply a secondary mask setting 17(a, .)) to zero these positions.
Sequential generation (right panels) reformulates image creation as progressive inpainting. Starting with pure Gaussian noise, we generate the image block-by-block by successively updating the operator masks. Unlike single-pass inpainting, this requires multiple forward passes-or one per block. For each pass, we apply a 1- f only pixels in the current generation block, maintaining appropriate values for previously generated and remaining noise regions.
CelebA and AFHQ-Cat. We present benchmark results for all methods across various image restoration tasks, evaluating the average peak signal-to-nois ratio (PSNR) and structural similarity index (SSIM) on 100 test images from each dataset: AFHQ-Cat (256 x 256) and CelebA (128 128). To assess the performance of our methodology, we employed two types of masking: square masks of sizes 40 40 and 80 x 80 with added Gaussian noise of standard deviation 0.05, and random masks covering 70% of image pixels with Gaussian noise of standard deviation 0.01. We benchmark
our method against four state-of-the-art interpolant- restoration methods: PnP-Flow Martin ct al. 2025), Flow-Priors Zhang et al. (2024). D-Flow [Bn--Hamu et al. (2024). OT-ODE Pokle et al. (2024).
As shown in Table 1 our method consistently ranks either first or second both reconstruction metrics across all tasks and datasets (with all values except the last taken from Martin al. (2025)). ). Regarding visual quality (Fig. [. our method generates realistic, artifact-free images, albeit with slight over-smoothing at times.
5.2 Posterior sampling in the of model
We apply our approach in the context of the øt model in d 2 spacetime dimensions, statistical lattice field theory where field configurations spatiotemporal extent) for details see Appendix B.2 This model poses sampling challenges due to its phase transition from disorder to full order, during which neighboring sites develop strong correlations in sign and magnitude Vierhaus 22010g Albergo et al. (2019).
E LXL represent the lattice state (L denotes
8

---

Table 1: PSNR and SSIM metrics for image inpainting methods CelebA and AFHQ-CCt datasets.
CelebA
APHO-CCt
Mothod
Random
Block
Random
Block
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
PSNR
SSIM
Decraded
11.82
0.190
22.12
0,742
13.35
00234
21.50
0.744
Pokle 2014
0.865
0.914
28.84
0,838
23.88
Wer- Limi
0224
33.07
0.938
29.70
0.893
31.37
0.888
0.833
Ziann of m
71-44
52.33
0.945
29.40
0.858
31.76
0.909
25.85
0.822
Mirtin
2122
33.54
0.953
30.59
0,943
32.98
0,930
26,87
0.944
Orignal
Compteed
Generated
Gonoratinn soquorce
Checker
Book
3
0
a
a
2
bot
7
2
3
2
Figure 3: Left: In painting on MNIST using various corruptions. Right: Image generation in arbitrary orders, starting from the same initial noise, with examples showing autoregressive, block-wise, and column- wise.
The øt model is specified by the following probability distribution (v = -- e- do
(23)
where Z -gaxe - E(d) db is normalization constant and E(Φ) is an energy function defined as
1 2x
<(b) + 1 Σ [o(a)) + 1/2 Σ [o(a))+
E(d)
Σ [o(a) a~d
(24)
where a, E [O,..., 1J2 denote the discrete positions on 2-dimensional lattice of size L x L, b denotes neighboring sites the lattice, and we assume periodic boundary conditions; X K E R and > are parameters. We perform MCMC simulations to generate configuration in parameter range close to the phase transition. We these data learn stochastic interpolant of the form (9) which allows us to perform unconditional generation of new field configurations as well arbitrary inpainting (conditional generation given partially observed configurations), as reported in Appendix B.2 It also allows us to test the formalism of Section 3.3 and perform sampling of the posterior distribution defined by adding a applied field h E IR the energy, i.e. using
(Ø) E()) h Σ p(a).
(25)
Original
Comupted
Generated
Original
Comrupted
Generated
Block
Brush Imp
Rand imp
M
AFHO-Cat
CelebA
Figure 4: Inpainting using various masks left panels: AFHIQ-Cat (256 x 256). Right panels: CelebA (128 x 128). Fixing block and random conrppiion are scored against related works in Table !| competitive or superior performance in all metrics.
showing

---

True (MCMC
Generated fiterppo ano
Hattgramm of
Migo (1/N) slu)
Trus date
Prior
Gen data
Tiue
Posteriar
= 0.07)
Figure 5: Simulating lattice o theory. Top left: L = 32 x L = 32 lattice configurations at the phase transition. Bottom left: lattice examples with drift parameter h = 0.02. Top middle: Generated lattice examples phase transition. Bottom middle: generated lattice examples with field h 0.02. Right: magnetization of 2000 lattice configurations.
The additional term plays the role a reward. The results of the generation based on Proposition 3.1 shown Figure which indicate that approach permits to valid sample configurations verified by their magnetization) of the posterior without retraining.
Planning and decision making in a maze
This section applies our framework to shortest path planning in maze enviromments, drawing from Jan ner al., (2022) and et al. 2024). Using the Hadamard product interpolant (9).
impose that the paths pass through arbitrary locations in the maze (by setting og = O at these locations), reformulating planning zero-shot impainting problem. Unlike traditional reinforcement learning approaches that generate paths sequentially through Markov decision processes, our method therefore produces entire trajectories simultaneously. also avoid additional guiding terms like Monte Carlo guidance used in Diffusion Forcing Chen et al. (2024).
For training.
paths of length 300 randomly extracted from the trajectory of length 2,000,000 from Chen et al.] (2024). For simplicity, we subsample these paths every six points, creating sparse recover paths of length 300 through linear interpolation
paths of length 50, from which between consecutive points. At inference, we perform zero-shot generation between any two points in the maze by enforcing that the trajectory passes through these points: the length of the path between these locations can be varied by pinning the first point by setting σn = 0, and the second point by setting a with value of [2, 50] that can be adjusted (see Appendix B.3 for details). Typical results are shown in Fig. 6 In terms of quality assessment, we check that the generated trajectories remain within allowed maze regions: all the 10,000 paths we generated between randomly chosen point pairs avoided the forbidden areas, demonstrating robust performance. More numerical experiments Appendix B.3 demonstrate that with a similar strategy, one can impose the pathway to take detour at will, even if it implies generating a longer path.
end
start
Figure 6: One-shot generation of pathways between two arbitrary points in the maze. The path length automatically tuned via a simple heuristic, see discussion Appendix B.3
10

---

Acknowledgements
We would like thank Yilun Du for many helpful discussions on the maze planning problem. MSA is supported by Junior Fellowship at the Harvard Society of Fellows as well as the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF Al Institute for Artificial Intelligence and Fundamental Interactions, htt::/aaifi.rrg.

---

References
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2022.
Michael S Albergo and Eric Vanden- Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference Learning Representations, 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in neural information processing systems, volume 33, pages 6840-6851, 2020.
Yang Song and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. arXiv:1907.085 1907.05600, 2020.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. In Advances in Neural Information Processing Systems, volume 34, pages 17695-17709, 2021.
Michael S Albergo, Nicholas M Boffi, and Eric Vanden Eiinden., Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303. 08797, 2023a.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEF/CVF confer- ence on computer vision and pattern recognition, pages 10684-10695, 2022.
François Mazé and Faez Ahmed. Diffusion models beat gans on topology optimization. In Proceed- ings of the AAAI conference on artificial intelligence, volume 37, pages 9108-9116, 2023.
Michael Alverson, Sterling G Baird, Ryan Murdock, Jeremy Johnson, Taylor D Sparks, et al. Generative adversarial networks and diffusion models in material discovery. Digital Discovery, (1):62-80, 2024.
Hyungjin Chung, Jeongsol Kim, Michael Thompson Mecann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In International Conference on Learning Representations, 2023. URL /
Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. In International Conference on Learning Representations, 2022. htRp:
Hengkang Wang, Xu Zhang, Taihui Li, Yuxiang Wan, Tiancong Chen, and Sun. DMPlug: A plug-in method for solving inverse problems with diffusion models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv:2011.134555 2020.
Michael S. Albergo, Mark Goldstein, Nicholas M. Boffi, Rajesh Ranganath, and Eric Vanden-Eijnden. Stochastic interpolants with data-dependent couplings. arXiv:2310.03725, 2023b.
Marcelo Pereyra, Philip Schniter, Emilie Chouzenoux, Jean- -han-Christophe Pesquet. Jean-Yves Tourneret, Alfred O Hero, and Steve McLaughlin. survey of stochastic simulation and optimization methods in signal processing. IEEE Journal of Selected Topics in Signal Processing, 10(2):224-241, 2015.
Bahjat Kawar, Michael Flad, Stefano Ermon, and Jiaming Song. Denoising diffusion restoration models. Advances in Neural Information Processing Systems, 35:23593-23600 2022.
12

---

Ségolène Martin, Anne Gagneux, Paul Hagemann, and Gabriele Steidl. Pnp-flow: Plug-and and-play image restoration with flow matching. In International Conference on Learning Representations, 2025.
Florentin Coeurdoux, Nicolas Dobigeon, and Pierre Chainais. Plug-and-play split gibbs sampler: embedding deep generative priors in bayesian inference. IEEE Transactions on Image Processing, 2024.
Yu Sun, Zihui Wu, Yifan Chen, Berthy T Feng, and Katherine I. Bouman. Provable probabilistic imaging using score-based generative priors. IEEE Transactions on Computational Imaging. 2024.
Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective solving inverse problems with diffusion models. arXiv preprint arXivv2305.04331, 2023.
Cagan Alkan, Julio Oscanoa, Daniel Abraham, Mengze Gao, Aizada Nurdinova, Kawin Setsompop, John M Pauly, Morteza Mardani, and Shreyas Vasanawala. Variational diffusion models for blind mri inverse problems. In NeurlPS 2023 workshop deep learning and inverse problems, 2023.
Giannis Daras, Yuval Dagan, Alexandros G Dimakis, and Constantinos Daskalakis. Score-guided intermediate layer optimization: Fast langevin mixing for inverse problems. arXiv preprint arXiv:2206.09104 09104, 2022.
Arpit Bansal, Eitan Borgnia, Hong-Min Chu, Jie S. Li. Hamid Kazemi, Furong Huang, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Cold diffusion: Inverting arbitrary image transforms without noise, 2022. htRps:ts::
org/ abs/ /als//220..e2 2208.09392
Keyu Tian, Yi Jiang, Zehuan Yuan, Bingyue Peng. and Liwei Wang. Visual autoregressive modeling: Scalable image generation via next-scale prediction, 2024. URL http:::///arrriw.rrr abs/
Tianhong Li, Qinyi Sun, Lijie Fan, and Kaiming He. Fractal generative models, 2025. URL
https://arxiv..rgg wwt.:oor//abb/////// 2502. 17437
Jiaxin Shi, Kehang Han, Zhe Wang, Amaud Doucet, and Michalis K. Titsias. Simplified and generalized masked diffusion for discrete data, 2025. URL httpt:///aaaii.oorrrrr org.abs//////
04329
Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions, 2025. URL https
// //aaxiv.org//// abs 06768
Yasi Zhang, Peiyu Yu, Yaxuan Zhu, Yingshan Chang, Feng Gao, Ying Nian Wu, and Oscar Leong. Flow priors for linear inverse problems via iterative corrupted trajectory matching. Advances in Neural Information Processing Systems, 37:57389-57444 57417, 2024.
Heli
Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D-flow: Differentiating through flows for controlled generation. In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 3462-3483. PMLR, 2024.
Ashwini Pokle, Matthew J. Muckley, Ricky T. Q. Chen, and Brian Karrer. Training- frren linear image inverses via flows. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
Ingmar Vierhaus. Simulation of de theory in the strong coupling expansion beyond the Ising Limit. PhD thesis, Humboldt University of Berlin, 07 2010.
M. S. Albergo, G. Kanwar, and P. E. Shanahan. Flow-based generative models for markov chain monte carlo in lattice field theory. Phys. Rev. D, 100:034515, Aug 2019. doi: 10. 103//hyssReID. 100.034515. URL https://link aps orw.oog//oooo btgg/di/110 1103 1wt/////////......mmm 100. .000.533155
Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis, 2022. URL /tttttsst//aiiiiii h/tp::///aaxx.e/ee org/ hatp://b///// abs//205.... 09991
13

---

Boyuan Chen, Diego Marti Monso, Yilun Du, Max Simchowitz, Russ Tedrake, and Vincent Sitzmann. Diffusion forcing: Next-token prediction meets full-sequence diffusion, 2024. URL https
01392
David Williams. Probability with Martingales. Cambridge Mathematical Textbooks. Cambridge University Press, Cambridge, 1991. ISBN 9780521406055. doi: 10.1017/CB0978051118800 18136558
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017. URL https://arxiv. org. abs/ 1112..9988
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomed- ical image segmentation. CoRR, abs/1505 04597, 2015.
04597
A Proofs
Definition 2.2 (Multipurpose drift). The drifts no, m :5:S% x H. + H. are given by 1,,a, 3,x) = ([zol/(a,, B) = x],
mi 3,2, x) E[ 3)= =]]
(2)
where E [(x, = x]] denotes expectation over the coupling (zo; .11 y conditioned on the event T [a,
Lemma 2.3 (Drift objective). Let v da, d3) be probability distribution whose support is S. Then the drifts 10,1 3) in Definition 2.2 can be characterized globally for all (a,B) S and all a E supp(po,s) via solution of the optimization problems
Voo argmin
E (a, (#0.2))~H
[Ilño(a, 3, [a, =rol
%o
7/1 arginin E (a,
(er 3, ((, 1111
(4)
f1 (x0,21)~4
where ||-|| denotes the norm in H.
Proof. The lemma is a simple consequence of the L characterization of the conditional expectation as ceast-squares-bes best predictor, see e.g. Section 9.3 in Williams (1991).
Lemma 2.4 (Score). Assume that H = Rd and that the probability distribution pa.3 of the stochastic
interpolant I (a, B) is absolutely continuous with respect to the Lebesgue measure with density (r) Assume also that Ty ~ [O, Id) and Ny 1 B1. Then the score da,)(x) = log Pa, (x) is related the drift nola, 3.x)) via
no,a, 3,x) -a)
(6)
Proof. The lemma follows from Stein's lemma (aka Gaussian integration by parts formula) that asserts that
E[ro|I(a, x] ((5,,(()))
(26)
well the definition of no(a, in (6)..
Proposition 2.5 (Probability flow). Let (av, Be) be any one-parameter family of operators
(on, Bi) E S. Assume that on, Br are differentiable for all t [0, 1). Then, for all E (0, 1), the law of I (ar, B.) is the same as the law of the solution X, to
X, = - Bc. Xt) + Bem (on, Br. Xt),
Xo d I(xo, 3o).
Proof. From the framework of standard stochastic interpolants Albergo and Vanden Vanden-Eijnden (2022); Albergo al. (2023a), we know that the law of I, I(a,, 3t) the same for all t E [0, as the law of X₁, l.e. the solution to the probability flow ODE
X, = (xe),,
X = 1:=0-
(27)
14

---

where
b((x) = B|4 14 = z²].
(28)
By the chain rule J, 0, 8, so that
bx(x) = dr 4 = x] + dr = x] = (a Be, + (an
(29)
where 10,1 are the drifts defined (T)). This means that (27
Proposition 2.6 (Diffusion). Assume that H = R and the probability distribution Pa. of the stochastic interpolant [(,,3) 3) is absolutely continuous with respect to the Lebesgue measure. Assume also that, in [(, 3), 2o is Gaussian and Zo 1 F1. Then, under the same conditions as in Proposi- tion 2.5 if ag is invertible, for all E =0, 1] and any et 2 O, the law of [av, 3c) is the same as the law of the solution X;
dxj ( "no Be, Xi)dt + Bemiaas, Br, Xi)dt 22,,WW,
Xú [(,, Bo). (8)
where We is a Wiener process in Rd
Proof. From the framework standard stochastic interpolants Albergo and Vanden- (2022)
Albergo et al. (2023a), we know that the law of the solution X₁ the probability flow ODE (27 the same for all f E [0, 1] as the law of X; solution to to SDE
dx; = be (X.)dtt (X))( + W2,,WW,
Xi=or d 1==0,
(30)
where se(x)) is the score of the probability density function of Le = I(a,, Since st(x) Sa (2)), from Lemma 2.4, we have
sx(r) -a, (are B1,
(31)
If we insert this expression in 30 and use Proposition 3.1. Let
22 we see that this SDE reduces to S
(o, (,x) = E xol (a, 3) = x],
m (/ = x]
(16)
be the drifts associated with the interpolant
[4 and
1o(,,, B, = [ox 3) = x of (ax, x,x)= =(( 3) =]]
(17)
be the drifts associated with the interpolant 115
If and 3 are invertible, then
(a, B,x) a "1,88',,,0000 B,, + a1 -(x -3,, xx,)
(18)
(a, B, = T71 (ar.Brr xr) as long as we can find pair (otr, BBr that satisfies
(19)
oT, a Br a
(20)
and is given by
xz = = a
(21)
We will prove this proposition as a corollary of:
Proposition A.1 (Posterior distributions). Let Ma,9 and P2. be the probability distributions of the stochastic interpolants defined in (14) and (15), respectively. Assume that and 3 are invertible, that the equations (20) for ar, 3r in Proposition have a solution, and that x, is given by (1).. Then these distributions are related, up constant independent of and vr,
(d. = |ay||a||- ((..BB Pa..3, (dxr))
(32)
where
R(a, 3,x) }/a, xx, aa 12
(33)
15

---

Proof. By definition of the probability distribution (a,o(da) of [, 5), given any integrable and bounded test function Φ IR I have
S (( (dx) (a,
(34)
Jaaxx Moro + (20) e 2mool doge ps (dr1)
If instead of zn we use as new integration variable x = 0200 + Ba,, this becomes
J (( (dx) = |a|- feene ((z) 4/2e H/- (a- (f dxe (m1) (1((11))
(35)
Similarly, for the probability distribution Ma (dx) of T(a, (), have
Jer ((x))a._(dd)) = |al-= / JREXR* (2 e Hla '(x "(x-B21)) dz (dx1)
(36)
If in this equation replace by ar, by Br, by Tr, and ((x) by R(. and multiply both side by |arl/ l/la/ it becomes:
/
Ja, |||]
((x))
Pars (dx,)
Rd
(37)
= |||
/ (())
(2T) d/2 6/22 - (0-30(0 dz (dx1).
Rd
We can now require that the right hand side of (37] be the same as at the right hand-side of 35 (so that, 1a (dx) = |ar||a| 1lolllol R(o, B, (dx,)) )) we arrive the requirement that
1 [ (zr, - 1,,x1) + R(a, x) = 1 -aa Br,))2 + 2 (21, Ari) + (b,
(38)
where we used r(x) = (z. Ax) + xb. Since (38 must hold for all ø1, we can expand both sides
of this equation, and equate the coefficient of order 2, and 0 in ø1. They completely equivalent
to 20 21.. and (33. respectively. So as long as we can find solutions to (20). (32) holds.
Proof of Proposition
3.1
By definition of the conditional expectation, we have
m 31, Sq x (x- Bm,)) (dæ1)
(39)
Jge e 2 (x- 11 (dx1)
(a, 3,x) = SE #1e 3iE Jg (x- 3 2 (^(ا)) μي (dx1)
(2- (())) μn (dx1)
(40)
in the first equality we replace a by a,, 3 by B,, and x by I,,, and assume that ay, 8r, and x satisfy (38). by construction obtain that (19 holds. To get (88) (19) as well as relation twice to deduce
T B,,) a (s nug (Ir, a, 3))
=a (s Bm (Ir, Or: 3r))
(41)
-1 (5 33, 1 (xr - ar ar (0, (3r
a 133, a, no ( Cr,, 3r) + a a B3, xr).
B Experimental details
Details for the experiments in Section are provided here.
16

---

Original
Comrupted
Generated
Figure 7: q'-moel: Inpainting of three different configurations.
B.1 Multitask inpainting and sequential generation
For all image generation experiments, the U-Net architecture originally proposed in Ho et al. (2020) used. The specification of architecture hyperparameters well training hyperparameters are given in T Training was done for 200 epochs on batches comprised of 30 draws from the target, and 50 time shices The objectives given B and were optimized using the Adam optimizer. The learning rate was set to .0001 and was dropped by factor of 2 every 1500 iterations of training. To integrate the ODE/SDE when drawing samples, we used simple Fuler integrator.
In order to progressively explore the space of the hypercube of a and B, we first learn a model in the diagonal of the hypercube, where all entries of a are all the same value. We then fine-tune the first model for matrices a uniformly distributed in [0, 1jd, We also fine-tune the first model for matrices a decomposed by blocks of 4 X where entries each blocks contains the same values in [0, 1]d.
MNIST
64
CelebA
AFHQ-Cat
Dimension # Training point
28x28
32x32 100,000
128x 128 190,000
256x256
60,000
5,000
Batch Size
50
100
128
64
Training Steps
4x10a
105
x 105
105
Attention Resolution
64
64
64
64
Learning Rate (LR) 0.0002
0.0002
0.0001
0.0001
LR decay (Ik epochs)
0.995
0.995
0.995
0.995
U-Net dim mult Learned embedding
[1,2,2.2]
[1,2,2.2]
[1,2,4.8]
[1,2,4.8]
Yes
Yes
Yes
Yes
# GPUs
Table 2: Hyperparameters and architecture for MNIST, t and maze datasets.
17

---

B.2 Details about the of Model
We define the discrete Fourier transform as
ē(k) L L Lirk ((a)
Φ(a) L L- Σ 2irk ((k)) (.)))))..
(42)
where a, k E 10....,L L 119 we can write the energy 24 E()) Eo(Ø) + (b)) with
= 1 Σ (H)))(()) M(k) = 2a d Σ cos - + So. (43)
where denotes the basis vectors on the lattice and Bu > 0 is an adjustable parameter; and
U(6) 2 2 no) Σ 1
1
[o(a)l*
4
(44)
1 (n kO) (Ø)
1²(k))2 2
1
Σ
L
L-1/2
4/2
o ((k)
2
4
k
where & and o are Fourier transform pairs defined in 42 E. (iff(()))))
the last term can be implemented via
Sampling using the Langevin SDE: To obtain the ground-truth samples from the of model, one option use the SDE
(k)(k) = N (k) (Is (k)) 78) v2dW.(k)
(45)
where we denote
3
(k)= L-4/2 4/2
Σ 2ixka//L L--d/2 E. 2ink ((kk))
(46)
which can be implemented via fft((ift(()) (3) This SDE may be quite stiff, however, a problem that can be alleviated by changing the mobility and using instead
(k) - (k:) (m-xo)) KO) ((*)) 7M
(47)
The discretized version of this equation reads
T (k) di,, (A:) - Al,, @e (k) + ko)- (,-----
(48)
+ 1/ (k),
where in is the Fourier transform of 1)n. (O, Id).
Computing the generator of the posterior distribution As illustrated with (25). one would like sample a slightly different of model with energy function, noted E,. E' and E, define respectively the prior and posterior distribution, via the Boltzmann's 's law (23). In this work, only relations of the form E, E (d, As) (h, sp) are studied, combining linear and quadratic term. assumed to be definite negative.
Lemma B.1. Assume that 3 a and let
n( x)
E(ro- a1laro + 11--01 (1-1) x| =noaa, a. (/ 1 a, a,
(49)
Then
no(a, 0, x) + a)) m (a, 1 - a, . = an(a
(50)
(51)
Proof of lemma B. Solving
and 49 in (7o. ?/) gives 50 and
The idea is to use the drift of the prior to sample from the posterior. The following proposition makes possible.
18

---

Proposition B.2 (Posterior drift). Assume
with a diagonal and invertible and let
(a, x) = E 21 lazr + (1 (x) ox)
2]] (a 1 -a, (ax 1 a,,))
(52)
where the drifts no and nf are defined in (8 and 19 respectively. Assume also that A is diagonal, non-positive definite, and invertible. Then, 3,
ar and 1f" can be expressed as:
1/ (a, x)) = n,ar, a
+ +a a- (x -x'),
(53)
where n(a, x)) is the drift of the prior defined in 49 and at and x, are given by
0r = V 2a + a
A) a
(54)
1 - 2a a²
az a)
0
T a² - a+)
(55)
Proof of the Proposition B.2
Using 50
in 18
and
one obtains:
na 3, x) = a a 33, a,, a, (m 3r, x,) + x,) + a -x (33,(x,).
(a, 2. I, a.n 3,)
Then, take the difference, and regrouping terms together eventually yields:
(a, B,x) (a, 3: (a Bar Brr Fr) n(ar xT) a 133 (a, 1)5y 3r a-
Ea-
--a-180
a a-tann 2r, + a tx -x,
The linear case Assume for now that A 0 and b = h. Note that one recovers the same case that in (25), that is, one applies uniform magnetic field of magnitude h over the whole lattice. What follows is simply a corollary of B.2
Proposition B.3, Assume E, = E (0,1). h). Then, or = a, Br = B, and
Φ, o + aax 3-T
(56)
Also, the posterior drift 1f" writes:
(a, 3, n,,,, (b)) -= 3-Th.
(57)
Proof of Proposition B.3 Since A = 0. 20 directly implies a,, = and B, B. Consequently,
(56) follow from 21 and 57) from 53
In summary, a simple shift proportional to appears in the posterior field Φ,. It clearly tends to favor the alignment with the magnetic field, which obeys common sense.
The quadratic case
Assume that = 0 and = 1d.
In similar fashion to the linear case, one derives analytical expressions for the quantities of interest.
Proposition B.4. Assume E, Ap) E k2 E(()), B, a, and 1 a.
Then
= a + / 2a + a (1 12)
(58)
1- 1-2a- + 2
0r (-a +
2x + a² +
(59)
/1 - 2a + a a + (1 2a 2 22
19

---

and
(a))),0)) +a v1 2a a= + h2
ma Sr,
(60)
1- 2a + a²
+a (1 (1 ( + -1- - 2a a² +1))++])]111111
(61)
VI 2a a² k2) 2a + a2k2)
Proof of the Proposition B.4 Given the assumptions, 20 yields:
(Id (ty) ) aa, a (Id a+) (Id a) o a) 1-
Observing that a and A are diagonals, o, is also diagonal. Furthermore, assuming that O, is proportional the identity, the above reduces to the scalar equation (keeping the same notation for conciseness):
)2) a)) k2
a
a²
After few elementary manipulations, one arrives at: (1-a, al = a1 - a)2 + k²a a2
This is quadratic equation that admits two solutions. Only one positive, and writes as:
a2 + a Lvvi 2a + a + k2)
=a + V 2x + a² + k2)
(62)
1 2a 22
1 1 1-2a- a2k²
quite casy to check that the discriminant is always positive, so it does not pose any problem Also, if k > and E [0, 1J, then ar 1]. This property is necessary, since + n( a) has been trained in the hypercube [0, 1]d.
After elementary simplifications and recalling that
Or, and 21 one has:
2 -
-x
/1 - 2a + a2 + k2)
a Cr
Φ
2
ar
1 2a 221.2
22 to² -ay 10110120222112
1-20 (1 -ayl-aal-- (1+6a)
Since 1 - a2
it yields:
1-20+a=12
1-2a+00k2
Φ (-a /1 2a a² + b2))²
1 a 2x + a² 122) -
(63)
1 2a 22k²
2x + a² + +2)
then factorizing by v 2a + a² 12) eventually gives (59).
Eventually, after replacing 62 and (59 into (3) (60) holds.
Assume that and 1-Id. In this case, the quadratic equation is:
)2
(1 - a)²
k2 a²
a²
or otherwise stated:
(1-a,) (1 a) a) o,20 2² = - kk)) 2a 1. Assuming it strictly positive, The discriminant of this polynomial is A a1 among the two solutions, only one is positive:
2a + a a² - 1.2) a
a, =
2x (ak))
The polynomial inside the square root is positive if and only if a & 1- To see that, see there exists always two real roots, since the discriminants is 41,2 > 0. Those roots T 1 and + > Since a, E [0, 1] must be respected for all a E [0, 1], only k: < can be considered with our method. Consequently, sampling using stochastic interpolants from a to a appears impossible with this method.
20

---

B.3 Details about the maze experiment
We use the Hadamard interpolant (9) and estimate the drift rfa, 2)) defined (10 by approximating it with a U-Net neural network Ho et al. (2020), trained with an Adam optimizer Kingma and Ba (2017). The U-Net comprises 4 stages with 48, 80, 160, and 256 channels respectively for the encoding flow. The decoder has the same architecture as the encoder but in reverse order, with added residual connections Ronneberger et al. (015). Each stage consists of residual blocks, with the first concatenated with self-attention block. The input vector has shape d x 2, where row contains the x and y coordinates of the i-th point in the trajectory.
In contrast to conventional U-Net architectures, we perform interpolation and max pooling operations independently each coordinate column to increase and reduce dimensions only along the trajectory length axis. The convolution kernel size is 5 2, processing each point's 's coordinates together with those of its two temporal predecessors and successors in the sequence. We add the necessary padding to maintain identical input and output dimensions, which amounts to padding by two rows at the top and bottom of the input vector.
Given a pair of randomly chosen points in the maze, we must determine where to constrain these points along the generated trajectory. If the constraint points placed too far apart in the sequence (large index difference), the resulting path will likely not be the shortest route; conversely, if placed close together (small index difference), the generated path has an increased chance of cutting through forbidden regions, making it inadmissible. To address this trade-off, we adopt the following heuristic. We fix the starting point at the beginning of the path (index = and employ a progressive search for the target point placement using the candidate indices [5, 10, 20, 30, 40, 45, 50J. We first generate a path with the target point constrained index (creating short trajectory). this path intersects forbidden regions, we increase the target index to 10 (allowing longer path), and continue this process until we generate a valid path that successfully avoids all obstacles.
On Figure 8, we impose paths to go by the bottom-right corner, the constraint is visible as small white dot. The path length adapts accordingly.
Figure Two generated pathways, one per row. Left: No constraint is imposed on the path, other than joining the two endpoints. Right: An additional constraint is imposed: the path must pass through the hottom-right corner, represented by white dot.
C Additional experimental results
Here we provide additional infilling image results, given in Figure
21

---

CelebA
AFHQ-Cat
Original
Comupled
Generaled
Original
Comupped
Generaled
Figure 9: Additional images demonstrating the inpainting task: Block inpainting is shown in the top panels, while random inpainting is displayed in bottom panels. The left panels depict images from the AFHQ dataset. and the right panels show images from the CelebA dataset.
22

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers
integral part of your paper submission. They visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] provvided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used").
general, answering "[No] " or "[NA]" is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to claborate. All supporting evidence can appear either in the main paper the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: The abstract claims that the formalism of stochastic interpolants have been adapted for zero-shot conditional generation, in-painting, and posterior sampling. The corresponding numerical experiments are explicitly displayed in their associated sections. Posterior sampling has been theoretically explored and numerically investigated in the ø² model with stochastic interpolants. We show that our model can learn generate images at multiple scales: pixel-wise, blockwise, and in all dimensions simultaneously.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals not attained by the paper.
23

---

2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: discussion on this matter is present the conclusion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification:
All proofs, with their set of assumptions, are provided in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]

---

The details of all the numerical experiments are discussed either in the main text of the paper or in the Appendix. especially concerning the one-shot conditional generation and inpainting.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (c.g.. a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All the code and data used for the numerical experiments will be exposed in Github repository. set of instructions to fully reproduce the results will be provided in README file.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https //nips http://wwww..//eeeeee/////eaaaaa
for more details.
25

---

The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
reproducible, they
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All test details are provided in Appendix (e.g. Table
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: The full distribution of the average magnetization has been studied for our o+ model in Figure
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
26

---

Answer: [Yes]
Justification: A Table containing all relevant information will be provided in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: Our paper presents theoretical and mathematical foundations for multitask generative modeling using operator-basee interpolants. The research primarily consists mathematical formulations and theoretical derivations. Our experiments are limited to publicly available benchmark datasets (MNIST digits) and physical simulation data, neither which contain sensitive information or raise ethical concerns. We do not collect or use personal data, conduct experimentation with potential for harm or misuse. All sources and prior work are properly cited, and our mathematical derivations, proofs, and experimental procedures
humans animals, develop technologies
presented transparently. The research was conducted with scientific integrity, honesty, and transparency throughout the process, fully adhering all aspects of the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Our paper presents a theoretical framework for operator-based interpolants in generative modeling, which primarily advances fundamental research in this area. The work is largely theoretical and mathematical in nature, focusing on developing a universal approach to multitask learning rather than specific applications with direct societal implica- tions. While we discuss potential technical benefits in the paper's 's conclusion, the theoretical nature and early stage this research makes specific societal impacts, whether positive or negative, difficult to assess meaningfully at this time.
Justification:
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to
27

---

any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This question is not applicable to our paper do not release any models datasets that present high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: In paper, properly credit the creators and original owners of all assets used, including the MNIST dataset and any referenced algorithms or methodologies. For the MNIST dataset, which is in the public domain, we acknowledge its source and cite the original publication. All assets are used in accordance with their intended research purposes and we have carefully respected all applicable terms of and licensing requirements throughout our research process. Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcon com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
28

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Our paper introduces new theoretical formulations and algorithms for operator- based interpolants, which are thoroughly documented within the paper itself. The mathemati- cal framework, definitions, lemmas, and propositions rigorously presented with complete derivations and proofs (in the appendix). For our experimental implementations, we provide detailed descriptions of the model architectures, training procedures, and inference methods in the paper and will release accompanying code with comprehensive documentation that ex- plains the implementation of our operator-based interpolant framework. This documentation includes clear instructions for reproducing our experiments, explanations of key parameters, and examples demonstrating how to apply our methods various tasks. All assets carefully documented to ensure transparency and reproducibility of our research.
Guidelines:
The answer NA means that the paper does release assets. Researchers should communicate the details of the dataset/code/mmodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This question is not applicable to our research as our paper does not involve any human subjects or study participants. Our experiments are conducted exclusively standard benchmark datasets (MNIST) and physics simulation data, with no human participation involved at any stage. Since no human subjects were part of this research, no IRB approvals equivalent reviews were necessary, and there were study-related risks to disclose or manage.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such the institution conducting the review.
15. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact core methodology, scientific rigorousness, originality of the research, declaration is not required.
29

---

Answer: [No]
Justification: No LLM or any kind of transformer architecture is involved in the numerical experiments displayed.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non ntn-standard components.
Please refer to our LLM policy (https: ptttps://neurip... wwwpouurppsccoooofeeeeeeeeeeeeee scc/oonferences/22222 /LLM/LLM for what should or should not be described.
30


Reviews:
[{"review_id": "bxgVIvtf16", "reviewer": "Reviewer_sQN9", "summary": "The aim of this paper is to train a single model that can perform various tasks. To do this, the paper generalizes the interpolants used in flow matching to accommodate diverse tasks. The suggested approach introduces additional terms in the training loss for the various interpolants but after training various tasks can be performed. Theoretical justifications are provided for the generalization of the interpolants and moderate-scale experiments are provided.", "questions": "Could the authors provide metrics that compare the multi-tasked trained models with single tasked based models in order to clarify the trade-offs involved?", "limitations": "Limitations are discussed.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "# Strengths\n\nThe approach is novel and allows a single model to perform diverse tasks. The paper's method is well justified, and the writing is clear.\n\n# Weaknesses\n\n1. The main weakness of this work is the limited scale of the experiments, although it is largely sufficient as a proof-of-concept. The work would be strengthened by a discussion of the trade-offs of having such a general model: how does it compare with a bespoke model for the various tasks considered?\n2. It appears that invertible maps $\\alpha$ and $\\beta$ are necessary for the results to hold. The work would be strengthened by having more discussions around the choices of $\\alpha$ and $\\beta$. I believe this introduces a limitation on the diversity of tasks the approach can handle.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive evaluation and acceptance recommendation. We want to emphasize that our work represents a **fundamental conceptual breakthrough** - moving beyond traditional time-based diffusion to operator-based interpolation, with **theoretically guaranteed unbiased sampling** unlike existing approximation-based methods.\n\n>***Larger-scale experiments:***\n\n Since submission, we have conducted extensive experiments on CelebA (128×128) and FFHQ (256×256) datasets to demonstrate scalability and  benchmark our approach against different methods in various image restoration tasks, evaluating the average Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index (SSIM) on 100 test images from each dataset. To assess the performance of our methodology using Hadamard product-based operators, we employed two types of masking: square masks of sizes $40\\times40$ and $80\\times80$ with added Gaussian noise of standard deviation 0.05, and random masks covering $70\\%$ of image pixels with Gaussian noise of standard deviation 0.01.\n\nThe results reported in the table below show that our method consistently ranks either first or second in both reconstruction metrics across all tasks and datasets. Regarding visual quality, our method generates realistic, artifact-free images.\n\n**CelebA Dataset (128×128)**\n| Method     | Random PSNR | Random SSIM | Block PSNR | Block SSIM |\n|------------|-------------|-------------|------------|------------|\n| Degraded   | 13.35       | 0.234       | 21.50      | 0.744      |\n| OT-ODE [1]     | 28.84       | 0.838       | 23.88      | 0.874      |\n| D-Flow [2]    | 31.37       | 0.888       | 26.69      | 0.833      |\n| Flow-Priors [3]| 31.76       | 0.909       | 25.85      | 0.822      |\n| PnP-Flow [4]  | 32.98       | 0.930       | 26.87      | 0.904      |\n| Ours       | **33.11**   | **0.945**   | **26.96**  | **0.914**  |\n\n**AFHQ-Cat Dataset (256×256)**\n| Method     | Random PSNR | Random SSIM | Block PSNR | Block SSIM |\n|------------|-------------|-------------|------------|------------|\n| Degraded   | 11.82       | 0.197       | 22.12      | 0.742      |\n| OT-ODE [1]    | 28.36       | 0.865       | 28.84      | 0.914      |\n| D-Flow [2]    | 33.07       | 0.938       | 29.70      | 0.893      |\n| Flow-Priors [3]| 32.33       | 0.945       | 29.40      | 0.858      |\n| PnP-Flow [4]  | 33.54       | 0.953       | **30.59**  | **0.943**  |\n| Ours       | **33.76**   | **0.967**   | 29.98      | 0.938      |\n\n\nThese experiments indicate that our operator-based interpolants maintain their versatility at scale, and confirm that our approach scales effectively while preserving the key advantage of training once and adapting at inference time.\n\nIn addition, we have preliminary results on a **robotics maze planning experiment**, as is done in [5]. In this setting, given starting and terminal coordinates in a maze, the aim is to construct a path between the points along the maze without hitting walls or obstacles. Our method allows the zero-shot generation of this path for any pair of points in the maze without re-training. In our preliminary results, we see a 100% success rate in this task, and are also in direct communication with the authors of [5] to produce their reward metrics, if that is of interest to the reviewers.\n\n>***Trade-off analysis***:\n\nThese quantitative results directly address your question about trade-offs. While our multitask model has  higher computational overhead during training, it achieves competitive or superior performance compared to specialized methods across tasks while eliminating the need for separate models. **Critically, unlike existing multitask approaches that introduce systematic biases through approximations, our method provides exact, unbiased sampling for all tasks.**\n\n>***Invertible operators $\\alpha$ and $\\beta$:*** \n\nWe appreciate this question, but note that invertibility is not required for all tasks - it primarily helps with computational efficiency by allowing us to train one η field and deduce the other from relation (5), rather than having to train both η₀ and η₁ separately. For our main applications (Hadamard products, Fourier domain operations), invertibility is naturally satisfied.\n\n\n**References:**\n\n[1] A. Pokle, M. Muckley, R. Chen, B. Karrer, \"Training-free linear image inverses via flows\" (TMLR)\n\n[2] H. Ben-Hamu, O. Puny, I. Gat, B. Karrer, U. Singer, Y. Lipman, \"D-Flow: Differentiating through Flows for Controlled Generation\" \n\n[3] Y. Zhang, P. Yu, Y. Zhu, Y. Chang, F. Gao, Y. Nian Wu, O. Leong, \"Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching\" (Neurips 2024)\n\n[4] S. Martin, A. Gagneux, P. Hagemann, G. Steidl, \"PnP-Flow: Plug-and-Play Image Restoration with Flow Matching\" (ICLR 2025)\n\n[5] B. Chen, D. Marti Monso, Y. Du, M. Simchowitz, R. Tedrake, V. Sitzmann \"Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion\" (NeurIPS 2024)"}, {"review_id": "D7lcKroigO", "reviewer": "Reviewer_MR8W", "summary": "This work extends stochastic interpolants to support multitask generation via extending the scalar interpolation to structured interpolation. The structured interpolation provide the flexibility that one could learn various tasks during training and assign a specific one (or some) during inference, eliminating the need of designing a distinct model for each task. Conceptually, this work demonstrates the framework would be capable to conduct arbitrary inpainting, posterior sampling with quadratic reward, multichannel denoising, and inference adaption. Empirically, they show the proposed framework is indeed capable to perform the arbitrary inpainting and posterior sampling tasks.", "questions": "See weaknesses", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths\n* The motivation for the proposed structured interpolation is clear\n* The high-level conceptual thinking of multi-purpose generative modeling is nice\n\n## Weaknesses\nIn my opinion, this paper is still on its primitive stage and it's not ready for publishing yet.\n\n* **Contextualization of related works and contribution of this paper**: \nAs the author said in line 128, the main focus is to exploit the flexibility of the extended structured interpolation. However, this is not the very first attempt and thus prior related works should be discussed. A non-exhaustive list is diffusion forcing [1] and MuLAN [2], which also employs certain heterogeneous diffusion process. With that being said, the author should highlight the technical contribution regarding to the extension, e.g., key steps that makes a notable difference from the scalar interpolation case, if any.\n\n* **Clarity of the presentation**: \nThe presentation in Section 3.4 is not great. The \"inference adaption\" problem itself is not well-defined, and it is difficult to tell what does \"bridge the distributions\" mean without any context for general audience. This can be addressed by adding a formal description of the problem itself on either main text or appendix.\n\n* **Limitation of the posterior-sampling application**:\nTo me, the relationship between $\\eta$ and $\\eta^r$ is largely attribute to the specific quadric structure of the reward and the relationship seems not depend much on whether the interpolation is scalar or structured. It seems like the benefit of the structured interpolation is still to support arbitrary inpainting when performing posterior sampling rather than being more capable at this setting itself.\n\n* **Evidence to support the claims**:\nI have to say the numerical experiments are not very convincing to show proposed framework is superior. In my opinion, at least there should be certain demonstration for the multichannel denoising and inference adaption that are demonstrated as conceptually plausible in Section 3. A quantitive comparison between previous methods (e.g., the approximated strategies on inference time) would be a plus, which is important to practitioners to get a sense of any possible trade-offs between versatility, quality, and learning difficulty.\n\n[1] Chen, Boyuan, et al. ‘Diffusion Forcing: Next-Token Prediction Meets Full-Sequence Diffusion’. (NeurIPS 2024).\n\n[2] Sahoo, Subham Sekhar, et al. ‘Diffusion Models With Learned Adaptive Noise’. (NeurIPS 2024).\n\n### Minor\n* In line 102, \"clean data corrupyed\" -> \"clean data corrupted\"\n* In line 146, \"They that\"?\n* In Eqn. (11), the brackets are not closed\n* In line 212, \"Wasserstein lenght\" -> \"Wasserstein length\"\n* In line 214, missing space after $t\\in[0,1]$", "was_revised": true, "final_justification": "The authors have successfully addressed my main concerns about contextualization and clarity issues. The reason preventing me from giving a higher rating is that, while the conceptual framework is novel, the execution (e.g., empirical justification, general scenario beyond quadratic reward constraints) could still be improved.", "rebuttal": "We appreciate the detailed feedback. We want to stress that our work represents a **completely novel conceptual paradigm**: while existing methods still operate within traditional time-based diffusion frameworks, we fundamentally reconceptualize multitask generation through operator-valued interpolants. **Crucially, our method provides theoretically unbiased sampling** unlike current approaches which often comes with inherent systematic biases.\n\nShould you find our reply satisfactory, we kindly invite you to reconsider your rating. WE would also be happy to answer any additional questions you have or provide clarification on our replies.\n\n\n>***Related work contextualization:*** \n\nYou are correct about relations to diffusion forcing and MuLAN, but **our approach represents a paradigmatic shift from temporal diffusion to operator-based interpolation.** While existing methods still rely on the concept of diffusion in time with task-specific modifications, we completely reconceptualize the problem. Our framework enables fine-tuning and posterior sampling capabilities these methods cannot handle, and **guarantees exact sampling from target distributions, eliminating the approximation errors that compromise existing multitask methods.**\n\n>***Section 3.4 clarity:*** \n\nWe have completely rewritten this section to address your concerns. **Inference adaptation** refers to the problem of optimizing the path ($\\alpha_t$, $\\beta_t$) used during generation to achieve specific objectives, such as minimizing computational cost, maximizing sample quality, or satisfying user constraints. This can be done in two ways: (1) **offline optimization** where we pre-compute optimal paths for different scenarios using objectives like Wasserstein length minimization in Eq. (22), and (2) **online adaptation** where paths are dynamically adjusted during generation based on intermediate results or user feedback.\n\nThe **\"bridge distributions\"** concept is now formally defined as finding paths ($\\alpha_t$, $\\beta_t$) that optimally transport probability mass from an initial distribution $\\mu_{\\alpha_0,\\beta_0}$ to a target distribution $\\mu_{\\alpha_1,\\beta_1}$, where optimality is measured by criteria such as the Wasserstein distance or other transport costs. Unlike fixed diffusion schedules, our framework allows these transport paths to be optimized post-training for specific applications.\n\n**Concrete examples** include: (1) adaptive mask scheduling for inpainting where $\\alpha_t$ is dynamically updated based on reconstruction quality, (2) multi-resolution generation where paths are optimized to generate coarse features first, then fine details, and (3) interactive editing where user feedback guides path selection in real-time.\n\n>***Posterior sampling limitation:***\n\nWhile you correctly note the quadratic structure requirement, this represents a large and practically important class, particularly in fine-tuning applications (Gaussian priors, L2 regularization, quadratic rewards). The key insight is that to exploit this class fully, it is important to use operator-based (as opposed to scalar) α and β. **The unified treatment allows the same trained model to handle both arbitrary inpainting AND posterior sampling with theoretical guarantees, which previous methods cannot achieve without retraining or biased approximations.**\n\n>***Larger-scale experiments:***\n\n Since submission, we have conducted extensive experiments on CelebA (128×128) and FFHQ (256×256) datasets to demonstrate scalability and  benchmark our approach against different methods in various image restoration tasks, evaluating the average Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index (SSIM) on 100 test images from each dataset. To assess the performance of our methodology using Hadamard product-based operators, we employed two types of masking: square masks of sizes $40\\times40$ and $80\\times80$ with added Gaussian noise of standard deviation 0.05, and random masks covering $70\\%$ of image pixels with Gaussian noise of standard deviation 0.01.\n\nThe results reported in the table below show that our method consistently ranks either first or second in both reconstruction metrics across all tasks and datasets. Regarding visual quality, our method generates realistic, artifact-free images.\n\n**CelebA Dataset (128×128)**\n| Method     | Random PSNR | Random SSIM | Block PSNR | Block SSIM |\n|------------|-------------|-------------|------------|------------|\n| Degraded   | 13.35       | 0.234       | 21.50      | 0.744      |\n| OT-ODE [1]     | 28.84       | 0.838       | 23.88      | 0.874      |\n| D-Flow [2]    | 31.37       | 0.888       | 26.69      | 0.833      |\n| Flow-Priors [3]| 31.76       | 0.909       | 25.85      | 0.822      |\n| PnP-Flow [4]  | 32.98       | 0.930       | 26.87      | 0.904      |\n| Ours       | **33.11**   | **0.945**   | **26.96**  | **0.914**  |\n\n**AFHQ-Cat Dataset (256×256)**\n| Method     | Random PSNR | Random SSIM | Block PSNR | Block SSIM |\n|------------|-------------|-------------|------------|------------|\n| Degraded   | 11.82       | 0.197       | 22.12      | 0.742      |\n| OT-ODE [1]    | 28.36       | 0.865       | 28.84      | 0.914      |\n| D-Flow [2]    | 33.07       | 0.938       | 29.70      | 0.893      |\n| Flow-Priors [3]| 32.33       | 0.945       | 29.40      | 0.858      |\n| PnP-Flow [4]  | 33.54       | 0.953       | **30.59**  | **0.943**  |\n| Ours       | **33.76**   | **0.967**   | 29.98      | 0.938      |\n\n\nThese experiments indicate that our operator-based interpolants maintain their versatility at scale, and confirm that our approach scales effectively while preserving the key advantage of training once and adapting at inference time.\n\nIn addition, we have preliminary results on a **robotics maze planning experiment**, as is done in [5]. In this setting, given starting and terminal coordinates in a maze, the aim is to construct a path between the points along the maze without hitting walls or obstacles. Our method allows the zero-shot generation of this path for any pair of points in the maze without re-training. In our preliminary results, we see a 100% success rate in this task, and are also in direct communication with the authors of [5] to produce their reward metrics, if that is of interest to the reviewers.\n\n>***Technical corrections:***\n\n All typos and formatting errors have been corrected.\n\n**References:**\n\n[1] A. Pokle, M. Muckley, R. Chen, B. Karrer, \"Training-free linear image inverses via flows\" (TMLR)\n\n[2] H. Ben-Hamu, O. Puny, I. Gat, B. Karrer, U. Singer, Y. Lipman, \"D-Flow: Differentiating through Flows for Controlled Generation\" \n\n[3] Y. Zhang, P. Yu, Y. Zhu, Y. Chang, F. Gao, Y. Nian Wu, O. Leong, \"Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching\" (Neurips 2024)\n\n[4] S. Martin, A. Gagneux, P. Hagemann, G. Steidl, \"PnP-Flow: Plug-and-Play Image Restoration with Flow Matching\" (ICLR 2025)\n\n[5] B. Chen, D. Marti Monso, Y. Du, M. Simchowitz, R. Tedrake, V. Sitzmann \"Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion\" (NeurIPS 2024)"}, {"review_id": "yGCcQHouDb", "reviewer": "Reviewer_KoX6", "summary": "This paper introduces a generalization of stochastic interpolant-based generative models by replacing the scalar time variable with operator-valued paths (vectors, matrices, or operators). This allows a single model to flexibly handle multiple generative tasks—such as inpainting, conditional generation, and posterior sampling—without retraining. The approach is supported by a solid theoretical framework and experiments on MNIST and a lattice field model.", "questions": "See weakness.", "limitations": "See weakness.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strength:**\n\nThis paper introduces a theoretically grounded approach to universal multitask generative modeling using operator-based interpolants. It clearly demonstrates how to select appropriate operators for different tasks, with all claims rigorously supported by detailed proofs. Furthermore, the method’s zero-shot generalization ability is demonstrated on both MNIST and physics datasets.\n\n**Weakness:**\n\nThe theoretical results in this paper are elegant. But the experimental cases are limited to on moderate-size datasets (MNIST, φ⁴). For more complex datasets, a key challenge lies in selecting an operator that is both stable to train and effective. But the paper didn't involve the discussion.", "was_revised": true, "final_justification": "The authors solve my main concerns.", "rebuttal": "Thank you for recognizing the theoretical contributions and experimental demonstrations. We want to emphasize that our approach represents a **paradigmatic shift from time-based diffusion to operator-based interpolation**, fundamentally reconceptualizing how multitask generation can be achieved with **guaranteed unbiased sampling**.\n\nShould you find our reply satisfactory, we kindly invite you to reconsider your rating.\n\n>***Larger-scale experiments:***\n\n Since submission, we have conducted extensive experiments on CelebA (128×128) and FFHQ (256×256) datasets to demonstrate scalability and  benchmark our approach against different methods in various image restoration tasks, evaluating the average Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index (SSIM) on 100 test images from each dataset. To assess the performance of our methodology using Hadamard product-based operators, we employed two types of masking: square masks of sizes $40\\times40$ and $80\\times80$ with added Gaussian noise of standard deviation 0.05, and random masks covering $70\\%$ of image pixels with Gaussian noise of standard deviation 0.01.\n\nThe results reported in the table below show that our method consistently ranks either first or second in both reconstruction metrics across all tasks and datasets. Regarding visual quality, our method generates realistic, artifact-free images.\n\n**CelebA Dataset (128×128)**\n| Method     | Random PSNR | Random SSIM | Block PSNR | Block SSIM |\n|------------|-------------|-------------|------------|------------|\n| Degraded   | 13.35       | 0.234       | 21.50      | 0.744      |\n| OT-ODE [1]     | 28.84       | 0.838       | 23.88      | 0.874      |\n| D-Flow [2]    | 31.37       | 0.888       | 26.69      | 0.833      |\n| Flow-Priors [3]| 31.76       | 0.909       | 25.85      | 0.822      |\n| PnP-Flow [4]  | 32.98       | 0.930       | 26.87      | 0.904      |\n| Ours       | **33.11**   | **0.945**   | **26.96**  | **0.914**  |\n\n**AFHQ-Cat Dataset (256×256)**\n| Method     | Random PSNR | Random SSIM | Block PSNR | Block SSIM |\n|------------|-------------|-------------|------------|------------|\n| Degraded   | 11.82       | 0.197       | 22.12      | 0.742      |\n| OT-ODE [1]    | 28.36       | 0.865       | 28.84      | 0.914      |\n| D-Flow [2]    | 33.07       | 0.938       | 29.70      | 0.893      |\n| Flow-Priors [3]| 32.33       | 0.945       | 29.40      | 0.858      |\n| PnP-Flow [4]  | 33.54       | 0.953       | **30.59**  | **0.943**  |\n| Ours       | **33.76**   | **0.967**   | 29.98      | 0.938      |\n\n\nThese experiments indicate that our operator-based interpolants maintain their versatility at scale, and confirm that our approach scales effectively while preserving the key advantage of training once and adapting at inference time.\n\nIn addition, we have preliminary results on a **robotics maze planning experiment**, as is done in [5]. In this setting, given starting and terminal coordinates in a maze, the aim is to construct a path between the points along the maze without hitting walls or obstacles. Our method allows the zero-shot generation of this path for any pair of points in the maze without re-training. In our preliminary results, we see a 100% success rate in this task, and are also in direct communication with the authors of [5] to produce their reward metrics, if that is of interest to the reviewers.\n\n>***Operator selection complexity:***\n\n This is indeed a key practical challenge. We have added extensive discussion and empirical guidelines for operator selection, including: (1) stability analysis for different operator classes, (2) empirical scaling laws for training complexity vs. operator diversity, (3) practical heuristics for choosing appropriate operator spaces for different application domains. The CelebA/FFHQ experiments provide concrete examples of operator selection at scale. **Importantly, once selected, our operator-based framework provides exact sampling without the approximation biases that plague existing methods.**\n\n>***Implementation details:***\n\n We have expanded the algorithmic section with detailed implementation guidelines, stability considerations, and practical tips for different operator types based on our large-scale experiments. Unlike existing multitask approaches that require complex approximation strategies, our mathematically rigorous framework provides clear implementation paths with theoretical guarantees.\n\n\n>***Key conceptual contribution:***\n\n While existing multitask methods still rely on traditional diffusion in time concepts with heterogeneous processes, our work represents a fundamental reconceptualization - treating all generative tasks as different geometric traversals through the same underlying operator space. This enables **theoretically exact, unbiased sampling** for all tasks without the systematic errors inherent in current approximation-based approaches.\n\n**References:**\n\n[1] A. Pokle, M. Muckley, R. Chen, B. Karrer, \"Training-free linear image inverses via flows\" (TMLR)\n\n[2] H. Ben-Hamu, O. Puny, I. Gat, B. Karrer, U. Singer, Y. Lipman, \"D-Flow: Differentiating through Flows for Controlled Generation\" \n\n[3] Y. Zhang, P. Yu, Y. Zhu, Y. Chang, F. Gao, Y. Nian Wu, O. Leong, \"Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching\" (Neurips 2024)\n\n[4] S. Martin, A. Gagneux, P. Hagemann, G. Steidl, \"PnP-Flow: Plug-and-Play Image Restoration with Flow Matching\" (ICLR 2025)\n\n[5] B. Chen, D. Marti Monso, Y. Du, M. Simchowitz, R. Tedrake, V. Sitzmann \"Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion\" (NeurIPS 2024)"}]

Author rebuttal:

Meta-review: This manuscript proposed a framework of operator-based interpolants, that aims to learn maps between probability distributions instead of flows as in existing flow-based generative models. Such framework enables construction of versatile models capable of multi tasks. The approaches proposed in this manuscript are sound in theory and also achieving nice empirical results in comprehensive numerical validation. The reviewers unanimously recommend acceptance after the discussion period. Thus the meta-reviewer recommends acceptance of the manuscript.
Decision: Accept (spotlight)