Stochastic interpolants with data-dependent couplings

1. Introduction

• A simple Gaussian base is often chosen for convenience, but represents an absence of prior knowledge about problem structure; adapted bases remain underexplored.Some generative-model formulations require a Gaussian base, while others permit more general choices.Source: §1 · p. 1
• The velocity field can be learned with an efficient, simulation-free square-loss regression problem analogous to standard data-agnostic interpolant and flow-matching algorithms.This learning approach is presented as part of the general stochastic-interpolant formulation.Source: §1 · pp. 1–2
• The framework constructs base samples conditionally on target samples, while allowing separate class or continuous-embedding information to condition the velocity field.The paper emphasizes that data-dependent coupling and conditioning are distinct, and reports improved performance in super-resolution and in-painting when both are used.Source: §1 · pp. 1–2
• The authors characterize transport under conditional, data-dependent generation and analyze its effects on the time-dependent density.The framework identifies a continuity equation obeyed by that density.Source: §1 · pp. 1–2
• The paper studies data-dependent bases and conditional information in image in-painting and super-resolution, reporting performance gains in both applications.The experiments focus on ImageNet.Source: §1 · p. 2

Couplings

• Prior coupling strategies include minibatch entropic optimal transport, learned conditional bases, and Schrödinger bridges, which can require costly iterative SDE solves.The paper notes that minibatch coupling estimates may become uninformative on large datasets, while learned conditional bases may not match the generation-time marginal.Source: §2 · p. 2
• The paper presents its approach as simpler and more flexible than related diffusion-bridge formulations, separating density bridging from generative-model construction.It identifies Liu et al. (2023a) and Somnath et al. (2023) as operationally similar but less general.Source: §2 · p. 2
• Figure 2 contrasts simple trajectories from data-dependent coupling with factorized class-conditioned flows and the complex transport of independent pairing.The comparison uses a three-mode Gaussian mixture transport problem.Source: Figure 2 · p. 3
• Standard flows and diffusions commonly use independent base-target pairings, whereas this framework supports couplings individualized to each generative task.Table 1 also describes learned conditional-base and minibatch optimal-transport approaches.Source: Table 1 · p. 3, §2 · p. 3

Conditional Diffusions and Flows for Images

• Conditional image-generation work includes diffusion methods for super-resolution and in-painting, while this paper revisits these tasks from a velocity-modeling perspective.Prior in-painting methods replace known coordinates during integration, with some work addressing approximation error using sequential Monte Carlo.Source: §2 · p. 3

3. Stochastic interpolants with couplings

• The framework replaces independent endpoint sampling with a joint coupling of base and target samples, while retaining a stochastic interpolant between their endpoint densities.External labels or embeddings may also condition the quantities in the framework.Source: §3 · p. 4

3.1. Transport equations and conditional generative models

• The interpolant density evolves from ρ0 to ρ1 under a transport equation, whose velocity and score are characterized by quadratic objectives with unique minimizers.The framework yields both deterministic probability-flow ODEs and stochastic generative models.Source: §3.1 · pp. 4–5

3.2. Designing data-dependent couplings

• A task-dependent base can be formed as m(x1) plus Gaussian noise, where m(x1) may be noisy, partial, or low-resolution and the noise design remains selectable.When m(x1) is deterministic, this conditional base is Gaussian with mean m(x1) and covariance σσ⊤.Source: §3.2 · p. 5

3.3. Reducing transport costs via coupling

• With m(x1)=x1, γt=0, and linear interpolation, the coupled transport-cost bound has left side dσ², below the independent bound 2E[|x1|²]+dσ².The paper states that designed couplings can lower the upper bound at no extra cost.Source: §3.1 · p. 5, §3.3 · p. 6

3.4. Learning and Sampling

• The velocity and score are approximated with neural networks by minimizing empirical versions of their objectives using stochastic gradient descent.Sampling initializes from the base, then numerically evolves the probability-flow ODE or forward SDE to produce target samples.Source: §3.4 · p. 6

4. Numerical experiments

• Experiments explore data-dependent couplings for conditional image generation and find the framework straightforward to scale directly in pixel space to high-resolution images, including comparisons across 64×64 to 256×256 and 256×256 to 512×512 resolutions.Source: §4 · p. 6, §4.2 · p. 7

