Zero-Shot Brain MRI Inpainting via 2.5D Unconditional Flow Priors

Why Brain MRI Needs Lesion "Inpainting"

Generative inpainting of brain MRI is playing an increasingly critical role in automated medical image analysis pipelines. When scans contain pathological regions such as tumors or lesions, downstream tasks like image registration, skull stripping, and segmentation are often disrupted, leading to errors. By synthesizing "healthy" brain tissue in pathological regions, the accuracy and reliability of these automated analysis applications can be significantly improved.

This paper published on arXiv (arXiv:2610.08983v1) addresses this need by proposing a zero-shot brain MRI inpainting framework based on 2.5D unconditional flow priors, aiming to strike a better balance between computational efficiency and output quality.

The Dilemma of Existing Approaches

Current brain MRI inpainting methods follow two main technical routes, each with significant drawbacks.

Full 3D methods preserve volumetric consistency well, but their computational cost is prohibitively high — the overhead of 3D convolutions makes them impractical in real-world applications. Efficient 2D slice-wise methods are fast, but suffer from severe inter-slice discontinuities: without spatial context constraints between adjacent slices, the reconstructed 3D structure exhibits abrupt breaks and jarring transitions between layers.

Beyond this, conventional models typically rely on conditional training, meaning they must learn from masked inputs for specific tasks, which limits generalization and flexibility. Each type of inpainting task may require retraining from scratch — a costly proposition.

2.5D Flow Priors: A Middle Ground Between 2D and 3D

The core innovation of this work lies in using a "2.5D" strategy to capture spatial context. Without incurring the full overhead of 3D convolutions, the approach models inter-slice spatial relationships along the superior-inferior axis.

Training Phase: Learning the Manifold of Healthy Anatomy

During training, the research team employs a flow matching model to learn the joint distribution of adjacent axial slice triplets. Rather than treating each slice in isolation, the model processes groups of three adjacent slices, capturing continuity information along the depth direction.

A key design choice is that the model explicitly excludes pathological regions from the loss function, modeling only healthy brain anatomy. This means the model learns the manifold of "what a healthy brain should look like," rather than being contaminated by lesion data. This is the foundation for zero-shot inpainting — the model does not need to have seen the specific inpainting task it will be applied to.

Flow Matching is a generative modeling paradigm that has gained traction in recent years. Like diffusion models, it belongs to the family of continuous-time generative models, but its training objective is more direct: it learns a velocity field that "flows" a simple noise distribution toward the target data distribution, and inference is performed by numerically integrating along this vector field. Compared to diffusion models that require hundreds or even thousands of denoising steps, flow matching typically requires far fewer function evaluations, making inference more efficient. In medical imaging, another advantage of flow matching is that its training objective is less sensitive to discrete step counts, which helps produce stable training signals on small datasets — particularly valuable for brain MRI where annotation costs are high. This paper applies the flow matching model to model the joint distribution of axial slice triplets, so the model learns not only the appearance statistics of individual slices but also implicitly encodes structural constraints between adjacent slices.

Inference Phase: Autoregressive Processing with Mask Constraints

At inference time, the model processes input triplets autoregressively along the depth axis, ensuring information propagates between layers. To adapt the unconditional prior to a specific inpainting task, the team introduces the Restora-Flow solver, which uses the input mask to impose constraints on the unconditional prior — enabling precise zero-shot inpainting without any task-specific training.

As a final step, the framework generates an ensemble of stochastic reconstructions and averages them as the final prediction, further improving stability.

The core challenge of zero-shot inpainting is: how can an unconditional generative model — one that has never seen masked-completion pairs — be made to respect pixel constraints outside the masked region at inference time? The approach used by the Restora-Flow solver is conceptually similar to the "replacement trick" or "posterior sampling" techniques in the diffusion model literature: at each iteration step, pixel values in the known region are "anchored" back to the observed values, while only the unknown region inside the mask is free to evolve along the generative trajectory. By repeatedly applying this constraint, the generation process converges to a reconstruction that is consistent with the known region while conforming to the healthy anatomy prior. Compared to conditional training, this inference-time constraint approach may require more convergence steps, but in return the model only needs to be trained once and can then adapt to masks of arbitrary shapes — a major gain in flexibility.

Experimental Results and Evaluation

The paper benchmarks the method on the official BraTS 2026 Inpainting Challenge validation set, with the following quantitative results:

• SSIM (Structural Similarity Index): 0.816 ± 0.112
• MSE (Mean Squared Error): 0.007 ± 0.005
• PSNR (Peak Signal-to-Noise Ratio): 22.923 ± 4.343

Qualitatively, the 2.5D strategy successfully resolves the structural discontinuity issues of 2D baseline methods, synthesizing plausible healthy tissue while maintaining volumetric consistency across all three planes — axial, sagittal, and coronal. This is precisely where purely 2D methods struggle the most.

The research team has open-sourced their code at https://github.com/imigraz/brats2026-inpainting for community reproduction and further research.

BraTS (Brain Tumor Segmentation) is one of the most influential public benchmarks in medical imaging, held annually since 2012. The dataset includes multi-modal MRI sequences (T1, T1ce, T2, FLAIR) and expert-annotated tumor segmentation masks. The Inpainting track newly introduced in the 2026 edition specifically evaluates the ability to synthesize healthy tissue in tumor regions. The evaluation logic works as follows: using known healthy images as reference, the tumor mask is applied to occlude the region, the model is asked to restore it, and the result is compared against the original healthy image using metrics like SSIM, MSE, and PSNR. SSIM (Structural Similarity Index) is sensitive to brightness, contrast, and structural changes as perceived by the human eye, with a range of 0 to 1 (higher is better). PSNR measures reconstruction error in decibels; for typical medical image reconstruction tasks, 20–30 dB represents a moderate level of performance. The public leaderboard on this validation set enables cross-method comparisons and provides a unified evaluation benchmark for future research.

Technical Significance and Outlook

The value of this work lies in its use of a relatively lightweight "2.5D" approach to sidestep the computational bottleneck of 3D methods while overcoming the continuity limitations of 2D methods. The combination of an unconditional prior with zero-shot inference also frees the model from dependence on task-specific conditional training — train once on the distribution of healthy anatomy, then adapt to different inpainting scenarios via mask constraints.

For medical image analysis, this efficient yet volumetrically consistent inpainting approach has real potential to improve the reliability of downstream tasks like registration and segmentation. The deployment of flow matching generative models in medical imaging once again demonstrates the potential of such generative methods to transfer from natural images to specialized domains.