Posterior-Mean Rectified Flow (PMRF)

• Posterior-Mean Rectified Flow (PMRF) is a two-stage framework that first estimates the minimum-MSE posterior mean and then refines it via a learned rectified flow to match the target distribution.
• It addresses the classic distortion–perception trade-off by separating fidelity from realism, enabling precise restorations in applications like blind face restoration and virtual contrast MRI synthesis.
• Empirical evaluations demonstrate that PMRF improves perceptual quality while maintaining competitive MSE, outperforming baseline methods on benchmarks such as CelebA-Test and volumetric MRI datasets.

Posterior-Mean Rectified Flow (PMRF) is a two-stage framework for image restoration and image synthesis that separates distortion minimization from perceptual refinement. In its original formulation, PMRF was introduced for photo-realistic image restoration as a practical approximation to the estimator that minimizes mean squared error (MSE) under a perfect perceptual constraint, namely $p_{\hat{X}} = p_X$ (Ohayon et al., 2024). The same principle was later adapted to volumetric medical imaging for synthesizing virtual contrast-enhanced T1-weighted brain MRI from non-contrast inputs, where it was used to manage the perception–distortion trade-off in clinically relevant settings (Brandstötter et al., 18 Aug 2025). Subsequent work extended PMRF into the latent space of a variational autoencoder (VAE), arguing that latent-space transport can align more closely with human perception than pixel-space transport while preserving PMRF’s minimum-distortion philosophy (Luo et al., 1 Jul 2025).

1. Problem setting and conceptual definition

PMRF addresses the classical distortion–perception trade-off. In photo-realistic image restoration, a method is expected to achieve low distortion by metrics such as MSE, PSNR, and SSIM, while also producing outputs whose distribution matches that of natural images, as evaluated by perceptual metrics such as FID, KID, NIQE, and precision (Ohayon et al., 2024). The central observation motivating PMRF is that a single model trained only for voxel-wise or pixel-wise fidelity tends to produce blurred but faithful reconstructions, whereas a model trained primarily for realism can generate sharper but less accurate outputs. In the medical-imaging adaptation, this tension is described as particularly consequential because lesions must remain structurally correct while enhancement patterns, edges, and vascular details must still appear realistic (Brandstötter et al., 18 Aug 2025).

The defining move in PMRF is to decompose the task into two successive stages. The first stage estimates the conditional posterior mean, which is the minimum-MSE estimator. The second stage uses rectified flow to transport that estimate toward the target image distribution. This decomposition is explicitly presented as the mechanism by which PMRF improves perceptual realism relative to a plain posterior-mean regressor and improves fidelity relative to a direct generative flow baseline (Brandstötter et al., 18 Aug 2025).

The framework was introduced in the context of blind face restoration, denoising, super-resolution, inpainting, and colorization (Ohayon et al., 2024). It was later reformulated for virtual contrast MRI synthesis from paired pre- and post-contrast T1w volumes (Brandstötter et al., 18 Aug 2025), and then generalized to a latent-space variant for higher-fidelity perceptual face restoration (Luo et al., 1 Jul 2025).

2. Two-stage construction and mathematical formulation

In the generic restoration setting, PMRF considers a degraded observation $Y$, a clean target $X$, and a reconstruction $\hat{X}$. The first stage trains a regressor $f_\omega$ to approximate the posterior mean by minimizing

$$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$

This produces a high-PSNR, blurry-but-accurate estimate of the clean image (Ohayon et al., 2024).

The second stage defines a source variable

$$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$

and then trains a rectified-flow vector field $v_\theta$ along the linear interpolation path

$Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$

The corresponding flow-matching objective is

$$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$

At inference, PMRF computes the posterior-mean estimate, optionally adds small Gaussian noise, and integrates the ODE

$Y$0

typically using Euler steps (Ohayon et al., 2024).

The MRI adaptation uses the same structure with volumetric notation. Let $Y$1 denote the non-CE T1w volume and $Y$2 the paired true CE T1w volume. Stage 1 uses a 3D residual U-Net $Y$3 to predict

$Y$4

trained with

$Y$5

Stage 2 perturbs this estimate as

$Y$6

defines

$Y$7

and learns a time-conditioned 3D U-Net $Y$8 with

$Y$9

Inference integrates

$X$0

using explicit Euler steps

$X$1

and outputs $X$2 (Brandstötter et al., 18 Aug 2025).

3. Theoretical interpretation and relation to alternative estimators

The theoretical core of PMRF is the constrained problem

$X$3

which seeks the minimum-MSE estimator among all estimators whose marginal output distribution exactly matches the ground-truth image distribution (Ohayon et al., 2024). A key theoretical result, attributed in the paper to Dror et al., is that the optimal estimator under this perfect perceptual constraint can be constructed by transporting the posterior mean $X$4 to the target distribution via an optimal transport map. PMRF is presented as a practical approximation to that construction.

