PaperlayerView exact revision on arXiv

Source-linked AI summary

Twist Flow for Inverse Problems

Shiqin Zeng, Zijun Deng, Felix J. Herrmann

arXiv:2610.09281v1cs.LGphysics.geo-ph

TL;DR

Inverse-problem samplers must capture plausible targets consistent with observations, yet paired conditional training can underuse latent noise and underrepresent posterior variation. Joint twist-flow learns an augmented transport that couples target samples with a likelihood-side coordinate. It better preserves posterior support in toy problems and is evaluated for observation-aware variability in image restoration and seismic inversion.

  • Problem

    Inverse problems can map one observation to multiple plausible targets, while paired training may lead conditional generators to underuse latent noise and represent too little posterior variation.

    Source: §1 · p. 1

  • Method

    Joint twist-flow learns a continuous transport from (zx, y) to (x, zy), using zx for posterior sampling and zy as an auxiliary likelihood-side coordinate.

    Source: §1 · p. 2, §3.2 · p. 5, §5 · p. 11

  • Results

    In toy problems, joint twist-flow reduces latent underuse, mode dropping, and support distortion relative to a direct conditional-flow baseline; image restoration shows increased variability in weakly constrained components while maintaining observation consistency, and the seismic experiment applies the formulation.

    Source: §4.1 · p. 6, §4.2 · p. 8, §5 · p. 11

  • Takeaways & Limitations

    The reported experiments support using the augmented transport to preserve posterior variation while keeping generated samples tied to the fixed observation.

    Source: §5 · p. 11

    The setup uses paired target–observation data, and the Gaussian-scale likelihood coordinate is motivated under a Gaussian observation model.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

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.

1 Introduction

Inverse-problem samplers must represent unresolved target uncertainty as well as observation consistency, but paired training can encourage conditional models to rely too little on latent noise. Joint twist-flow addresses this with an augmented transport that couples posterior samples to an observation-side coordinate.

1 Introduction

  • A useful posterior sampler must capture multiple plausible targets compatible with one observation, because observation consistency alone does not represent unresolved uncertainty.When distinct solutions are compatible with the same observation, a sampler may miss solutions, distort them, or create unrealistic transitions.Source: §1 · pp. 1–2
  • Paired training can satisfy the objective by relying mainly on the observation and weakly using latent noise, producing too little posterior variation.With multiple compatible solutions, this can lead to missed solutions, distortion, or unrealistic transitions.Source: §1 · pp. 1–2
  • Joint twist-flow learns an augmented transport from (zx, y) to (x, zy), rather than a direct conditional map from latent noise and observation to x.zx drives posterior sampling, while zy is an auxiliary terminal coordinate associated with the observation branch.Source: §1 · p. 2, Figure 1 · p. 2
  • For a fixed observation, samples generated from independent zx draws form an ensemble whose x-variability represents posterior uncertainty unresolved by the observation.The added observation-side degree of freedom is intended to reduce nearly deterministic dependence of x on y.Source: §1 · p. 2
  • The paper evaluates posterior variation and observation consistency, including reference-posterior comparisons on low-dimensional problems and image restoration and seismic velocity-model inversion.The reported low-dimensional experiments show reduced latent underuse and better preservation of posterior support.Source: §1 · p. 2

2 Related work

The paper situates joint twist-flow among learned-prior inverse methods, continuous-time flow approaches, Gaussian-coordinate sampling methods, and models of shared posterior and likelihood conditionals. Its distinction is an augmented flow-matching transport for fixed-observation inference.

Learned-prior methods for inverse problems

  • Plug-and-play, RED, and diffusion or score-based inverse methods combine learned priors or dynamics with observation-consistency or guidance mechanisms.Joint twist-flow instead learns an amortized posterior sampler as an augmented flow-based transport for fixed-observation inference.Source: §2 · pp. 2–3

Flow-based posterior sampler parameterizations

  • Normalizing and continuous normalizing flows, flow matching, rectified flow, and stochastic interpolants provide transport and velocity-field foundations for the method.Joint twist-flow applies the continuous-time transport viewpoint to (zx, y) → (x, zy), rather than a direct conditional map to x alone.Source: §2 · p. 3

Randomized least squares and transport for Bayesian inverse problems

  • Gaussian-coordinate and residual-based approaches motivate the two branches: zx is a reference coordinate for the unknown, and zy a likelihood-side residual coordinate.Unlike RTO-style methods, joint twist-flow learns the relation by flow matching rather than solving a randomized optimization problem for each sample.Source: §2 · p. 3

