Solving Inverse Problems with Flow-based Models via Model Predictive Control
This addresses the computational bottleneck in using flow-based models for inverse problems like image restoration, though it is an incremental improvement over existing optimal control approaches.
The paper tackles the challenge of computationally intensive training-free conditional generation in flow-based models for inverse problems by proposing MPC-Flow, a model predictive control framework that formulates it as sequential control sub-problems, enabling practical guidance and demonstrating strong performance on image restoration tasks including scaling to FLUX.2 (32B) on consumer hardware.
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.