---
title: Interpolative Multi-Marginal Flow Matching (IMMFM)
url: https://www.emergentmind.com/topics/interpolative-multi-marginal-flow-matching-immfm
type: topic
---

# Interpolative Multi-Marginal Flow Matching (IMMFM)

Interpolative Multi-Marginal Flow Matching (IMMFM) is a conditional flow/score matching framework for learning a continuous stochastic process from sparse, irregularly sampled longitudinal observations. In its published formulation, IMMFM targets settings in which each subject is observed at only a few non-uniform time points, the observations are high-dimensional, and the objective is to recover a continuous-time family of marginals while retaining subject-specific dynamics. Its defining construction is a conditional Gaussian probability path whose mean is a piecewise-quadratic interpolant through multiple observed states and whose variance follows a Brownian-bridge-like schedule; neural networks are then trained to match the corresponding probability-flow drift, score, and a learned diffusion coefficient [2510.03569].

## 1. Problem setting and intended scope

IMMFM is formulated for trajectory snippets
\[
z_{1:M}=(x_{t_1},\ldots,x_{t_M}),
\]
with non-uniform times
\[
t_0 < t_1 < \cdots < t_M \in [0,1], \qquad x_{t_i}\in \mathbb{R}^D.
\]
The paper assumes that the data lie on or near a lower-dimensional smooth manifold embedded in ambient space, and denotes by \(\rho_i\) the marginal distribution at observation time \(t_i\). The learning target is a continuous-time family of marginals \(\{p_t\}_{t\in[0,1]}\) such that
\[
p_{t_i}=\rho_i \quad \text{for all } i.
\]
The intended application is longitudinal modeling rather than one-step prediction: the goal is to infer a continuous stochastic process that explains the entire trajectory across time, not merely a single transition between adjacent visits [2510.03569].

The paper motivates this formulation by arguing that pairwise modeling is inadequate in sparse longitudinal settings. Independent pairwise fits can lose global trajectory coherence, ignore dependencies beyond adjacent visits, become highly ambiguous when observations are far apart, and induce velocity discontinuities when local interpolants are chained. IMMFM therefore adopts a multi-marginal viewpoint in which multiple observed time points jointly constrain the learned process. In the paper’s usage, “multi-marginal” is narrower than a fully general multi-marginal optimal transport or Schrödinger bridge formulation: it refers to a conditional interpolation path on each interval that depends on more than its two endpoints, rather than to a globally optimized coupling over all times [2510.03569].

## 2. Stochastic process formulation and probability-flow decomposition

IMMFM models latent dynamics as an Itô SDE
\[
d x_t = u_t(x_t)\,dt + g(t,x_t)\,dW_t,
\]
where \(u_t\) is the drift, \(g\) is the diffusion coefficient, and \(W_t\) is standard Brownian motion. The corresponding marginals satisfy the Fokker–Planck equation
\[
\partial_t p_t(x) = -\,\nabla \cdot \big(p_t\,u_t\big) +\frac{1}{2}\,\Delta\!\big(g(t, x_t)^2\,p_t\big).
\]
To obtain a simulation-free training objective, the framework introduces a deterministic probability-flow ODE with the same marginals,
\[
\partial_t p_t = -\nabla \cdot (p_t u_t^\circ),
\]
and uses the identity
\[
u_t(x_t) = \underbrace{v_t(x_t) + \frac{1}{2}\nabla g(t, x_t)^2}_{\text{Prob. flow drift }u_t^\circ(x_t)} + \frac{g(t,x_t)^2}{2} \nabla \log p_t(x_t).
\]
In implementation, the term \(\tfrac12 \nabla g^2\) is absorbed into the learnable drift. The learned objects are therefore
\[
v_\theta(t,x,c) \approx u_t^\circ(x), \qquad
s_\theta(t,x,c) \approx \nabla \log p_t(x), \qquad
g_\theta(t,x,c),
\]
and the SDE drift used at generation time is assembled as
\[
u_\theta(t,x,c)=v_\theta(t,x,c)+\frac{g_\theta(t,x,c)^2}{2}s_\theta(t,x,c).
\]
The conditioning variable \(c\) is subject context; in the forecasting procedure it is taken to be the previous observation \(x_{t_{i-1}}\) when available [2510.03569].

