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

TL;DR

MPC-Flow reframes conditioning flow-based generative models for inverse problems as a sequence of short-horizon optimal control sub-problems, enabling training-free guidance during inference with reduced memory demands. The approach offers theoretical guarantees linking the model-predictive scheme to the underlying optimal control objective and presents two design regimes: receding-horizon control and Delta-t horizon control, including a memory-efficient single-step variant. Empirically, MPC-Flow delivers strong performance on linear and nonlinear image restoration tasks and scales to large models like FLUX.2 (32B) on consumer hardware, demonstrating practical viability for high-capacity flow-based priors in real-world imaging applications.

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.

Figures (19)

• Figure 1: MPC-Flow strategy for guiding flow-based generative models toward a target objective. Starting from the current state, MPC-Flow plans a sequence of velocity adjustments that steer the flow toward the objective while keeping the intervention small (A). Only the initial part of the plan is applied (B), after which we re-plan from the new state and repeat the process (C).
• Figure 2: Receding-Horizon Control (RHC)
• Figure 3: Comparison of MPC-RHC with varying $K$ compared to the global optimal control solution. All approaches use the same initial value ${\bm{x}}_0$ and $\lambda=2500$.
• Figure 4: Qualitative results on the CelebA dataset for the image super-resolution task with noise level $\sigma = 0.05$ and $\times 2$ upscaling.
• Figure 5: Example images generated by FLUX.2 with training-free style transfer guidance by FlowChef (middle row) and MPC (bottom row). The top row shows the image generated without conditioning as well as the prompt used and the reference style image.

Theorems & Definitions (7)

• Remark 2.1
• Theorem 3.1: Optimality of Receding-Horizon Control
• Theorem 3.2: $\Delta t$-Horizon optimality
• Proposition 2.1: Existence of optimal control
• proof
• proof
• proof