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

1. Problem Statement

Inverse problems involve recovering a signal $x \in \mathbb{R}^d$ from noisy measurements $y = A(x) + \epsilon$, where $A$ is a potentially non-linear forward operator and $\epsilon$ is additive Gaussian noise. While flow-based generative models (specifically Continuous Normalising Flows trained via Flow Matching) provide powerful unconditional priors for such tasks, guiding their dynamics to satisfy conditional constraints (i.e., data fidelity) remains challenging.

Existing training-free guidance methods often rely on heuristic strategies that lack theoretical guarantees regarding consistency with measurements and can suffer from computational instability. Conversely, framing conditional generation as a deterministic optimal control problem offers a principled approach but introduces prohibitive computational and memory costs. Solving the resulting trajectory optimization typically requires differentiating through the entire flow dynamics or solving adjoint equations, which scales poorly with model size and trajectory length, making it impractical for large-scale architectures.

2. Methodology: MPC-Flow

The authors propose MPC-Flow, a framework that applies Model Predictive Control (MPC) to guide pre-trained flow-based generative models. Instead of optimizing the entire trajectory at once, MPC-Flow decomposes the control problem into a sequence of short-horizon sub-problems.

Core Mechanism

At each time step $t$, the framework solves a local optimal control problem over a planning horizon $[t, t+H]$. The control $u(\tau)$ is optimized to minimize a cost function comprising control energy and a terminal (or intermediate) loss, subject to the flow dynamics:
$$\frac{dx(\tau)}{d\tau} = v_\theta(x(\tau), \tau) + u(\tau)$$
Only the initial portion of the computed control (corresponding to a short interval $\Delta t$) is applied to the system. The state is updated, the horizon is shifted forward, and the process is repeated (receding-horizon control).

Two Design Regimes

The paper analyzes two specific instantiations of this framework:

1. Receding-Horizon Control (MPC-RHC):

   • Horizon: $H = 1 - t$ (the full remaining trajectory).
   • Mechanism: At every step, the controller re-plans the entire remaining path.
   • Theoretical Guarantee: If the sub-problems are solved to optimality, MPC-RHC recovers the global optimal control defined by the original trajectory optimization problem.
   • Implementation: To ensure tractability, the remaining horizon is discretized into $K$ steps.
     • Multi-step ($K > 1$): Requires backpropagation through $K-1$ nested evaluations of the flow model, which is memory-intensive.
     • Single-step ($K = 1$): The horizon is discretized into a single step. This effectively linearizes the flow path, allowing the control to be solved without backpropagating through the generative model $v_\theta$. This variant drastically reduces memory requirements and enables scaling to massive models.

2. $\Delta t$-Horizon Control (MPC-$\Delta t$):

   • Horizon: $H = \Delta t$ (only the immediate next step).
   • Mechanism: A greedy approach optimizing control for the next immediate time step.
   • Cost Function: Requires an intermediate cost function $\Phi_{MPC}$. Theoretically, the optimal choice is the value function (cost-to-go). Since the true value function is intractable, the authors propose a heuristic approximation using a single Euler step (similar to the Tweedie estimator in diffusion models), which avoids backpropagation through the full trajectory.

3. Key Contributions

• Framework Proposal: Introduction of MPC-Flow, the first conditional generation method explicitly rooted in optimal control that scales to massive state-of-the-art architectures via training-free guidance.
• Theoretical Analysis: Formal proofs linking MPC-Flow to the underlying optimal control objective. Specifically, Theorem 3.1 establishes that MPC-RHC recovers the global optimal control, and Theorem 3.2 defines the conditions under which $\Delta t$-horizon control yields the optimal policy.
• Memory-Efficient Variants: Demonstration that the single-step RHC variant ($K=1$) eliminates the need for backpropagation through the generative trajectory, making it feasible to run on consumer hardware with tight memory constraints.
• Unified Spectrum: Showing that different choices of horizon length and control parameterization yield a spectrum of guidance algorithms, ranging from globally optimal (but expensive) to greedy (but efficient) approaches.

4. Experimental Results

The authors evaluated MPC-Flow on benchmark image restoration tasks (denoising, deblurring, super-resolution, inpainting) and large-scale conditional generation tasks.

• Image Restoration (CelebA):
  • MPC-Flow variants consistently outperformed global optimal control baselines (FlowGrad, OC-Flow) in terms of both performance (PSNR/SSIM) and runtime.
  • MPC-$\Delta t$ achieved state-of-the-art or near-state-of-the-art performance across linear tasks, often ranking second only to PnP-Flow, but with significantly lower computational cost than global control methods.
  • MPC-RHC ($K=1$) offered the fastest inference times but showed limitations on tasks with large null spaces (e.g., certain inpainting or super-resolution scenarios) where the single-step approximation struggled to fill missing information.
• Large Vision Models (FLUX.2 and SD3.0):
  • The framework was successfully applied to the FLUX.2 (32B) model in a 4-bit quantized setting on consumer hardware (NVIDIA RTX 3090 with 24GB VRAM).
  • Style Transfer: MPC-Flow produced more realistic images than FlowChef, preserving the original generation trajectory while adopting stylistic characteristics, whereas FlowChef introduced significant artifacts.
  • Image Colorization: MPC-Flow achieved higher PSNR (23.8 dB) and SSIM (0.891) compared to FlowChef and unconditional baselines, demonstrating the ability to steer the flow toward observation-consistent images without deviating from the learned manifold.
  • Super-Resolution: The method demonstrated improved perception-distortion trade-offs compared to existing baselines.

5. Significance and Claims

The paper claims that MPC-Flow bridges the gap between principled optimal control and practical inference-time guidance for large generative models. By decomposing the global optimization into manageable sub-problems, it enables:

1. Scalability: The ability to guide massive models (e.g., 32B parameters) on consumer hardware, a feat previously hindered by the memory costs of adjoint methods or full-trajectory backpropagation.
2. Training-Free Operation: The method does not require retraining the base model, leveraging pre-trained unconditional priors directly.
3. Theoretical Rigor: Unlike heuristic guidance strategies, MPC-Flow is grounded in optimal control theory, providing formal guarantees on optimality under specific discretization and horizon conditions.

The authors acknowledge limitations, noting that the single-step RHC ($K=1$) may not always yield optimal results for linear inverse problems with large null spaces and that the method relies on differentiable objectives. However, they position MPC-Flow as a significant step toward practical, controllable, and scalable generative modeling for inverse problems.