Joint conditional generation

  • The shared-joint view treats posterior and likelihood conditionals as conditionals of one joint distribution; joint twist-flow realizes it through an augmented continuous-time transport.Here zx drives posterior sampling, while zy is paired with the observation branch as a likelihood-side residual coordinate.Source: §2 · pp. 3–4

3 Method

The method frames inverse problems as ambiguous mappings from paired targets and observations, then learns a flow over an augmented state. Its x-marginal supplies posterior samples while a Gaussian likelihood-side coordinate is jointly generated.

3.1 Problem setup

  • Paired target–observation data define the forward relation, but a fixed observation may correspond to multiple plausible targets.Under a Gaussian observation model, the whitened residual motivates representing observation compatibility on a Gaussian scale.Source: §3.1 · p. 4

3.2 Joint twist-flow

  • A one-sided conditional flow uses the observation only as input and must encode both target variability and observation compatibility in its x output.In paired inverse-problem training, this parameterization can favor a narrow reconstruction and underrepresent conditional variability.Source: §3.2 · pp. 4–5
  • Joint twist-flow instead transports the augmented state (zx, y) to (x, zy), with a Gaussian sampling coordinate and fixed observation at the source and target plus likelihood-side coordinate at the terminal.During training, zy is sampled independently at a prescribed Gaussian scale and paired with x; the joint velocity field is trained along a probability path between the states.Source: §3.2 · p. 5
  • For fixed y, the learned transport generates x posterior samples and an associated likelihood-side coordinate, while posterior sampling uses the x-marginal.The auxiliary coordinate is amortized through training, so sampling requires zx and the fixed observation rather than a new randomized optimization.Source: §3.2 · p. 5
  • The prescribed scale for zy gives its likelihood-side coordinate an interpretable Gaussian scale while allowing x to vary and keeping the coordinate organized around that scale.The parameter σy controls the prescribed scale, and the confidence-level ball is defined using a chi-square quantile.Source: §3.2 · p. 5

4 Experiments

The experiments compare joint twist-flow with CondFlow on toy posteriors with MCMC references, image restoration without true posterior samples, and an RTM-conditioned seismic inversion. Across these settings, joint twist-flow better preserves posterior support or increases observation-compatible variability.

4.1 Toy inverse problems: posterior support preservation

  • Across two_moons, ring, and triangular cluster, joint twist-flow better matches MCMC marginals and preserves feasible posterior support than CondFlow.CondFlow misses curved branches, collapses to narrow ring arcs, or places mass near or between separated clusters; joint twist-flow better preserves these distinct supports.Source: §4.1 · pp. 6, 8

4.2 Image restoration: observation-aware posterior variability

  • 0.069→0.202 image-space pairwise RMSE accompanies Ir improving from 0.710 to 0.881 in blur-downsampling; across restoration tasks, variability rises while observation consistency remains high.Joint twist-flow increases spread in weakly constrained image components while keeping the observed component consistent with the fixed observation.Source: §4.2 · pp. 8–9, Table 2 · p. 9

4.3 Seismic inverse problem

  • Joint twist-flow broadens uncertainty around a comparable velocity reconstruction, concentrating spread in weakly constrained geology while maintaining shot-record coverage under the fixed RTM observation.Posterior-mean RMSE and SSIM remain close to CondFlow; the wider generated shot-record band better covers the ground-truth trace in this example.Source: §4.3 · pp. 9–10, Figure 4 · p. 10, Table 3 · p. 10

5 Conclusions

Joint twist-flow learns an augmented transport for posterior sampling, using separate coordinates for target variability and observation-side residual information. Experiments support improved preservation of posterior variation across toy inverse problems, image restoration, and seismic inversion.

5 Conclusions

  • Joint twist-flow transports the augmented state (z_x, y) to (x, z_y), with z_x driving posterior sampling and z_y providing a likelihood-side residual coordinate.Posterior uncertainty remains measured by variability across generated x-samples.Source: §5 · p. 11
  • In controlled toy inverse problems, joint twist-flow reduces latent underuse, mode dropping, and support distortion relative to a direct conditional-flow baseline.The low-dimensional experiments use reference posterior samples.Source: §5 · p. 11
  • Image-restoration experiments show increased variability in weakly constrained components while suppressing variability in components that determine the observation.The seismic experiment applies the formulation to a high-dimensional inverse problem with multiple plausible subsurface structures.Source: §5 · p. 11
  • The paper identifies improved calibration diagnostics, adaptive choices of the observation-side scale σ_y, and reverse-transport diagnostics as directions for future work.Source: §5 · p. 11

