# Trajectory Stitching for Solving Inverse Problems with Flow-Based Models

arXiv: https://arxiv.org/abs/2602.08538 · HTML: https://mlsanity.com/paper/2602.08538/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 57 references that are in the map.
- First posted on arXiv 9 Feb 2026.

## Builds on (most cited first, 30 of 36)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [MedMNIST v2 -- A large-scale lightweight benchmark for 2D and 3D biomedical image classification](https://mlsanity.com/paper/2110.14795/) (2021)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [PULSE: Self-Supervised Photo Upsampling via Latent Space Exploration of Generative Models](https://mlsanity.com/paper/2003.03808/) (2020)
- [Compressed Sensing using Generative Models](https://mlsanity.com/paper/1703.03208/) (2017)
- [Rectified Flow: A Marginal Preserving Approach to Optimal Transport](https://mlsanity.com/paper/2209.14577/) (2022)
- [Analyzing Inverse Problems with Invertible Neural Networks](https://mlsanity.com/paper/1808.04730/) (2018)
- [D-Flow: Differentiating through Flows for Controlled Generation](https://mlsanity.com/paper/2402.14017/) (2024)
- [PnP-Flow: Plug-and-Play Image Restoration with Flow Matching](https://mlsanity.com/paper/2410.02423/) (2024)
- [Invertible generative models for inverse problems: mitigating representation error and dataset bias](https://mlsanity.com/paper/1905.11672/) (2019)
- [Training-free Linear Image Inverses via Flows](https://mlsanity.com/paper/2310.04432/) (2023)
- [Intermediate Layer Optimization for Inverse Problems using Deep Generative Models](https://mlsanity.com/paper/2102.07364/) (2021)
- [FlowDPS: Flow-Driven Posterior Sampling for Inverse Problems](https://mlsanity.com/paper/2503.08136/) (2025)
- [DMPlug: A Plug-in Method for Solving Inverse Problems with Diffusion Models](https://mlsanity.com/paper/2405.16749/) (2024)
- [Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching](https://mlsanity.com/paper/2405.18816/) (2024)
- [Norm-guided latent space exploration for text-to-image generation](https://mlsanity.com/paper/2306.08687/) (2023)
- [From Facial Parts Responses to Face Detection: A Deep Learning Approach](https://mlsanity.com/paper/1509.06451/) (2015)

## Built on by (most cited first, 2 of 2)

- [A Stability Benchmark of Generative Regularizers for Inverse Problems](https://mlsanity.com/paper/2605.10076/) (2026)
- [Spanning Diffusion and Flow-based Posterior Sampling via Exact Gaussian Conditioning](https://mlsanity.com/paper/2610.13812/) (2026)

## Abstract (arXiv)

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.