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

2026-10-06

AI Summary

This study addresses the computational expense of 3D brain MRI inpainting and the inter-slice discontinuity inherent in 2D approaches by proposing a zero-shot, training-free inpainting framework. The method leverages a 2.5D unconditional flow matching prior to capture spatial context, processing adjacent slice triplets autoregressively. A dedicated Restora-Flow solver is designed to enforce input masking constraints, explicitly excluding pathological regions from the loss function to synthesize healthy tissue. This architecture effectively balances computational efficiency with volumetric consistency while avoiding the overhead of full 3D convolutions. Evaluated on the BraTS 2026 validation set, the model achieves an SSIM of 0.816, an MSE of 0.007, and a PSNR of 22.923 dB, significantly preserving multi-planar structural continuity.

Abstract

Generative inpainting of brain MRI volumes is essential for synthesizing healthy tissue in pathological regions, improving the accuracy and reliability of automated downstream brain analysis applications such as image registration, brain extraction, and segmentation. However, standard 3D approaches are computationally prohibitive, while efficient 2D slice-wise methods suffer from severe inter-slice discontinuities. Furthermore, traditional models rely on conditional training, requiring task-specific learning of masked inputs. We propose a zero-shot brain MRI inpainting framework utilizing 2.5D unconditional flow priors to capture spatial context along the superior-inferior axis without the overhead of full 3D convolutions. During training, our flow matching model learns the joint distribution of adjacent axial slice triplets, modeling the manifold of healthy brain anatomy while explicitly excluding pathological regions from the loss function. At inference, the model processes the input triplets autoregressively along the depth axis. We employ the Restora-Flow solver to constrain the unconditional prior using the input mask, achieving accurate zero-shot inpainting. Evaluations show our 2.5D strategy resolves the structural discontinuities of 2D baselines, synthesizing plausible healthy tissue while maintaining volumetric consistency across the axial, sagittal, and coronal planes. As a final step, we generate and average an ensemble of multiple stochastic reconstructions to form the final prediction. Quantitative results benchmarked on the official BraTS 2026 Inpainting Challenge validation set demonstrate the effectiveness of our proposed approach, yielding an SSIM of 0.816 $\pm$ 0.112, MSE of 0.007 $\pm$ 0.005, and PSNR of 22.923 $\pm$ 4.343. Code is available at https://github.com/imigraz/brats2026-inpainting.

Problem

• Brain MRI inpainting
• Zero-shot generation
• Inter-slice discontinuity
• Computational efficiency
• Unconditional prior

Innovation

• Zero-shot inpainting
• 2.5D unconditional flow priors
• Flow matching
• Brain MRI
• Restora-Flow solver