Twist Flow for Inverse Problems

Shiqin Zeng; Zijun Deng; Felix J. Herrmann

cs.LGOct 7, 2026

Twist Flow for Inverse Problems

Authors: Shiqin Zeng, Zijun Deng, Felix J. Herrmann

Organizations: Georgia Institute of Technology

Abstract

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.

Figures & tables

Figure 1 : Overview of conditional flow matching and joint twist-flow. (a) Conditional flow matching samples $x$ from a one-way conditional map with fixed observation $y_{\mathrm{obs}}$ , which may underrepresent posterior ambiguity. (b) Joint twist-flow learns an augmented transport $(z_{x},y_{\mathrm{obs}})\mapsto(x,z_{y})$ , where $x$ is the posterior sample and $z_{y}$ is a Gaussian likelihood-side coordinate for the observation branch. This augmented transport is designed to preserve broader posterior variability under observation consistency.

Figure 2 : Toy posterior comparison for three inverse problems. Each row corresponds to one structured prior and one fixed observation $y_{0}$ . The first column shows samples from the prior in $x$ -space, with color indicating the scalar observation value $y$ generated by the forward model. The next three columns compare the MCMC posterior reference, conditional-flow samples, and joint twist-flow samples for the fixed $y_{0}$ ; gray points indicate the prior support for context. The last two columns compare the one-dimensional posterior marginals of $x_{1}$ and $x_{2}$ .

Setting	Method	$W_{1}(x_{1})$	$W_{1}(x_{2})$	KS $(x_{1})$	KS $(x_{2})$
Two moons $c_{1}=1.0,\ c_{2}=0.5$	Conditional Flow	$0.282\pm 0.246$	$0.663\pm 0.327$	$0.269\pm 0.127$	$0.311\pm 0.094$
	Joint Twist Flow (ours)	$\mathbf{0.132\pm 0.030}$	$\mathbf{0.151\pm 0.072}$	$\mathbf{0.121\pm 0.014}$	$\mathbf{0.074\pm 0.016}$
Ring $c_{1}=1.2,\ c_{2}=0.8$	Conditional Flow	$0.181\pm 0.019$	$0.188\pm 0.041$	$0.133\pm 0.022$	$0.161\pm 0.085$
	Joint Twist Flow (ours)	$\mathbf{0.098\pm 0.019}$	$\mathbf{0.133\pm 0.012}$	$\mathbf{0.056\pm 0.007}$	$\mathbf{0.101\pm 0.025}$
Triangular cluster $c_{1}=1.5,\ c_{2}=1.1$	Conditional Flow	$0.220\pm 0.104$	$0.271\pm 0.238$	$0.147\pm 0.051$	$0.231\pm 0.219$
	Joint Twist Flow (ours)	$\mathbf{0.142\pm 0.020}$	$\mathbf{0.129\pm 0.072}$	$\mathbf{0.091\pm 0.025}$	$\mathbf{0.059\pm 0.015}$

Table 1: Quantitative comparison against the MCMC posterior reference on the toy inverse problems. The coefficients

(c_{1},c_{2})

used in the nonlinear observation model are listed below each setting.

Panel of Figure 3 — Figure 3 : Image restoration under three observation operators.

Task	Observation $y=F(x)$	Weak component $u$	Model	$D_{x}\uparrow$	$S_{x}\uparrow$	$D_{r}\uparrow$	$I_{y}\uparrow$	$I_{r}\uparrow$
Low-frequency	$y=F_{\mathrm{LF}}(x)=\mathcal{F}^{-1}(P_{\Omega}\mathcal{F}x)$	$u=x-F_{\mathrm{LF}}(x)$	CondFlow	0.062	0.036	0.062	0.821	0.594
			Twist	0.156	0.093	0.126	0.977	0.848
Blur-downsampling	$y=F_{\mathrm{BD}}(x)=U(D(B(x)))$	$u=x-F_{\mathrm{BD}}(x)$	CondFlow	0.069	0.041	0.069	1.000	0.710
			Twist	0.202	0.121	0.193	0.996	0.881
Blur-masking	$y=F_{\mathrm{M}}(x)=MB(x)+(1-M)c$	$u=(1-M)x$	CondFlow	0.036	0.022	0.036	1.000	0.845
			Twist	0.212	0.127	0.165	0.998	0.889

Table 2 : Observation-aware posterior variability on CelebA restoration tasks. Each task uses observations of the form

y=F(x)

, while

u

denotes the task-specific weakly constrained component. Here

\mathcal{F}

denotes the Fourier transform,

P_{\Omega}

is a centered low-frequency Fourier mask,

B

is the blur operator,

D

and

U

denote downsampling and upsampling,

M

is the observed-region mask, and

c

is the constant fill value in the masked region. Metrics are averaged over three test cases.

D_{x}

is image-space pairwise RMSE across samples,

S_{x}

is mean sample standard deviation, and

D_{r}

is pairwise RMSE in

u

I_{y}

and

I_{r}

are empirical 90% interval inclusion rates in the observed component and in

u

, respectively.

Figure 4 : Seismic posterior sampling under a fixed RTM observation. The top rows show the ground-truth velocity and RTM image, posterior means, standard-deviation maps, and RMSE maps for CondFlow and joint twist-flow. Joint twist-flow preserves the main velocity structure while assigning larger uncertainty to reflector interfaces, structural transitions, and deeper weakly constrained regions. The bottom row forward models generated velocity samples under the same acquisition; shot-record bands and receiver-wise coverage profiles provide an observation-space consistency check.

	Per-sample		Posterior mean
