Trajectory Stitching for Solving Inverse Problems with Flow-Based Models

Alexander Denker, Zeljko Kereta, Carola-Bibiane Schönlieb, Moshe Eliasof
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:24478-24499, 2026.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-denker26a, title = {Trajectory Stitching for Solving Inverse Problems with Flow-Based Models}, author = {Denker, Alexander and Kereta, Zeljko and Sch\"{o}nlieb, Carola-Bibiane and Eliasof, Moshe}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {24478--24499}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/denker26a/denker26a.pdf}, url = {https://proceedings.mlr.press/v306/denker26a.html}, 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.} }
Endnote
%0 Conference Paper %T Trajectory Stitching for Solving Inverse Problems with Flow-Based Models %A Alexander Denker %A Zeljko Kereta %A Carola-Bibiane Schönlieb %A Moshe Eliasof %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-denker26a %I PMLR %P 24478--24499 %U https://proceedings.mlr.press/v306/denker26a.html %V 306 %X 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.
APA
Denker, A., Kereta, Z., Schönlieb, C. & Eliasof, M.. (2026). Trajectory Stitching for Solving Inverse Problems with Flow-Based Models. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:24478-24499 Available from https://proceedings.mlr.press/v306/denker26a.html.

Related Material