This decomposition places IMMFM close to simulation-free Schrödinger bridge and stochastic flow-matching lines of work, but its stated emphasis is longitudinal trajectory modeling with explicit subject conditioning. The decomposition also makes clear that the framework is stochastic at the model level even though training avoids backpropagating through full SDE trajectories.

## 3. Piecewise-quadratic multi-marginal interpolation path

The defining object of IMMFM is a conditional Gaussian path
\[
p_t(x\mid z) = \mathcal{N}\!\left(x\bigm|\mu_t(z),\;\sigma^2(t)I\right).
\]
As in standard conditional flow matching, this yields analytic conditional targets:
\[
u^\circ_t(x \mid z) = \frac{\sigma'(t)}{\sigma(t)}(x - \mu_t(z)) + \mu'_t(z),
\]
and
\[
\nabla_x \log p_t(x \mid z) = \frac{\mu_t(z)-x}{\sigma^2(t)}.
\]

The novelty lies in the mean path. For \(t\in[t_i,t_{i+1}]\), IMMFM defines a piecewise-quadratic mean
\[
\mu_t(z) = x_{t_i} + v_i (t - t_i) + \tfrac{1}{2}\alpha_t (v_i - v_{i+1})(t - t_i),
\]
with
\[
v_i = \frac{x_{t_{i+1}} - x_{t_i}}{t_{i+1} - t_i}, \qquad
\alpha_t = \frac{t_{i+1}-t}{t_{i+1}-t_i}.
\]
The appendix gives the equivalent form
\[
\mu_t = x_i + v_i(t - t_i) + \frac{1}{2}(v_i - v_{i+1})\frac{(t - t_i)(t_{i+1} - t)}{t_{i+1} - t_i}.
\]
Its variance schedule is Brownian-bridge-like,
\[
\sigma(t) = \sigma_0 (t - t_i)\,\alpha_t
\]
or equivalently
\[
\sigma(t)=\sigma_0\cdot\frac{(t-t_i)(t_{i+1}-t)}{t_{i+1}-t_i},
\]
with derivative
\[
\sigma'(t)=\sigma_0(2\alpha_t-1).
\]

The path is called multi-marginal because the interpolation on \([t_i,t_{i+1}]\) depends not only on \(x_{t_i}\) and \(x_{t_{i+1}}\), but also on the next segment velocity \(v_{i+1}\), which depends on \(x_{t_{i+2}}\). Its derivative is
\[
\mu'_t = v_i + \frac{1}{2}(v_i - v_{i+1})(2\alpha_t - 1),
\]
so the conditional target drift becomes
\[
u^\circ_t(x \mid z) =
v_i + \frac{1}{2}(v_i - v_{i+1})(2\alpha_t - 1)
+ \frac{\sigma'(t)}{\sigma(t)}\bigl(x - \mu_t(z)\bigr).
\]
The stated purpose is to avoid the discontinuous velocities induced by concatenated linear interpolants and to encode local curvature using more than two observed points [2510.03569].

This notion of “multi-marginal” is more localized than the conditioning used in some adjacent methods. In Momentum Multi-Marginal Schrödinger Bridge Matching, the conditional acceleration on an early segment depends explicitly on all future pinned positions in a phase-space bridge law [2506.10168]. In OT-potential FM, a single continuous-time flow is softly steered toward all observed intermediate marginals through time-localized optimal-transport potentials [2606.05327]. IMMFM instead uses a lighter-weight analytic construction: on each interval, the path incorporates the next segment’s velocity and therefore one additional future observation, without introducing a global multi-time bridge or a forced Hamilton–Jacobi system.

## 4. Learning objective, diffusion estimation, and stability condition

The intractable ideal objective over the true marginals is replaced by a tractable conditional objective on the constructed path:
\[
\mathcal{L}_{\rm CSDE}(\theta) =
\mathbb{E}_{\substack{t \sim \mathcal{U}(0,1)\\ z \sim q\\ x \sim p_t(x \mid z)}}
\Bigl[
\bigl\lVert v_\theta(t,x,c) - u^\circ_t(x \mid z) \bigr\rVert_2^2
+\lambda(t)^2 \bigl\lVert s_\theta(t,x,c) - \nabla_x \log p_t(x \mid z) \bigr\rVert_2^2
\Bigr].
\]
The score weight is chosen as
\[
\lambda(t)=\frac{2\sigma(t)}{\sigma_0^2}.
\]
Training is therefore simulation-free in the standard conditional flow/score matching sense: one samples \(t\), a trajectory snippet \(z\), and a noisy point \(x\sim p_t(x\mid z)\), then regresses onto analytic conditional targets [2510.03569].

IMMFM augments this with a learned diffusion coefficient. The uncertainty objective is
\[
\mathcal{L}_{\rm uncertainty}(\theta) =
\mathbb{E}_{t, z, x}
\Biggl[
\Biggl\lVert
g_\theta(t,x_t,c)^2 -
\underbrace{\lVert x_i + (t_{i+1}-t_i)u_\theta(t,x_t,c) - x_{t_{i+1}} \rVert_2^2}_{\text{Squared error of predictive construction}}
\Biggr\rVert_2^2
\Biggr],
\]
and the full loss is
\[
\mathcal{L}_{\rm IMMFM}(\theta) =
\mathcal{L}_{\rm CSDE}(\theta) + \beta \mathcal{L}_{\rm uncertainty}(\theta).
\]
The paper states that \(g:[0,1]\times\mathbb R^d\to\mathbb R_+\), that \(g_\theta\) outputs multiple scalars, and that the practical interpretation is closer to scalar or diagonal diffusion than to a full covariance model. Positivity is enforced implicitly by using \(g_\theta^2\) in both the drift assembly and the uncertainty loss [2510.03569].

The paper’s main theoretical claim is a stationary-point result. Under mild regularity conditions, every stationary point of the core SDE objective \(\mathcal L_{\mathrm{SDE}}\) is also a stationary point of \(\mathcal L_{\mathrm{IMMFM}}\). The supporting lemma defines the residual
\[
r_\theta(t,x,z) = x_i + \Delta t\,u_\theta(t,x,c) - x_{t_{i+1}}, \qquad \Delta t := t_{i+1}-t_i,
\]
and states that if \(u_\theta(t,x,c)=u_t(x)\), then
\[
\mathbb{E}_{z|x,t}[\,r_\theta(t,x,z)\,]=0.
\]
The significance of this proposition is limited but specific: the residual-based diffusion objective does not alter the stationary solutions of the original drift/score learning problem. This suggests that uncertainty estimation is introduced as an auxiliary stochasticity mechanism rather than as a competing transport criterion [2510.03569].

## 5. Coupling construction, conditioning, inference, and implementation

Training uses sampled trajectory snippets together with an offline coupling construction. The paper states that real datasets rarely provide full trajectories with exact pairings across marginal sets, so it constructs the trajectory coupling \(q(z)\) via multi-marginal OT with additive cost
\[
C(x_{t_0}, \ldots, x_{t_M}) = \sum_{i=0}^{M-1} k(x_{t_i},x_{t_{i+1}}).
\]
With pairwise-additive cost, the problem decomposes into pairwise OT. Using squared Euclidean cost after spatial alignment,
\[
\pi_{i,i+1}^{*} = \argmin_{\pi \in \Pi(\rho_i^{\dagger}, \rho_{i+1}^{\dagger})} \int \|x_{t_i} - x_{t_{i+1}}\|_2^2 \, d\pi(x_{t_i},x_{t_{i+1}}),
\]
and the resulting joint coupling is
\[
q(z) = \pi^{*}(x_{t_0}, \ldots, x_{t_M})
= \frac{\prod_{i=0}^{M-1} \pi_{i,i+1}^{*}(x_{t_i}, x_{t_{i+1}})}{\prod_{i=1}^{M-1} \rho_i^{\dagger}(x_{t_i})}.
\]
The paper describes this MMOT coupling as an offline preprocessing step reused during training [2510.03569].

The training loop samples a mini-batch of trajectories, samples \(t\sim \mathcal U(t_0,t_M)\), finds the active segment \([t_j,t_{j+1})\), computes local velocities \(v_j\) and \(v_{j+1}\), constructs \(\mu_t\) and \(\sigma(t)\), samples
\[
x\sim \mathcal N(\mu_t,\sigma(t)^2 I),
\]
and then minimizes \(\mathcal L_{\rm CSDE}\). The conditioning variable is set to the previous frame,
\[
c \leftarrow x_{t_{j-1}},
\]
or to zero if \(j=0\). This previous-frame conditioning is central to the paper’s claim of subject specificity: it is said to provide an implicit initial velocity and to distinguish different subject-specific evolution patterns [2510.03569].

At inference, IMMFM supports both deterministic and stochastic rollout. In ODE mode,
\[
x_{\text{new}} = x + u_t \Delta t, \qquad u_t = v_\theta(t,x,c),
\]
with Euler integration. In SDE mode,
\[
u_t = v_\theta(t,x,c)+\frac{g_\theta(t,x,c)^2}{2}s_\theta(t,x,c),
\]
\[
x_{\text{new}} = x + u_t \Delta t + g_t \sqrt{\Delta t}\mathbf z, \qquad \mathbf z\sim \mathcal N(0,I),
\]
with Euler–Maruyama integration. The paper evaluates three variants: O-IMMFM, S-IMMFM, and SU-IMMFM, corresponding respectively to deterministic ODE rollout, SDE rollout with fixed diffusion, and SDE rollout with learned uncertainty-driven diffusion [2510.03569].

The implementation is organized around a pretrained autoencoder and a U-ViT-based flow regressor. The autoencoder is U-Net-based with skip connections from encoder to decoder removed to enforce compression. The flow regressor outputs \(v_\theta\), \(s_\theta\), and \(g_\theta\), uses sinusoidal time embeddings, and conditions on the current latent, the previous latent \(c\), and optionally projected non-imaging covariates. Reported training cost is one NVIDIA A100 40GB, roughly 12–18 hours for IMMFM itself, and around 3000 GPU hours for the full project pipeline [2510.03569].

## 6. Empirical evaluation and reported behavior

The paper evaluates IMMFM on synthetic 2D Gaussian trajectories, Starmen, ADNI, Brain MS, and Brain GBM. Baselines include ImageFlowNet, \(I^2SB\), TFM, L-MMFM, and M-MMFM, alongside the IMMFM variants O-IMMFM, S-IMMFM, and SU-IMMFM. The main image-level metrics are PSNR, SSIM, and MSE; clinically localized downstream evaluation uses DSC and HD [2510.03569].

The reported quantitative pattern is that IMMFM improves over prior longitudinal baselines, with the strongest gains in clinically meaningful ROI metrics and in settings where stochasticity and sparse sampling matter. On Starmen, IMMFM variants achieve best MSE \(0.09\), compared with \(0.11\) for L-MMFM and \(0.12\) for M-MMFM. On ADNI, SU-IMMFM reports PSNR \(37.52\), SSIM \(0.97\), MSE \(0.02\), DSC \(0.92\), and HD \(6.50\); on Brain MS it reports PSNR \(36.67\), SSIM \(0.95\), MSE \(0.01\), DSC \(0.73\), and HD \(21.03\); on Brain GBM it reports PSNR \(31.94\), SSIM \(0.93\), MSE \(0.01\), DSC \(0.46\), and HD \(135.08\). The paper summarizes its overall improvements as \(1\)–\(4.4\%\) Dice, \(1.5\)–\(2.2\) dB PSNR, and \(1.2\)–\(4.5\%\) SSIM [2510.03569].

Ablations identify the piecewise-quadratic path as the dominant component. The reported gains from the quadratic path reach \(+3.7\%\) Dice and \(+2.1\) dB PSNR; previous-frame conditioning contributes up to \(+2.1\%\) Dice and \(+1.09\) dB PSNR in ADNI; learned diffusion helps most on difficult stochastic data, especially GBM, with Dice improved by \(+1.5\%\) and HD reduced by over \(6.29\). The paper treats these ablations as evidence that the interpolation path is a substantive modeling choice rather than a cosmetic reparameterization [2510.03569].

The downstream prognosis result in ADNI is unusually prominent. Using observed ventricular area at the second visit yields \(71.7\%\) classification accuracy for AD versus cognitively normal subjects, whereas using the IMMFM-forecasted trajectory to the last visit yields \(80.8\%\). The paper presents this as evidence that forecasted longitudinal trajectories sharpen clinically meaningful group separation [2510.03569].

## 7. Relationship to adjacent multi-marginal methods, misconceptions, and limitations

IMMFM occupies one point within a broader family of multi-time flow and bridge methods. Relative to Momentum Multi-Marginal Schrödinger Bridge Matching, it is less globally coupled: 3MSBM lifts dynamics to phase space, conditions on several future points, learns an acceleration field for a controlled diffusion, and claims globally optimal coupling at convergence together with preservation of intermediate marginals throughout training [2506.10168]. Relative to OT-potential FM, it is more explicit about subject-specific longitudinal conditioning but less global in its multi-time constraint: OTP-FM learns a single continuous-time flow between endpoints and softly steers it toward intermediate marginals through OT-based potential terms, with a hard-potential limit that recovers piecewise CFM [2606.05327]. Relative to SplineFlow and MMSFM, IMMFM uses a simpler analytic path construction: a piecewise-quadratic interpolant rather than a global B-spline through all observations or OT-coupled transport splines over overlapping windows [2601.23072] [2508.04351].

A recurring misconception is to read “multi-marginal” in IMMFM as a claim of globally conditioned all-time interpolation. The paper’s actual construction is more local. On interval \([t_i,t_{i+1}]\), the path depends on \(v_i\) and \(v_{i+1}\), hence on \(x_{t_i}\), \(x_{t_{i+1}}\), and \(x_{t_{i+2}}\); this is a multi-observation path, but not a globally optimized multi-time bridge. That distinction matters when comparing IMMFM with methods whose conditioning spans all future pinned points or whose action contains explicit potentials for every observed intermediate marginal.

The paper also states several limitations. Performance depends on high-quality data and on the coverage and diversity of training trajectories. Severe input artifacts can degrade prediction quality. Long-horizon forecasting deteriorates, especially for highly unpredictable diseases such as GBM. The latent representation could be improved by temporally aware encoders, and future directions mentioned in the paper include biophysical constraints, multimodal covariates, and causal or counterfactual conditioning [2510.03569].

Within the contemporary literature, IMMFM is therefore best understood as a simulation-free, conditional, stochastic multi-time flow/score matching method whose central innovation is a piecewise-quadratic path that smooths local multi-visit interpolation while remaining tractable in high-dimensional longitudinal settings. Its multi-marginality is real but local; its stochasticity is explicit but implemented through conditional analytic targets; and its strongest empirical evidence is in subject-specific neuroimaging trajectories rather than in general-purpose multimarginal transport.

Source: https://www.emergentmind.com/topics/interpolative-multi-marginal-flow-matching-immfm