Flow-Mask Inverse Dynamics Model (FM-IDM)

• Flow-Mask Inverse Dynamics Model (FM-IDM) is a training-free, mask-guided image restoration framework that uses a pretrained flow matching prior with mask-guided trajectory correction to recover degraded images.
• It integrates a mask-guided fusion mechanism with a correction step to enforce data fidelity and enhance restoration performance in tasks like inpainting, denoising, and super-resolution.
• FM-IDM achieves state-of-the-art perceptual quality while operating significantly faster than diffusion-based and plug-and-play models, enabling efficient high-resolution image restoration.

The Flow-Mask Inverse Dynamics Model (FM-IDM) is a training-free, mask-guided image restoration framework that leverages a pretrained unconditional flow-matching prior and enforces data fidelity via mask-guided trajectory correction. The method—essentially the Restora-Flow approach—achieves state-of-the-art perceptual quality in restoration tasks (inpainting, super-resolution, denoising) and operates an order of magnitude faster than prevailing diffusion and flow-based models. FM-IDM introduces mask-guided fusion and a correction mechanism into the flow matching paradigm, making use of a degradation mask at sampling time to integrate observed data and maintain consistency with degraded inputs (Hadzic et al., 25 Nov 2025).

1. Flow-Matching Framework

FM-IDM operates atop a pretrained unconditional flow-matching generative model. Let $p(x)$ denote a data distribution over $\mathbb{R}^d$. The flow-matching method learns a time-dependent velocity field $v_{\theta,t} : \mathbb{R}^d \to \mathbb{R}^d$ that parameterizes the deterministic transport of a Gaussian noise distribution at $t=0$ to $p(x)$ at $t=1$ via the ordinary differential equation:

$$\frac{d x(t)}{dt} = v_{\theta,t}\bigl(x(t)\bigr), \quad x(0) \sim \mathcal{N}(0, I)$$

The simulation-free conditional loss

$$\mathcal{L}_{\rm FM}(\theta) = \mathbb{E}_{t \sim U[0,1],\,x_1 \sim p(x),\,x_0 \sim \mathcal{N}(0, I)} \left\| v_{\theta,t}\left(\Psi_t(x_0)\right) - (x_1 - x_0) \right\|^2,$$

with $\Psi_t(x_0) = (1-t)x_0 + t x_1$, aligns $v_{\theta, t}$ with the optimal transport path. Sampling proceeds via the explicit Euler scheme:

$$x_{t+\Delta t} = x_t + \Delta t\,v_{\theta, t}(x_t), \quad t=0, \Delta t, \dots, 1-\Delta t$$

This paradigm supports high-quality unconditional generation and forms the backbone of FM-IDM.

2. Mask-Guided Conditioning for Inverse Problems

For masked inverse problems, the observation model is

$$z = H x + \xi, \quad \xi \sim \mathcal{N}(0, \sigma^2 I)$$

with $H$ a linear operator represented for masking as $(Hx)_i = x_i$ if $m_i=1$, and $0$ otherwise, where $m \in \{0,1\}^d$ encodes the degradation mask. At each time step, FM-IDM fuses observed data using:

• Noising: Observations are noised to the corresponding latent scale via $z' = t\, z + (1-t)\, \varepsilon$, $\varepsilon \sim \mathcal{N}(0, I)$.
• Mask-guided fusion: The current state is locally clamped by the available (noised) observations: $x'_t = m \odot z' + (1-m)\odot x_t$.
• Conditional ODE step: Next state is updated under the fused context: $$x_{t+\Delta t} = x'_t + \Delta t\,v_{\theta,t}(x'_t)$$.

Observed regions thus adhere to $z'$, while unobserved regions evolve under the generative prior.

3. Trajectory Correction Mechanism

A trajectory correction step addresses misalignments induced at the interface between masked and unmasked regions. For each ODE iteration, a correction cycle is performed as follows:

• Forward extrapolation: Progresses the current state toward the clean image manifold:

$$\widetilde x_1 = x_{t+\Delta t} + (1 - (t+\Delta t)) v_{\theta, t+\Delta t}(x_{t+\Delta t})$$

• Re-noising: Introduces appropriate stochasticity for the next time step:

$$x_t = t\,\widetilde x_1 + (1-t)\,\eta, \quad \eta \sim \mathcal{N}(0, I)$$