This point distinguishes PMRF from posterior sampling. Posterior sampling satisfies $X$5, but the paper states that its MSE is exactly twice the MMSE: $X$6 PMRF is therefore not framed as a posterior sampler; it is framed as a transport of the MMSE estimate that preserves a stronger fidelity anchor than direct sampling (Ohayon et al., 2024).

The original PMRF paper also states a proposition under the assumptions $X$7 and unique ODE solvability. Under those assumptions, PMRF achieves perfect perceptual index $X$8, has MSE no larger than posterior sampling, and has strictly smaller MSE under a non-degeneracy condition on $X$9 (Ohayon et al., 2024). The paper gives a toy Gaussian denoising example,

$\hat{X}$0

for which the posterior mean is

$\hat{X}$1

and the optimal estimator under the perfect perceptual constraint is

$\hat{X}$2

The paper states that PMRF with $\hat{X}$3 exactly recovers $\hat{X}$4 in this example (Ohayon et al., 2024).

Relative to GAN-weighted objectives, PMRF is positioned as avoiding heuristic balancing of distortion and realism losses. Relative to a flow that starts directly from degraded input $\hat{X}$5, PMRF is presented as more principled because its source is already a posterior-mean estimate in the same image domain as the target (Ohayon et al., 2024). A later extension argues, however, that pixel-space PMRF is limited in how efficiently it aligns with human perception, because human perception is better approximated by distances in representation space than in raw pixel space (Luo et al., 1 Jul 2025).

4. Empirical behavior in photo-realistic image restoration

The original PMRF study evaluates the method mainly on blind face image restoration and also on denoising, super-resolution, inpainting, and colorization (Ohayon et al., 2024). On blind face restoration, benchmarks include CelebA-Test, LFW-Test, WebPhoto-Test, CelebAdult-Test, and WIDER-Test. The reported metrics span perceptual quality, distortion, geometry, and real-world proxies, including FID, KID, NIQE, Precision, PSNR, SSIM, LPIPS, Deg, LMD, and IndRMSE (Ohayon et al., 2024).

On CelebA-Test, PMRF achieves the best FID, best KID, best PSNR, and best SSIM. The table reported in the paper gives FID $\hat{X}$6, KID $\hat{X}$7, PSNR $\hat{X}$8, and SSIM $\hat{X}$9 (Ohayon et al., 2024). On real-world benchmarks, PMRF achieves the best IndRMSE or is among the best while maintaining competitive perceptual quality; the reported best IndRMSE values are $f_\omega$0 on LFW-Test, $f_\omega$1 on WIDER-Test, $f_\omega$2 on WebPhoto-Test, and $f_\omega$3 on CelebAdult-Test (Ohayon et al., 2024).

Controlled ablations compare PMRF with flow conditioned on $f_\omega$4, flow conditioned on $f_\omega$5, and posterior sampling. The reported finding is that PMRF is generally best overall; in denoising, PMRF and flow from $f_\omega$6 are often close, while posterior sampling and flow conditioned on $f_\omega$7 are worse in MSE, consistent with the theory (Ohayon et al., 2024).

The implementation reported for blind face restoration uses SwinIR for the posterior-mean predictor, HDiT for the vector field, AdamW with learning rate $f_\omega$8, batch size $f_\omega$9, EMA decay $$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$0, mixed precision bfloat16, and $$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$1. Inference uses Euler steps, with default $$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$2 in the main blind-face experiments. The paper states that larger $$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$3 improves perceptual quality but can slightly hurt distortion (Ohayon et al., 2024).

5. Adaptation to volumetric contrast-enhanced MRI synthesis

In the MRI study, PMRF is used as a two-stage synthesis framework for generating virtual contrast-enhanced T1-weighted brain MRI from non-contrast T1-weighted inputs while explicitly managing the perception–distortion trade-off (Brandstötter et al., 18 Aug 2025). The model is trained on paired non-CE and CE T1w volumes from the BraTS 2023–2025 datasets, totaling 7,196 paired volumes. The data include high-grade gliomas, low-grade gliomas, metastases, meningiomas, and post-surgical cavities. Images are skull-stripped, aligned to isotropic $$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$4, and z-score normalized separately for non-CE and CE images. The split is 80% training, 15% validation, and 5% held-out test, with 360 test volumes (Brandstötter et al., 18 Aug 2025).

Because full-volume 3D training is memory-intensive, both stages are trained on random $$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$5 patches and run patch-wise at inference. Full volumes are reconstructed with overlapping patches and Hann-window blending to reduce seams. The paper states that larger overlap improves both fidelity and the quality of the refinement stage (Brandstötter et al., 18 Aug 2025).

