Solving Inverse Problems with Flow-based Models via Model Predictive Control

TL;DR

This paper introduces MPC-Flow, a model predictive control framework that guides flow-based generative models for inverse problems, improving efficiency and scalability without requiring backpropagation through the model.

Contribution

MPC-Flow formulates inverse problem solving as a sequence of control sub-problems, enabling practical, training-free guidance of flow models at inference time.

Findings

Strong performance on image restoration benchmarks

Scalable to large models like FLUX.2 (32B)

Avoids backpropagation through the generative trajectory

Abstract

Flow-based generative models provide strong unconditional priors for inverse problems, but guiding their dynamics for conditional generation remains challenging. Recent work casts training-free conditional generation in flow models as an optimal control problem; however, solving the resulting trajectory optimisation is computationally and memory intensive, requiring differentiation through the flow dynamics or adjoint solves. We propose MPC-Flow, a model predictive control framework that formulates inverse problem solving with flow-based generative models as a sequence of control sub-problems, enabling practical optimal control-based guidance at inference time. We provide theoretical guarantees linking MPC-Flow to the underlying optimal control objective and show how different algorithmic choices yield a spectrum of guidance algorithms, including regimes that avoid backpropagation through the generative model trajectory. We evaluate MPC-Flow on benchmark image restoration tasks, spanning linear and non-linear settings such as in-painting, deblurring, and super-resolution, and demonstrate strong performance and scalability to massive state-of-the-art architectures via training-free guidance of FLUX.2 (32B) in a quantised setting on consumer hardware.