This cycle is repeated once ($C=1$) per iteration. Enabling more corrections ($C > 1$) is possible but incurs additional computational cost. This mechanism enforces data fidelity and ensures consistency across mask boundaries.

4. Architecture, Hyperparameters, and Sampling Procedure

FM-IDM employs a pretrained unconditional flow-matching network (e.g., U-Net with time embeddings). Key aspects include:

• Network input: Mask-fused image $x'_t$ plus a positional encoding of $t$. The mask $m$ and the noised observation $z'$ are handled externally, not as network inputs.
• Hyperparameters: Typical step counts are $N=64$ (denoising, box inpainting $128^2$), $N=128$ (2× super-resolution, random inpainting $128^2$), $N=256$ (4× super-resolution $256^2$), with $C=1$ correction per step.
• Sampling algorithm:

# Pseudocode for FM-IDM sampling: # Input: pretrained flow vθ, observation z, mask m, steps N, corrections C=1 x = Normal(0, I) for t in {0, Δt,…,1−Δt}: for c in 0,…,C-1: ε = Normal(0, I) z_prime = t * z + (1 - t) * ε x_prime = m * z_prime + (1 - m) * x x = x_prime + Δt * vθ,t(x_prime) if c >= 1 and t < 1 - Δt: η = Normal(0, I) x_forward = x + [1 - (t + Δt)] * vθ,t+Δt(x) x = t * x_forward + (1 - t) * η return x

5. Quantitative Performance

Evaluation on CelebA ($128\times128$) benchmarks yields superior results relative to prior art. The table below summarizes key metrics (LPIPS↓, SSIM↑, PSNR↑) and per-image runtime for FM-IDM versus relevant baselines.

Task	Restora-Flow (FM-IDM)	Best Baseline (Method)
Denoising, $\sigma=0.2$	LPIPS=0.019, SSIM=0.922, PSNR=33.09 dB, 0.58 s	LPIPS=0.056, SSIM=0.910, PSNR=32.12 dB, 4.60 s (PnP-Flow)
Box Inpainting (40×40)	LPIPS=0.018, SSIM=0.964, PSNR=30.91 dB, 2.06 s	LPIPS=0.016, SSIM=0.967, PSNR=30.81 dB, ≈33 s (RePaint)
2× Super-Resolution	LPIPS=0.014, SSIM=0.952, PSNR=33.59 dB, 3.63 s	LPIPS=0.014, SSIM=0.946, PSNR=32.59 dB, ≈33 s (RePaint)
Random Inpainting (70% miss)	LPIPS=0.015, SSIM=0.947, PSNR=32.71 dB, 3.63 s	LPIPS=0.022, SSIM=0.954, PSNR=33.55 dB, 4.60 s (PnP-Flow)

Analogous speed and/or perceptual advantages are observed on AFHQ-Cat, COCO, and X-ray-Hand datasets. FM-IDM achieves sub-5 s runtimes for $256 \times 256$ images ($N \leq 256$) on A100 GPUs.

6. Training-Free Operation, Computational Complexity, and Limitations

FM-IDM is intrinsically training-free: it operates by reusing a fixed, pretrained unconditional flow-matching prior without fine-tuning on degraded images. The overall sampling complexity is $O(N)$ network calls per sample.

Limitations include:

• Mask out-of-distribution: Irregular or atypical masks can induce boundary artifacts if $N$ is insufficient.
• Under-correction in extreme scenarios: A single correction per step may be insufficient for highly corrupted inputs; increasing $C$ ameliorates this at the expense of runtime.
• Generalization failure: Substantial deviations from the training data (e.g., rare poses in human faces) or objects outside the prior's support can yield failures.

A plausible implication is that the approach's efficacy is contingent on the representational breadth of the prior and the mask's adherence to the conditions the prior was exposed to during original training.

7. Significance Within Image Restoration

FM-IDM provides an efficient, flexible alternative to iterative diffusion-based and plug-and-play models for image restoration under mask-based degradations. Its trajectory correction mechanism and plug-and-play compatibility with unconditional flow-matching priors distinguish it within the domain, offering a favorable trade-off between speed and perceptual quality across tasks such as denoising, inpainting, and super-resolution. FM-IDM exemplifies the integration of mask-guided fusion and generative priors in training-free restoration pipelines (Hadzic et al., 25 Nov 2025).