Trajectory Stitching for Solving Inverse Problems with Flow-Based Models

2026-02-09

AI Summary

This work addresses the high memory cost and numerical instability associated with conventional flow-based methods for solving inverse problems, which typically require backpropagation through the entire generative trajectory. To overcome these limitations, the authors propose MS-Flow, a novel approach that models the generative trajectory as a sequence of intermediate latent states. By imposing flow dynamics constraints locally and coupling adjacent segments via a trajectory-matching penalty term, MS-Flow enables an alternating optimization strategy that updates latent states and enforces data consistency without full-trajectory backpropagation. This design substantially reduces memory consumption while maintaining numerical stability. Experiments on image inpainting, super-resolution, and CT reconstruction demonstrate that MS-Flow achieves superior reconstruction quality compared to existing methods, effectively balancing computational efficiency and accuracy.

Technology Category

• Machine Learning: Learning with Manifolds
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

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.

Problem

• inverse problems
• flow-based models
• trajectory optimization
• memory efficiency
• numerical stability

Innovation

• trajectory stitching
• flow-based models
• inverse problems
• intermediate latent states
• memory-efficient optimization