Method	MSE / PSNR / SSIM	FID / KID
Residual U-Net (posterior mean)	0.0450 / 34.82 / 0.9524	39.76 / 0.0314
RF baseline	0.0623 / 33.30 / 0.9434	11.66 / 0.0068
PMRF (ours)	0.0570 / 33.64 / 0.9442	12.46 / 0.0071

These results locate the three methods at distinct points on the perception–distortion trade-off. The posterior-mean U-Net gives the best MSE, PSNR, and SSIM but poor FID and KID. The RF baseline gives the best perceptual scores but worst distortion. PMRF provides an intermediate operating point: FID drops from $$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$6 to $$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$7 relative to the posterior mean, about $$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$8, while MSE rises from $$\omega^* = \arg\min_{\omega} \mathbb{E}\big[\|X - f_\omega(Y)\|^2\big].$$9 to $$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$0, about $$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$1 (Brandstötter et al., 18 Aug 2025).

The ablation study varies patch overlap and the number of Euler steps $$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$2. The tested overlaps are 4, 16, and 32 voxels, and the tested step counts are $$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$3. The best reported configuration is overlap $$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$4 and $$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$5, with MSE $$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$6, PSNR $$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$7, SSIM $$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$8, FID $$Z_0 = f_{\omega^*}(Y) + \sigma_s \epsilon,\quad \epsilon \sim \mathcal{N}(0,I),$$9, and KID $v_\theta$0. The paper also reports that, for overlap 4, FID improves until around $v_\theta$1 and then worsens; for overlap 16, the best FID occurs around $v_\theta$2; and for overlap 32, increasing $v_\theta$3 keeps improving or maintaining quality through $v_\theta$4 (Brandstötter et al., 18 Aug 2025).

Qualitatively, the posterior-mean U-Net localizes lesions correctly and estimates overall enhancement but appears overly smooth and “denoised.” The RF baseline produces sharper-looking images but can miss lesion enhancement strength or fine anatomical detail. PMRF sharpens lesion boundaries, restores subtle vascular structures, recovers rim heterogeneity and margin details, and remains conservative in the sense that it can only refine what the posterior mean already suggests; if Stage 1 suppresses a subtle enhancement pattern, Stage 2 cannot reliably invent it (Brandstötter et al., 18 Aug 2025).

6. Latent-space reformulation and subsequent developments

Latent-PMRF reformulates PMRF in the latent space of a VAE and is motivated by the claim that pixel-space modeling limits PMRF’s ability to align with human perception (Luo et al., 1 Jul 2025). The extension preserves the two-stage logic: first estimate the posterior mean $v_\theta$5, then encode it and learn a rectified flow from that latent source to the latent distribution of high-quality images.

A central technical issue in the latent formulation is the choice of source distribution. The paper contrasts $v_\theta$6 with $v_\theta$7 and argues that the second option should be used. Its stated reason is that $v_\theta$8 has zero error relative to the posterior mean if the VAE reconstructs perfectly, whereas $v_\theta$9 would require linearity properties that do not hold for deep nonlinear VAEs (Luo et al., 1 Jul 2025). The paper therefore states that the minimum distortion is bounded by the VAE’s reconstruction error, making VAE design central rather than incidental.

The latent method introduces Sim-VAE, with the following reported design changes: a simplified ResBlock, Pixel-wise LayerNorm instead of GroupNorm, removal of self-attention in middle layers in favor of $Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$0 convolutions, redistribution of parameters in resizing layers, and removal of adversarial loss when the VAE is sufficiently strong. The reported training loss is

$Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$1

with $Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$2 (Luo et al., 1 Jul 2025).

Empirically, Latent-PMRF is reported to achieve a $Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$3 speedup over PMRF in terms of FID and a $Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$4 speedup in DISTS. On CelebA-Test, the reported comparison is PSNR $Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$5 versus $Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$6, MS-SSIM $Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$7 versus $Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$8, LPIPS $Z_t = tX + (1-t)Z_0,\quad t \sim U[0,1].$9 versus $$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$0, DISTS $$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$1 versus $$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$2, FID $$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$3 versus $$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$4, $$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$5 $$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$6 versus $$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$7, and $$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$8 $$\theta^*=\arg\min_\theta \int_0^1 \mathbb{E}\Big[\big\|(X-Z_0)-v_\theta(Z_t,t)\big\|^2\Big]\,dt.$$9 versus $Y$00, with the paper interpreting the overall result as an improved perception–distortion trade-off and much faster perceptual convergence (Luo et al., 1 Jul 2025).

The latent-space extension also introduces a limitation absent from the original motivation for leaving pixel space: test-time runtime is slightly worse than PMRF because encoding and decoding with the VAE add overhead (Luo et al., 1 Jul 2025). This suggests a broader interpretation of PMRF’s trajectory: the original framework supplies a theory-driven decomposition of fidelity and realism, while later work explores which representation space makes that decomposition most effective for a given application.