Twist Flow for Inverse Problems

The gist The authors propose joint twist-flow, an augmented flow-matching formulation that learns a continuous transport from an augmented source state (zx, y) to the augmented terminal state (x, zy) for posterior sampling in inverse problems.

Joint Twist Flow Formulation

Joint twist-flow is introduced as an augmented flow-matching formulation that learns a continuous transport from the augmented source state (zx, y) to the augmented terminal state (x, zy) where x is the target variable and y is the observation The setup involves learning an augmented transport Sθ(zx, y) = (x, zy), where zx ∼ N (0, I) drives posterior sampling and zy is an auxiliary terminal coordinate associated with the observation This added coordinate gives the transport an observation-side degree of freedom, reducing the tendency to fit paired training data through a nearly deterministic dependence of x on y For a fixed observation y, the learned flow maps independent draws of zx ∼ N (0, I) to an ensemble of target samples x

Key Components and Interpretation

The formulation introduces two key coordinates: zx and zy The source state contains a Gaussian sampling coordinate and the fixed observation, while the terminal state contains the generated target and a Gaussian likelihood-side coordinate zy is sampled independently from N (0, σy squared y I) during training and paired with x at the terminal state This gives the likelihood-side branch a controlled Gaussian scale, while the learned transport couples the observation branch with posterior sampling For confidence level q ∈ (0, 1), a ball Bq(σy) gives an interpretable Gaussian scale for the likelihood-side coordinate generated by the learned transport The posterior sampler is then defined as the x-marginal of this augmented distribution, pθ(x y) = Z pθ(x, zy y) dzy

Validation on Toy Inverse Problems

The method was validated on low-dimensional inverse problems with reference posterior samples, where joint twist-flow better preserves multimodal posterior support than a direct conditional-flow baseline In controlled two-dimensional inverse problems, joint twist-flow reduces latent underuse and improves support preservation relative to a direct conditional-flow baseline The table shows that joint twist-flow reduces both Wasserstein and Kolmogorov–Smirnov discrepancies for x1 and x2 in all three settings This indicates that the twist-flow marginals more closely follow the MCMC reference than the conditional-flow marginals

Performance in Image Restoration

In image restoration tasks, joint twist-flow produces larger sample standard deviation and broader PCA clouds, indicating a wider generated ensemble Its F(x) standard-deviation maps remain small relative to the image-space spread, and the Iy values in Table 2 remain high, showing that generated samples remain consistent with the fixed observation This indicates that twist-flow explores multiple observation-compatible solutions rather than collapsing to the paired reference

Application to Seismic Inversion

The method was further evaluated on a seismic inverse problem where the target variable x is a two-dimensional subsurface velocity model Joint twist-flow produces larger and more structured uncertainty, concentrated around reflector interfaces, structural transitions, and deeper regions where the RTM image provides weaker constraints The generated velocity samples are forward modeled under the same acquisition geometry, and the resulting shot-record bands are compared with the ground-truth shot record This indicates that joint twist-flow produces a broader set of velocity models that remain compatible with the observed data

Conclusion

The experiments support three related conclusions: first, in controlled toy inverse problems, joint twist-flow reduces latent underuse, mode dropping, and support distortion relative to a direct conditional-flow baseline second, in image restoration, the augmented transport produces observation-aware posterior variability: uncertainty increases in weakly constrained components while remaining suppressed in components that determine the observation third, the seismic experiment shows that the same formulation can be applied to a high-dimensional scientific inverse problem where a single observation may be compatible with multiple plausible subsurface structures The main risk is that generated samples may be over-interpreted as calibrated posterior truth

--- Page 5 ---

The gist The authors propose joint twist-flow, an augmented flow-matching formulation that learns a continuous transport from an augmented source state (zx, y) to the augmented terminal state (x, zy) for posterior sampling in inverse problems.

Application to Seismic Inversion

The method was further evaluated on a seismic inverse problem where the target variable x is a two-dimensional subsurface velocity model <ref:2610.

1. The improved system can perform posterior sampling in inverse problems by learning a continuous transport from the augmented source state (zx, y) to the augmented terminal state (x, zy) rather than relying on direct conditional maps that may under-represent posterior variability. This allows for posterior sampling requires generating samples that are consistent with given observations while capturing the range of plausible solutions even when the observation operator is many-to-one.

2. The system will generate posterior samples that preserve broader posterior variability under observation consistency, specifically addressing mode distortion and artificial transitions between distinct feasible solutions, as demonstrated by its ability to better match MCMC reference samples in toy problems.

3. In image restoration tasks, the improved system can achieve observation-aware posterior variability: uncertainty increases in weakly constrained components while remaining suppressed in components that determine the observation, leading to reconstructions that explore multiple observation-compatible solutions rather than collapsing to the paired reference.

4. For scientific inverse problems, the system can produce ensembles where larger and more structured uncertainty is assigned to areas like reflector interfaces, structural transitions, and deeper regions where the RTM image provides weaker constraints, offering a more physically meaningful uncertainty representation.

In Bayesian inverse problems, posterior sampling requires generating samples that are consistent with given observations while capturing the range of plausible solutions. Direct conditional generative models introduce latent noise to model this ambiguity, but paired inverse-problem training can still encourage an almost deterministic map from the observation to the target. As a result, generated samples may be observation-consistent while under-representing posterior variability, especially when the posterior is multimodal, leading to undercoverage, mode distortion, or artificial transitions between distinct feasible solutions. We propose joint twist-flow, an augmented flow-matching formulation that learns a continuous transport from the augmented source state (z x, y) to the augmented terminal state (x, z y). Here x is the target variable, y is the observation, z x is the Gaussian reference coordinate for posterior sampling, and z y is a Gaussian likelihood-side coordinate associated with the observation branch. Under a Gaussian observation model, z y is motivated by the normalized observation residual associated with observation compatibility. Its role is not to replace uncertainty in x, but to couple generated samples of x to observation consistency, helping reduce likelihood-inconsistent variation while preserving variability in weakly constrained directions. We validate the method on low-dimensional inverse problems with reference posterior samples, where joint twist-flow better preserves multimodal posterior support than a direct conditional-flow baseline. We further evaluate the method on image restoration and seismic subsurface velocity-model inversion, showing increased posterior variability while maintaining observation consistency.