Metric	CondFlow	Twist	CondFlow	Twist
RMSE	$0.101\pm 0.001$	$0.118\pm 0.005$	$0.0919$	$0.0921$
MAE	$0.0565\pm 0.0006$	$0.0680\pm 0.0027$	$0.0512$	$0.0564$
SSIM	$0.834\pm 0.003$	$0.777\pm 0.011$	$0.857$	$0.841$

Table 3 : Seismic reference metrics over 64 posterior samples

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Asset	Use in this paper	License or terms of use
CelebA	Image-restoration experiments.	Non-commercial research use only; images are not redistributed.
Devito	Acoustic wave-equation simulation.	MIT License.
JUDI.jl	Seismic modeling and RTM-related simulation workflow.	MIT License.
PyTorch	Neural-network training and inference.	BSD-style open-source license.
torchdiffeq	ODE integration for continuous-time sampling.	MIT License.
NumPy, SciPy, scikit-learn, Matplotlib	Numerical computation, metrics, and plotting.	Open-source scientific Python packages under permissive licenses.

Appendix

Table 4: Existing assets used in this work.

Explore similar work

Jun 23, 2026cs.CV

What Do Flow-Based Inverse Solvers Approximate? A Posterior-Transport View

A growing family of training-free solvers -- FlowDPS, FLOWER, PnP-Flow and their diffusion ancestors (DPS, DAPS) -- repurpose a pretrained flow-matching prior to solve imaging inverse problems by adding a measurement-guidance term to the deterministic probability-flow ODE. Despite strong empirical results, what these per-step corrections actually approximate -- and how far the resulting samples are from the true posterior

p(x\mid y)

-- has not been characterized. We give a posterior-transport account of flow-based inverse problem solving. Our starting point is a simple but consequential fact: for a \emph{deterministic} flow prior, Bayesian conditioning is realized entirely by a \emph{reweighting of the source distribution}, not by a drift correction; pushing the reweighted source through the \emph{unmodified} velocity field yields exact posterior samples. From this we show that trajectory-guidance solvers can be read as the minimum-kinetic-energy \emph{correction} field needed to morph the unconditional source into the posterior, and that FlowDPS / FLOWER / PnP-Flow correspond to distinct zeroth-order / Gaussian / proximal approximations of this single object; we bound the resulting posterior bias in Wasserstein distance. A controlled

2

D study with a closed-form posterior confirms the theory decisively: source reweighting matches the true posterior to the Monte-Carlo floor on every metric, whereas trajectory guidance incurs

200

800\times

larger error and collapses posterior modes, \emph{regardless of guidance strength}. Guided by the analysis we propose a cheap, principled velocity-correction solver that is competitive across two in-domain priors (AFHQ, CelebA) and two out-of-distribution settings while, unlike point-estimate source-space optimizers, producing diverse posterior samples with uncertainty that correlates with reconstruction error.

Jian Xu, Delu Zeng, John Paisley +1

1RIKEN iTHEMS · 2RIKEN AIP · 3South China University of Technology +1

Bayesian Inverse Problems Posterior Sampling Velocity Field Lineages

Ask PDF

Sep 28, 2026cs.CV

SNaP: One-Step Posterior Sampling for Noisy Inverse Problems

Diffusion and flow-matching models can produce high-quality posterior samples for inverse problems, but typically require tens to thousands of network evaluations per draw. MeanFlow enables one-step generation, yet applying it to inverse problems leaves no intermediate steps at which to enforce measurement consistency. We introduce SNaP, a one-step MeanFlow posterior sampler for linear inverse problems with Gaussian noise. Its central innovation is a measurement-adapted source: a Gaussian distribution whose mean and anisotropic covariance are determined by the measurement operator, observation, and noise level. The source anchors well-measured directions while preserving variation where the measurements are weak or uninformative. We show that the exact conditional flow transports this source to the true posterior. Across natural-image restoration and multi-coil MRI, SNaP produces diverse, high-quality samples with one network evaluation per draw, 30 to 2250

\times

faster than iterative samplers.

Shirin Shoushtari, Edward P. Chandler, Xiao Shi +1

WashU · UW-Madison

Posterior Sampling Bayesian Inverse Problems Inverse Problem Meanflow Covariance Magnetic Resonance Imaging

Ask PDF

Feb 9, 2026eess.IV

Trajectory Stitching for Solving Inverse Problems with Flow-Based Models

Flow-based generative models have emerged as powerful priors for solving inverse problems. One option is to directly optimize the initial latent code (noise), such that the flow output solves the inverse problem. However, this requires backpropagating through the entire generative trajectory, incurring high memory costs and numerical instability. We propose MS-Flow, which represents the trajectory as a sequence of intermediate latent states rather than a single initial code. By enforcing the flow dynamics locally and coupling segments through trajectory-matching penalties, MS-Flow alternates between updating intermediate latent states and enforcing consistency with observed data. This reduces memory consumption while improving reconstruction quality. We demonstrate the effectiveness of MS-Flow over existing methods on image recovery and inverse problems, including inpainting, super-resolution, and computed tomography.

Alexander Denker, Zeljko Kereta, Carola-Bibiane Schönlieb +1

Department of Computer Science, University College London · Department of Applied Mathematics and Theoretical Physics, University of Cambridge.

Flow Models Inverse Problem Intermediate Latent States

Ask PDF