Generative Modeling from Black-box Corruptions via Self-Consistent Stochastic Interpolants

AI summary

Overview

Research area: Generative modeling / inverse problems — specifically, learning a generative model of clean data when only corrupted observations are available, using stochastic interpolants (SI) and transport-based generative models.

Technical level: Advanced. The paper combines stochastic differential equations, measure-theoretic inverse problems over probability spaces, Wasserstein and KL contraction analysis, and neural network training.

Scope: The paper introduces and analyzes Self-Consistent Stochastic Interpolants (SCSI), a fixed-point training scheme that inverts a black-box corruption channel at the level of distributions, and establishes Wasserstein and KL convergence guarantees under stated assumptions.

What This Paper Is About

Standard generative models (diffusion and flow models) require clean samples of the target distribution, but in many scientific and engineering settings only measurements corrupted by a noisy, ill-conditioned forward map are observed. The authors frame this as an inverse problem at the level of distributions — inverting the linear operator equation Kπ = μ — and propose an iterative method that learns a transport map from corrupted measurements back to clean data using only the corrupted dataset plus black-box access to the corruption channel. The goal is to recover the prior distribution π without ever observing clean samples, clean likelihoods, or gradients through the forward model.

Key Contributions

1. A self-consistent iterative training scheme (SCSI). The method alternates between constructing a stochastic interpolant from the current transport map through the black-box forward model, and retraining the velocity/denoiser fields on that interpolant, converging toward a self-consistent transport map that inverts the corruption channel.

2. Black-box applicability. Unlike prior inverse generative modeling work that assumes explicit, differentiable, often linear forward models (with rank conditions), SCSI requires only a simulator that produces samples from P(dy|x). This extends the method to nonlinear forward models (e.g., motion blur), non-differentiable operators (e.g., JPEG compression), and non-Gaussian noise (e.g., Poisson noise), with no backpropagation or posterior sampling through F.

3. Theoretical guarantees. The paper proves that π is the only admissible fixed point under injectivity of the channel operator K (Proposition 4), establishes linear Wasserstein contraction (Fact 6), and proves KL contraction under a Lipschitz stability assumption and a finite condition number (Theorem 9). It also proves a finite condition number for compact hypothesis classes (Proposition 8) and covers a class of Fokker-Planck channels.

4. Computational efficiency relative to variational alternatives, with public source code at https://github.com/modichirag/SCSI.

Main Findings

• Fixed-point uniqueness: If the integral operator K is injective at the distribution level and the iterative scheme converges to a fixed point Θ*, then π_{Θ*} = π (Proposition 4). Injectivity holds for the AWGN channel for any σ < ∞, for random projection channels (covering tomography and inpainting) under a spanning condition on the support of the projection distribution, and for compositions of injective channels such as Cryo-EM-adjacent channels (Remark 3).

• Wasserstein linear convergence: Under Assumption 5 (an R-Lipschitz stability condition on the map from a measure to its backward transport map), W₂²(π, π^(k)) ≤ R^k W₂²(π, π^(0)), giving exponential convergence whenever R < 1. Injectivity of the channel is shown to be necessary for R < 1. This bound applies to both the SDE setting (ε > 0) and the ODE setting (ε = 0).

• KL contraction with an error floor: Under Assumption 7 (KL Lipschitz stability of the drift function f_π = b_π + ε s_π) and with a finite regularized condition number χ, if L_ε χ < 4 then with η = 1 + L_ε/4 − χ⁻¹, the KL divergence satisfies KL(π‖π^(k)) ≤ 2η^k KL(π‖π^(0)) + O(ε⁻³) δ²/(1 − √η)², where δ bounds the drift/score approximation error at every iteration. Theorem 9 gives global convergence guarantees for a non-convex objective by replacing gradient descent with a Picard-type iteration.

• Condition number characterization: The condition number is defined as χ = sup_{ρ∈S_λ} KL(π‖ρ) / KL(Kπ‖Kρ), and by the data-processing inequality χ ≥ 1 always. The regularized inverse problem is non-singular when χ < ∞. Proposition 8 shows χ < ∞ when K is injective, the parameter space is compact with continuous parametrization, and π cannot be exactly represented by the model (a misspecified setting).

• Approximation–estimation tradeoff: The KL bound captures a tension: a smaller function class yields a smaller condition number (better contraction rate) but larger approximation error δ. The authors state that quantitative analysis of this tradeoff for specific function classes is beyond the scope of the work.

• Role of the diffusion coefficient: The learning-error term scales as O(ε⁻³), while ε also affects the Lipschitz constant L_ε. The authors explicitly state the KL contraction results do not cover the ε → 0 limit where transport is performed with the probability flow ODE.

• Empirical claims without reported numbers: The abstract states the method demonstrates superior performance on inverse problems in natural image processing and scientific reconstruction, but the provided paper content does not report specific datasets, benchmark names, or numerical results.

Methodology in Plain English

