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

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 analysis 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.

We are interested in solving inverse problems in imaging, i.e., reconstructing images from noisy or incomplete observations. To recover high-quality images, one typically introduces regularisation or prior knowledge into the reconstruction process. Flow-based generative models provide a strong prior for these tasks. They work by starting from a simple base distribution and transporting it to the data distribution. To use flow-based models for inverse problems, the transport process must be adapted so that the generated image is both high quality and consistent with the observations. One way to achieve this is through optimal control, where the generative trajectory is steered to satisfy specific constraints. However, existing optimal control approaches are often computationally and memory intensive because they require solving a high-dimensional optimisation problem. Instead, we propose MPC-Flow, a framework based on model predictive control that makes this problem tractable. The key idea is to decompose the full optimisation problem into a sequence of smaller and easier sub-problems that are solved iteratively. We demonstrate that MPC-Flow scales to recent large-scale generative models, including FLUX.2 with 32 billion parameters.