A.1 Topological Obstruction in Conditional Inverse Transport

The appendix formalizes a topological obstruction for direct conditional transport and motivates joint twist transport as a way to reduce parameterization pressure behind empirical failure modes, without claiming to remove the obstruction in projected x-space.

A.1 Topological Obstruction in Conditional Inverse Transport

  • The appendix formalizes a topological obstruction for direct conditional transport in inverse problems.Source: §A.1 · p. 14
  • Joint twist transport couples the target branch with a likelihood-side coordinate to reduce parameterization pressure associated with the obstruction.Source: §A.1 · p. 14
  • The proposed augmentation is not claimed to remove the topological difficulty in projected x-space; it is intended to make empirical failures such as mode dropping, spurious bridges, or off-support mass less likely.Source: §A.1 · p. 14

A.1.1 Setup

The setup considers inverse problems with multiple separated explanations for an observation and contrasts direct conditional maps with the geometry of multimodal posterior support. Under stated structural assumptions, connected latent regions cannot match disconnected posterior high-density regions without support distortion.

A.1.1 Setup

  • The inverse observation model uses a continuous forward operator h from target space X to observation space Y, with observational noise ε.Source: §A.1.1 · p. 14
  • The assumed regime allows one observation to admit multiple well-separated explanations, with positive prior density near both explanatory sets.The appendix describes this as posterior multimodality induced by the inverse problem.Source: §A.1.1 · p. 15
  • A response gap away from the explanatory neighborhoods yields a disconnected posterior high-density region under sufficiently small observational noise.Source: §A.1.1 · p. 15
  • A direct conditional generator is assumed continuous in its latent variable, while its latent high-density region is connected, as for a unimodal Gaussian base distribution.Source: §A.1.1 · p. 15
  • Under these assumptions, the posterior high-density set is disconnected while the image of a connected latent high-density region under a continuous direct conditional map remains connected, forcing support distortion.Source: §A.1.1 · p. 15

A.1.2 Multimodal posterior geometry

Under the appendix assumptions, sufficiently small observation noise can produce a posterior superlevel set with at least two distinct connected components near separate explanations. The proof relies on high likelihood near both explanations and a response gap elsewhere.

A.1.2 Multimodal posterior geometry

  • For sufficiently small σ, some posterior superlevel set has at least two connected components near the distinct explanatory sets A1 and A2.Source: §A.1.2 · p. 16
  • The likelihood is maximized on A1 and A2, where the prior is positive nearby, supporting high posterior density around both branches.Source: §A.1.2 · p. 16
  • Outside the disjoint explanatory neighborhoods, the forward response stays at least η from y0, so under Gaussian noise the corresponding posterior density bound tends to zero as σ decreases.Source: §A.1.2 · p. 16
  • Together, these conditions permit a threshold λ for which the posterior superlevel set contains disconnected high-density branches near A1 and A2.Source: §A.1.2 · p. 16

A.1.3 Topological obstruction for direct conditional transport

Under Assumptions A.4–A.7, direct conditional transport cannot reproduce a disconnected posterior superlevel set when its latent high-density set is connected. Representing multiple posterior branches then entails mode dropping or connecting them through low-density regions.

A.1.3 Topological obstruction for direct conditional transport

  • Under Assumptions A.4–A.7 and the specified λ, a connected direct-conditional image cannot match a posterior superlevel set with at least two connected components.Theorem A.9 states the obstruction for the image of the latent high-density set.Source: §A.1.3 · p. 16
  • The corollary identifies three unavoidable outcomes when direct conditional transport attempts to represent multiple true branches: mode dropping, a spurious bridge, or off-support spurious mass.These outcomes follow under Theorem A.9's assumptions.Source: §A.1.3 · pp. 16–17
  • If the model covers multiple true branches while remaining connected, its connecting path must pass through low-density regions, creating a spurious bridge or off-support mass.If it avoids connecting the branches, it must omit at least one branch.Source: §A.1.3 · p. 17

A.1.4 How joint twist transport changes the parameterization pressure

