[论文解读] Trajectory Stitching for Solving Inverse Problems with Flow-Based Models

Alexander H. Denker, Moshe Eliasof+2 more|arXiv (Cornell University)|Feb 9, 2026

Generative Adversarial Networks and Image Synthesis被引用 0

一句话总结

MS-Flow introduces a multiple-shooting approach that represents the flow trajectory as intermediate latent states, enabling memory-efficient and stable inverse problem solving with flow-based models.

ABSTRACT

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.

研究动机与目标

• Motivate solving inverse problems with flow-based generative priors in imaging.
• Address high memory and conditioning issues of single-shot latent optimization in flow models.
• Propose a multiple-shooting framework that decouples dynamics, data consistency, and regularization.
• Enable scalable, parallelizable optimization by local trajectory stitching penalties.

提出的方法

• Represent the flow trajectory as K intermediate latent states (shooting points) instead of a single initial code.
• Formulate an MS-Flow objective with data consistency, radial Gaussian latent regularization, and trajectory consistency penalties.
• Use an alternating minimization: (i) update shooting points via coordinate descent with Jacobian-free gradient approximations, (ii) update the terminal image estimate via a proximal/regularized data-consistency step.
• Discretize dynamics with explicit Euler and employ local stitching constraints to avoid backpropagating through the full time horizon.
• Provide convergence guarantees for the trajectory updates under regularity assumptions and show memory complexity independent of time steps.
• Demonstrate scalability and parallelism by enabling batch evaluations over shooting intervals.

Figure 1 : Peak GPU memory for D-Flow vs MS-Flow (Ours) using Euler discretization of the ODE on CelebA. D-Flow scales linearly with the number of timesteps, while MS-Flow is constant.

实验结果

研究问题

• RQ1 How does MS-Flow compare to single-shot flow methods in terms of memory usage and reconstruction quality?
• RQ2 Can a multi-shooting formulation improve conditioning and convergence for latent-space optimization in flow-based inverse problems?
• RQ3 What are the computational and memory benefits of Jacobian-free trajectory updates versus exact gradients?
• RQ4 How does MS-Flow perform on diverse imaging inverse problems (inpainting, super-resolution, CT) and latent-flow models?
• RQ5 Is MS-Flow scalable to large latent-flow architectures (e.g., Stable Diffusion) while preserving reconstruction quality?

主要发现

• MS-Flow achieves constant memory with respect to ODE discretization, unlike single-shot approaches whose memory scales with timesteps.
• The coordinate-descent trajectory updates monotonically decrease the trajectory objective and converge to stationary points under regularity assumptions.
• Jacobian-free trajectory updates offer substantial speedups with preserved convergence guarantees under a sufficient-decrease condition.
• MS-Flow matches or surpasses state-of-the-art flow-based inverse solvers on CelebA-based tasks (deblurring, inpainting, super-resolution) in PSNR/SSIM, and demonstrates robustness across hyperparameters.
• When applied to latent-flow models (e.g., Stable Diffusion), MS-Flow outperforms FlowDPS in PSNR/SSIM for Gaussian deblurring across tested noise levels.
• The approach enables parallel processing across shooting intervals, facilitating scalable computation on modern hardware.

(a) Changing the regularization terms.