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.

Inverse problems aim to reconstruct signals or images from noisy, incomplete, or indirect measurements. Recent advances in flow-based generative models have made them promising tools for these tasks because they can capture realistic image structure. One common approach is to optimize the model’s latent variables so that the generated output matches the observed measurements. However, this requires backpropagating through the entire generative process, leading to high memory usage and numerical instability. In this work, we propose MS-Flow, a more memory-efficient approach that represents the generative process using a sequence of intermediate latent states rather than a single starting point. By enforcing consistency between neighboring states and the observed measurements, MS-Flow avoids the need to backpropagate the full trajectory during optimization. This reduces memory consumption while also improving reconstruction quality across several inverse problems.