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

What Do Flow-Based Inverse Solvers Approximate? A Posterior-Transport View
Jun 23, 2026Jian Xu, Delu Zeng, John Paisley +1Bayesian Inverse Problems Posterior Sampling
1RIKEN iTHEMS · 2RIKEN AIP · 3South China University of Technology +1
SNaP: One-Step Posterior Sampling for Noisy Inverse Problems
Sep 28, 2026Shirin Shoushtari, Edward P. Chandler, Xiao Shi +1Posterior Sampling Bayesian Inverse Problems
WashU · UW-Madison
Trajectory Stitching for Solving Inverse Problems with Flow-Based Models
Feb 9, 2026Alexander Denker, Zeljko Kereta, Carola-Bibiane Schönlieb +1Flow Models Inverse Problem
Department of Computer Science, University College London · Department of Applied Mathematics and Theoretical Physics, University of Cambridge.

Jun 23, 2026, cs.CVJ/K move · Enter open · S save

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

Jian Xu, Delu Zeng, John Paisley, Qibin Zhao

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

Bayesian Inverse Problems Posterior Sampling Velocity Field Lineages

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

SNaP: One-Step Posterior Sampling for Noisy Inverse Problems

Trajectory Stitching for Solving Inverse Problems with Flow-Based Models