Flow Priors for Linear Inverse Problems via Iterative Corrupted Trajectory Matching
Proposes ICTM, an efficient MAP approximation for high-dimensional linear inverse problems using flow models, outperforming existing flow matching methods.
Key Findings
Methodology
This paper introduces the Iterative Corrupted Trajectory Matching (ICTM) algorithm, which decomposes the global MAP objective into N local objectives. It leverages Tweedie’s formula to perform sequential gradient updates, avoiding costly ODE backpropagation. The method partitions the flow trajectory into segments, approximates the MAP via local objectives, and uses path interpolation to ensure convergence. Theoretical analysis shows convergence as NFEs tend to infinity, and experiments demonstrate superior performance in super-resolution, deblurring, and inpainting tasks, outperforming baseline flow matching approaches.
Key Results
- On CelebA-HQ, ICTM achieves PSNR of 30.65 and SSIM of 0.894 in super-resolution, surpassing DPS-ODE's 29.57 PSNR and 0.872 SSIM.
- In medical compressed sensing, PSNR improves by 4-6 dB over baselines, with consistent qualitative gains.
- Ablation studies confirm that local objectives and Tweedie’s formula significantly contribute to the robustness and efficiency of the method under limited NFEs.
Significance
This work addresses the computational bottleneck of applying flow models to high-dimensional inverse problems, providing a scalable, theoretically grounded solution. It bridges the gap between high-fidelity probabilistic priors and practical optimization, enabling real-world applications in medical imaging and remote sensing. The approach offers a new paradigm for efficient MAP estimation, combining deep generative models with advanced optimization techniques, thus advancing both theoretical understanding and practical capabilities.
Technical Contribution
The paper introduces a novel decomposition of the MAP objective into local objectives, justified by theoretical convergence proofs. It integrates Tweedie’s formula for efficient gradient estimation, reducing reliance on costly ODE backpropagation. The method extends flow models’ applicability to high-dimensional inverse problems, providing a scalable, convergent framework that balances computational efficiency with reconstruction quality.
Novelty
This is the first work to decompose the flow-based MAP objective into a sum of local objectives, enabling efficient optimization via iterative gradient steps. The combination of path interpolation, local MAP formulation, and Tweedie’s formula constitutes a significant innovation over prior methods that rely solely on full ODE backpropagation or sampling. It offers a new perspective on leveraging invertible generative models for inverse problems.
Limitations
- The approach assumes the trajectory exactly follows a predefined interpolation path, which may not hold perfectly in learned models, potentially affecting accuracy.
- The Lipschitz continuity requirement on the velocity field limits the complexity of models that can be used.
- Performance in highly nonlinear or non-Gaussian noise scenarios remains to be validated, and robustness under extreme conditions needs further study.
Future Work
Future research will explore adaptive path interpolation, robustness to model misspecification, and extension to nonlinear inverse problems. Integrating learned priors with dynamic or multimodal data, as well as reducing hyperparameter sensitivity, are promising directions to broaden the method’s applicability.
AI Executive Summary
High-dimensional linear inverse problems, such as image super-resolution, deblurring, and inpainting, pose significant computational challenges. Traditional optimization methods struggle with the high cost of backpropagating through ODE solvers used in flow-based models, limiting their practical application. Recent advances in flow matching models have demonstrated their ability to efficiently generate high-resolution images by defining smooth, invertible probability paths. However, leveraging these models as priors for MAP estimation remains computationally intensive due to the need for full trajectory likelihood calculations.
This paper introduces the Iterative Corrupted Trajectory Matching (ICTM) algorithm, a novel approach that decomposes the global MAP objective into N local objectives, each corresponding to a segment of the flow trajectory. By employing path interpolation and Tweedie’s formula, ICTM performs sequential gradient updates that approximate the MAP solution efficiently. The key insight is that as the number of function evaluations (NFEs) increases, the sum of local objectives converges to the true MAP, providing a solid theoretical foundation.
Empirical results on datasets like CelebA-HQ and the Human Connectome Project demonstrate that ICTM outperforms existing flow matching-based methods, achieving higher PSNR and SSIM scores across various tasks. For instance, in super-resolution, the method reaches a PSNR of 30.65, surpassing baseline methods by over 1 dB. In medical imaging, PSNR improvements of 4-6 dB highlight its practical value. Ablation studies confirm the importance of local objectives and Tweedie’s formula, showing robustness even with limited NFEs.
Overall, this work offers a scalable, theoretically justified framework for high-quality image reconstruction using flow priors. It paves the way for real-time, high-dimensional inverse problem solutions in medical imaging, remote sensing, and beyond. Future directions include extending to nonlinear models, adaptive path strategies, and multimodal data integration, promising broader impact and enhanced robustness.
Deep Analysis
Background
线性逆问题在图像重建、遥感、医学成像等领域广泛存在。传统方法如正则化最小二乘和贝叶斯推断在低维场景中表现良好，但在高维空间中计算成本激增。近年来，正则化流（Normalizing Flows）通过可逆变换实现高效概率建模，特别是流匹配模型在高分辨率图像生成中表现突出，成为研究热点。这些模型的优势在于路径直线、可逆性强，便于直接计算图像似然，但在逆问题中应用时，MAP估计的计算仍受ODE反向传播限制，成为瓶颈。此前的研究多依赖采样或近似，缺乏高效逼近MAP的系统方案。
Core Problem
核心难题在于高维流模型在MAP估计中的计算复杂性，尤其是反向传播ODE求解器，耗时长、内存大，难以快速获得高质量重建。现有方法多依赖逐步采样或简化，效果有限，且缺乏理论保证。如何在保证重建质量的同时大幅降低计算成本，成为亟待解决的问题。
Innovation
本文提出的ICTM算法创新点包括：1）将全局MAP目标拆解为N个局部目标，减少ODE反向传播频次；2）引入路径插值和Tweedie公式，利用梯度逐步优化局部目标；3）在有限NFEs下逼近最优解，兼顾效率与准确性。该方法突破了传统流模型在高维逆问题中的计算瓶颈，为高效MAP估计提供新思路。
Methodology
- �� 使用预训练流模型作为图像先验，定义全局MAP目标。• 将轨迹划分为N段，每段对应局部MAP目标，利用路径插值连接起始点和目标图像。• 利用Tweedie公式，将梯度估算转化为可计算表达式，避免ODE反向传播。• 设计逐步优化流程：在每个时间步，通过梯度下降调整轨迹点，使其符合局部目标。• 结合路径插值和局部目标，逐步逼近全局MAP，确保重建质量。• 理论证明在无限NFEs时，局部目标与全局目标一致，保证收敛。
Experiments
在CelebA-HQ和HCP数据集上验证算法性能。设置不同噪声水平和采样次数，比较PSNR、SSIM指标，分析有限NFEs下的效果。采用对比基线如OT-ODE、DPS-ODE，进行消融分析，验证局部目标分解和Tweedie公式的贡献。参数调优包括路径插值系数、步长和正则化权重。
Results
在超分任务中，PSNR达30.65，优于DPS-ODE的29.57，SSIM提升至0.894。医学成像中，PSNR平均提升5dB。消融实验确认局部目标和Tweedie公式的关键作用。算法在有限NFEs下仍表现优异，验证其实用性和鲁棒性。
Applications
适用于高分辨率图像重建、医学成像、遥感数据恢复等。结合预训练流模型和ICTM算法，可实现快速高质量重建，降低硬件成本，提升工业应用效率。
Limitations & Outlook
假设轨迹严格遵循插值路径，实际中可能偏差影响效果。对路径平滑和Lipschitz条件要求较高，限制模型复杂度。在极端噪声或非线性模型中，鲁棒性尚待验证。未来需扩展非线性路径和多模态逆问题研究。
Plain Language Accessible to non-experts
想象你在厨房做一道复杂的菜肴。传统做法是逐步尝试，费时又不一定成功。现在，有个聪明的助手能提前规划每一步，边做边调整，确保每个环节都尽善尽美。这个助手就像本文的算法，把大任务拆成许多小任务，逐步优化，最后做出完美的菜。它用一种聪明的技巧，避免每次都从头开始，节省时间，又保证效果。这就像用一种新厨艺，让你轻松做出高水平的菜肴。
ELI14 Explained like you're 14
想象你在玩拼图游戏，目标是拼出一幅完整的图片。传统方法是每次都从零拼，慢又难。现在，有个朋友告诉你：可以先把拼图分成几块，每块都拼得差不多，然后逐步调整，直到拼出完整的图。这就像本文的方法，把复杂的重建任务拆成很多小步骤，每次只优化一部分，然后逐渐逼近最终效果。这样既快又准，就像用聪明的拼图技巧，轻松拼出漂亮的画。
Glossary
Normalizing Flow (正则化流模型)
一种可逆的深度生成模型，通过变换简单分布到复杂数据分布，便于概率密度计算。
论文中用作图像先验，直接计算对数似然。
MAP (最大后验估计)
在已知观测数据条件下，找到最可能的隐藏变量值的统计推断方法。
结合流模型先验，用于优化图像重建。
Tweedie Formula (Tweedie公式)
一种估算后验梯度的数学工具，可在噪声模型下高效计算得分函数。
避免ODE反向传播，提升效率。
Local MAP (局部最大后验)
将全局MAP目标拆解为多个局部目标，逐步优化逼近整体最优。
算法核心思想，减少ODE反向传播频次。
Open Questions Unanswered questions from this research
- 1 如何在非线性或非高斯噪声模型中扩展，鲁棒性和泛化能力仍待验证。
- 2 未来应探索多模态和动态场景中的逆问题解决方案。
Applications
Immediate Applications
高分辨率医学图像重建
利用预训练流模型结合ICTM，快速恢复MRI或CT图像，减少辐射或扫描时间，提升临床效率。
遥感影像修复
修复卫星图像中的遮挡或模糊区域，提升地理信息系统精度，应用于环境监测。
Long-term Vision
智能图像处理平台
结合深度学习和优化算法，打造实时高质量图像重建系统，推动自动驾驶、虚拟现实等。
Abstract
Generative models based on flow matching have attracted significant attention for their simplicity and superior performance in high-resolution image synthesis. By leveraging the instantaneous change-of-variables formula, one can directly compute image likelihoods from a learned flow, making them enticing candidates as priors for downstream tasks such as inverse problems. In particular, a natural approach would be to incorporate such image probabilities in a maximum-a-posteriori (MAP) estimation problem. A major obstacle, however, lies in the slow computation of the log-likelihood, as it requires backpropagating through an ODE solver, which can be prohibitively slow for high-dimensional problems. In this work, we propose an iterative algorithm to approximate the MAP estimator efficiently to solve a variety of linear inverse problems. Our algorithm is mathematically justified by the observation that the MAP objective can be approximated by a sum of $N$ ``local MAP'' objectives, where $N$ is the number of function evaluations. By leveraging Tweedie's formula, we show that we can perform gradient steps to sequentially optimize these objectives. We validate our approach for various linear inverse problems, such as super-resolution, deblurring, inpainting, and compressed sensing, and demonstrate that we can outperform other methods based on flow matching. Code is available at https://github.com/YasminZhang/ICTM.