The central idea is to turn a consistency condition into a training loop. If a transport map correctly inverts the corruption channel, then pushing corrupted samples through it to get clean samples, and then pushing those clean samples back through the known corruption channel, should reproduce the original corrupted distribution. The method uses this as a fixed-point target rather than as a differentiable loss through the channel.

Each outer iteration builds a stochastic interpolant — a continuous path that linearly blends (with time-dependent schedules α_t, β_t, γ_t) a current estimate of the clean sample, its corruption passed through the black-box simulator, and independent Gaussian noise. This interpolant is sampleable using only corrupt data and the simulator. The inner loop then fits standard stochastic interpolant least-squares objectives: a velocity field b that predicts the time derivative of the interpolant, and a denoiser g (equivalently a score s) that predicts the injected noise. Samples are generated by integrating the reverse-time SDE (or the probability flow ODE when ε = 0) backwards from the observation distribution to the clean distribution.

The structure mirrors expectation-maximization, with an "E" step that constructs the interpolant from the current map and an "M" step that retrains the fields, so the paper calls the scheme self-consistent. For efficiency, the inner problem is not solved to convergence: parameters are initialized from the previous outer iterate and updated for T_tr gradient steps. When T_tr = 1, the scheme is equivalent to treating the interpolant as a constant with a stop-gradient applied.

The authors contrast this with model-based approaches: the method behaves more like model-free reinforcement learning, which learns by interacting with a simulator, than like model-based control, which requires explicit knowledge and differentiability of the physics. They also note that the concurrent work of reference [19] can be viewed as an instance of their scheme with a conditional stochastic interpolant that directly models the posteriors P^(k)(dx|y). Once the marginal prior π is recovered, downstream tasks follow by composing the transport map with standard procedures: training a generative model for μ and restoring its outputs, training a generative model directly on the restored dataset, or training a conditional transport model on pairs (x, F(x)) for posterior inference.

Why This Matters

Impact on research: The work extends transport-based generative modeling to regimes where clean data simply do not exist, while replacing restrictive assumptions on the forward operator (linearity, differentiability, known likelihood) with black-box access. It also contributes convergence analysis for a bi-level, EM-like non-convex training scheme, complementing prior theoretical work on inverse problems over measure spaces.

Real-world applications (as named in the paper):

• Medical imaging, such as tomographic projections of internal structures.
• Astronomical observation through atmospheric distortion.
• Image restoration under nonlinear or non-differentiable degradations, including motion blur and JPEG compression.
• Scientific measurement subject to non-Gaussian noise, such as Poisson noise.

Industry relevance: Any pipeline where a physical simulator or degradation process exists but clean ground-truth data is scarce or uncollectible — imaging hardware vendors, remote sensing, scientific instrument manufacturers, and learned reconstruction products — can in principle use the black-box simulator already in hand instead of building differentiable forward models or collecting paired clean data.

Future Directions

• Extending the KL contraction analysis to the ε → 0 probability-flow ODE regime, which the paper explicitly leaves uncovered.
• Quantifying the approximation–estimation tradeoff (condition number versus approximation error δ) for specific function classes.
• Reducing the Lipschitz constant L_ε by going beyond linear interpolants, for example through better "preconditioning" to the channel K, which the authors flag as future work.
• Addressing time discretization, since the theoretical analysis is carried out in the idealized continuous-time limit.

Target Audience

Researchers and graduate students in machine learning theory and generative modeling, particularly those working on inverse problems, diffusion and flow matching, score-based methods, and stochastic interpolants. It is also relevant to applied scientists in imaging, astronomy, and computational physics who have simulators of their measurement process but no clean data. The mathematical content — measure-space operators, Wasserstein and KL analysis, Girsanov-based bounds — makes the theory sections most accessible to readers with a background in probability and optimal transport, though the algorithmic description is readable at an intermediate level.

Authors’ abstract

Transport-based methods have emerged as a leading paradigm for building generative models from large, clean datasets. However, in many scientific and engineering domains, clean data are often unavailable: instead, we only observe measurements corrupted through a noisy, ill-conditioned channel. A generative model for the original data thus requires solving an inverse problem at the level of distributions. In this work, we introduce a novel approach to this task based on Stochastic Interpolants: we iteratively update a transport map between corrupted and clean data samples using only access to the corrupted dataset as well as black box access to the corruption channel. Under appropriate conditions, this iterative procedure converges towards a self-consistent transport map that effectively inverts the corruption channel, thus enabling a generative model for the clean data. We refer to the resulting method as the self-consistent stochastic interpolant (SCSI). It (i) is computationally efficient compared to variational alternatives, (ii) highly flexible, handling arbitrary nonlinear forward models with only black-box access, and (iii) enjoys theoretical guarantees. We demonstrate superior performance on inverse problems in natural image processing and scientific reconstruction, and establish convergence guarantees of the scheme under appropriate assumptions. Our source code is publicly available at https://github.com/modichirag/SCSI