Joint twist transport maps between augmented source and target spaces, pairing generated targets with a terminal observation-side coordinate. This changes the parameterization pressure relative to an x-valued conditional decoder, without eliminating the difficulty of representing multimodal posteriors after projection to x.

A.1.4 How joint twist transport changes the parameterization pressure

  • The augmented state couples each generated target to a likelihood-side coordinate, reducing pressure to encode observation compatibility and posterior variation solely through x.This changes the parameterization but does not remove the topological difficulty after projection onto x.Source: §A.1.4 · pp. 17–19
  • Assumption A.11 requires a continuously differentiable joint map that is locally invertible on the region of interest.The text motivates this condition for endpoint maps of sufficiently regular flows or ODEs.Source: §A.1.4 · p. 18
  • Joint twist transport learns an augmented map from (zx, y) to (x, zy), rather than representing conditional variation through an x-valued decoder alone.For fixed y0, the restricted map sends zx to the pair (x, zy).Source: §A.1.4 · pp. 17–18
  • Assumption A.12 requires residual conditional variability to be locally represented in augmented coordinates, with some latent directions changing the observation branch while changing x little.The stated local sensitivity condition is motivated by the residual view of inverse problems.Source: §A.1.4 · p. 18
  • The auxiliary zy coordinate adds no information about x beyond y; under a Gaussian observation model, it is motivated by the normalized residual and serves as the terminal observation-branch coordinate.The branch can vary locally alongside the generated target.Source: §A.1.4 · pp. 17–18
  • The joint parameterization is compatible with flow matching, which can in principle learn the joint map when the data distribution admits a sufficiently regular transport representation.The cited compatibility is with flow matching [16].Source: §A.1.4 · p. 19

A.1.5 A least-squares residual interpretation of twist flow

The learned flow’s endpoint can be characterized as minimizing an augmented least-squares residual, but its velocity need not be a classical least-squares gradient. The same augmented transport supports fixed-observation posterior sampling and a reverse direction.

1. Flow matching learns a joint transport. Let

  • Joint twist-flow transports the augmented source (z_x, y) to (x, z_y), rather than decoding x alone from latent noise and the observation.For fixed y, independent draws of z_x yield an ensemble of target samples whose variability represents posterior uncertainty unresolved by the observation.Source: §1 · p. 20
  • The learned transport also defines a reverse direction, although the paper’s main use is fixed-observation posterior sampling.The reverse map is conditional on the learned endpoint map being sufficiently regular and locally invertible.Source: §1 · p. 20

A.1.5 A least-squares residual interpretation of twist flow

  • Each posterior sample is an endpoint of a learned transport whose inverse defines an augmented residual minimized at that endpoint.The residual includes a learned observation-consistency term and a term enforcing consistency with the latent coordinate indexing posterior variability.Source: §A.1.5 · p. 19, §1 · pp. 20–21
  • The twist velocity is interpretable as a learned preconditioned least-squares update direction, not necessarily as the ordinary Euclidean gradient of the energy.The flow-matching steps are not claimed to be identical to classical least-squares gradient steps.Source: §A.1.5 · p. 19, §1 · p. 21

A.2.1 Data generation and analytical priors

  • The toy priors yield distinct posterior-support structures, including curved branches, annular support, and separated cluster-like modes.The two-moons prior concentrates density near two curved branches, while the triangular mixture has three dominant modes.Source: §A.2.1 · pp. 21–22

A.2.2 Reference posterior sampling with MCMC

  • Low-dimensional toy experiments use MALA samples as evaluation references; this analytical-prior and reference-sampling construction is specific to those experiments.The retained MCMC samples provide marginal Wasserstein distances, KS statistics, mean and covariance errors, and posterior overlays.Source: §A.2.2 · p. 22

A.3 Seismic data simulation and RTM construction

  • The seismic model conditions on an RTM image derived from simulated shot records, rather than directly on raw seismic data.The RTM image is computed using a fixed smooth background model and serves as the generative model’s observation.Source: §A.3 · p. 23, §A.3.1 · p. 23

A.4 Additional reproducibility details

  • Learning-based experiments compare joint twist-flow with a direct conditional-flow baseline using matched training pairs and, where possible, the same preprocessing, optimizer, architecture width, and evaluation protocol.Image restoration uses three observation operators: low-frequency restoration, blur-downsampling, and blur-masking.Source: §A.4 · p. 23, §A.5 · p. 24
Loading 2610.09281v1…