Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Trajectory Stitching for Solving Inverse Problems with Flow-Based Models

Alexander Denker, Moshe Eliasof, Zeljko Kereta, Carola-Bibiane Schönlieb

TL;DR

MS-Flow tackles memory and conditioning issues in flow-based inverse problem solvers by introducing a multiple-shooting formulation that represents the generative trajectory with intermediate latent states. It enforces local trajectory consistency with soft penalties, enabling Jacobian-free updates and constant memory w.r.t discretization that allow finer temporal resolution. The approach yields competitive or superior reconstruction across image restoration tasks and scales to latent-flow models like Stable Diffusion, demonstrating robustness and scalability beyond single-shot methods. By trading a larger, better-conditioned optimization for improved stability and memory efficiency, MS-Flow broadens the practical use of flow priors for high-resolution inverse problems and suggests future extensions to adaptive penalties and SDE-based generators.

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.

Trajectory Stitching for Solving Inverse Problems with Flow-Based Models

TL;DR

MS-Flow tackles memory and conditioning issues in flow-based inverse problem solvers by introducing a multiple-shooting formulation that represents the generative trajectory with intermediate latent states. It enforces local trajectory consistency with soft penalties, enabling Jacobian-free updates and constant memory w.r.t discretization that allow finer temporal resolution. The approach yields competitive or superior reconstruction across image restoration tasks and scales to latent-flow models like Stable Diffusion, demonstrating robustness and scalability beyond single-shot methods. By trading a larger, better-conditioned optimization for improved stability and memory efficiency, MS-Flow broadens the practical use of flow priors for high-resolution inverse problems and suggests future extensions to adaptive penalties and SDE-based generators.

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.
Paper Structure (32 sections, 1 theorem, 41 equations, 8 figures, 6 tables, 1 algorithm)

This paper contains 32 sections, 1 theorem, 41 equations, 8 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Preliminaries
  3. Flow-based Models
  4. ODE-constrained Optimization
  5. Multiple-Shooting Flow
  6. From Single- to Multiple-Shooting
  7. Optimizing the MS-Flow Objective
  8. Complexity of MS-Flow
  9. Properties of MS-Flow
  10. Stability of Multiple-Shooting
  11. Convergence of Trajectory Updates
  12. Experiments
  13. Empirical Convergence and Efficiency Analysis
  14. Image Recovery Tasks
  15. Latent Flow Models
  16. ...and 17 more sections

Key Result

Proposition 4.1

Under regularity assumption (Assumption ass:gs_convergence), the iterates $\{X^{(\ell)}\}_{\ell\ge 0}$ generated by eq:gs_block_update satisfy:

Figures (8)

  • 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.
  • Figure 2: Evaluation on the OrganCMNIST for sparse-angle CT. Left: The effect of regularization terms. Right: Comparison of MS-Flow with D-Flow across temporal resolutions. Shaded areas represent the standard deviation over the 10 images.
  • Figure 3: Convergence of the trajectory loss in \ref{['eq:traj_loss']} for coordinate descent (with and without the Jacobian-free approximation) and full gradient descent.
  • Figure 4: Comparison of the reconstructions for the three image recovery tasks on CelebA. First row: Gaussian deblurring. Second row: Inpainting. Third row: Super-resolution.
  • Figure 5: Reconstruction for Gaussian deblurring on FFHQ with $\sigma_\text{noise}=0.01$.
  • ...and 3 more figures

Theorems & Definitions (3)

  • Proposition 4.1: Monotone descent and stationarity with MS-Flow
  • Remark 4.2: Inexact and Jacobian-free updates
  • proof : Proof of Proposition \ref{['prop:gs_stationarity']}