4.1. In-painting

• Inpainting keeps observed pixels fixed and initializes masked pixels with independent Gaussian noise; randomly selected tiles form the training masks.The model can enforce fixed observed pixels by masking velocity outputs, without inference-time replacement or MCMC corrections.Source: §4.1 · pp. 6–7

4.1. In-painting; 4.2. Super-resolution on Imagenet

• Inpainting samples provide realistic fills conditioned on observed content, while super-resolution samples visibly increase resolution from 64×64 to 256×256.The inpainting aim is conditionally valid filling, not recovery of the reference image's precise content.Source: §4.1 · p. 7, §4.2 · p. 7

4.2. Super-resolution on Imagenet

• For super-resolution, the base is the upsampled low-resolution image plus Gaussian noise, and the low-resolution image conditions the model throughout transport.Positive noise smooths the base over the higher-dimensional image space; with zero noise, it would lie on a lower-dimensional manifold.Source: §4.2 · p. 7

5. Discussion, challenges, and future work

• The framework constructs data-dependent couplings between base and target densities within the stochastic interpolant formalism.The inpainting examples aim for conditionally valid fills rather than exact recovery of the reference image.Source: §5 · p. 7, Figure 3 · p. 8
• The paper suggests choosing a Gaussian base whose mean and covariance adapt to target samples, and reports using such couplings in inpainting and super-resolution.The super-resolution illustration compares low-resolution input, model output, and ground truth.Source: §5 · p. 8, Figure 4 · p. 8
• Potential applications include fixing a chemical backbone while varying functional groups, and correcting errors from variational autoencoders.For the latter, the paper proposes using inputs as the target density and autoencoder outputs as the base density.Source: §5 · pp. 8–9

Conditioning setup

• The setup allows ξ to represent a discrete class or continuous embedding, with data distributed according to ρ1(x1|ξ)η(ξ).The interpolant uses a conditionally drawn pair (x0, x1|ξ) and independent standard Gaussian noise.Source: A. Omitted proofs with conditioning… · p. 12

Velocity and score objectives

• The conditional velocity and score are defined through expectations given the interpolant state and ξ, and are unique minimizers of their respective objectives.The objectives can be estimated from samples of the coupled pair, noise, and conditioning variable.Source: A. Omitted proofs with conditioning… · pp. 12, 14

Transport equation with coupling and conditioning

• With conditioning ξ, the interpolant density has endpoints ρ0(x|ξ) and ρ1(x|ξ) and solves the transport equation with velocity b_t(x, ξ).The score identity holds when γ_t is nonzero.Source: A. Omitted proofs with conditioning… · p. 12

Probability flow and diffusions

• The transport characterization extends to probability-flow ODEs and forward and backward SDEs with coupling and conditioning.The appendix focuses on the deterministic probability-flow ODE for simplicity.Source: A. Omitted proofs with conditioning… · p. 14

Control of transport cost

• The appendix also gives a proposition controlling the transport cost of the probability-flow ODE.Its proof invokes Jensen's inequality and the tower property of conditional expectation.Source: A. Omitted proofs with conditioning… · p. 15

• The velocity model uses a U-Net with class-label embeddings, and image-shaped conditioning appends upsampled low-resolution images and in-painting masks to the input at each time step.Source: B. Further experimental details · p. 15
• The U-Net configuration uses 8 ResNet block groups, learned sinusoidal conditioning with dimension 32, attention head dimension 64, and no random Fourier features.Source: B. Further experimental details · p. 15
• Optimization uses Adam at learning rate 2e-4 with StepLR scaling by 0.99 every 1,000 steps, no weight decay, and gradient-norm clipping at 10,000.Source: B. Further experimental details · p. 15
• Sampling uses the Dopri solver from torchdiffeq, with PyTorch and Lightning Fabric used for parallelism.Source: B. Further experimental details · p. 15
• Additional illustrations show 256 × 256 in-filling, 64-to-256 super-resolution, and 256-to-512 super-resolution examples, including temporal probability-flow slices for the first two.Source: B. Further experimental details · p. 15, Figure 5 · p. 16, Figure 6 · p. 17