RARF: Region-Aware Rectified Flows for 3D Brain MRI Inpainting

AI summary

Overview

• Research area: Medical image analysis and generative computer vision, specifically 3D brain MRI inpainting with rectified flows.
• Technical level: Advanced. The paper assumes familiarity with flow matching, diffusion sampling, U-Net architectures, and medical image evaluation metrics.
• Scope: The paper introduces RARF, a region-aware rectified flow framework for masked data generation, and instantiates it as a submission to the BraTS Inpainting Challenge 2026 for synthesizing healthy brain tissue inside tumor regions.

What This Paper Is About

Many automated brain MRI analysis tools assume the input shows healthy anatomy, which breaks down in neuro-oncology where lesions are present and healthy images of the same patient do not exist. The BraTS Local Inpainting task asks a model to fill in a masked tumor-affected region with anatomically plausible, tumor-free tissue. RARF addresses this by restricting a rectified-flow generative process to the masked region only, leaving the observed anatomy untouched as patient-specific context.

Key Contributions

1. A task-agnostic region-aware rectified flow framework. RARF defines a localized interpolation path in which only the masked region evolves along the flow, while visible voxels are copied from the voided input and remain fixed during both training and inference. The paper states the framework supports arbitrary inpainting masks and multiple training and inference strategies.

2. A 3D region-aware flow-matching objective with auxiliary supervised terms. The training loss combines a masked flow-matching term with a normalized endpoint mean-absolute-error term and a masked structural-similarity term, with supervision restricted to the healthy-tissue mask rather than the full inpainting mask.

3. An inference-strategy ablation on the official BraTS validation platform. Using a single trained checkpoint with an exponential moving average of parameters (decay 0.999), the authors compare four-sample strategies and regular single-sample inference: 30-sample averaging, 30-sample closest to the mean, kernel density steering (KDS), minimum Bayes-risk (MBR) selection, and regular inference.

4. A released challenge submission. The main submission uses a separately trained checkpoint with averaging of 50 samples per case, and the model weights and source code are publicly released.

Main Findings

• Averaging many samples gives the best distortion metrics. On the official BraTS validation set with 30 samples, averaging achieved MSE 0.006, MAE 0.018, PSNR 23.830, and SSIM 0.823.

• Representative-sample selection improves over single-sample inference but not over averaging. The 30-sample closest to the mean gave MSE 0.009, MAE 0.021, PSNR 21.911, SSIM 0.780; KDS gave MSE 0.009, MAE 0.021, PSNR 22.096, SSIM 0.784; MBR gave MSE 0.009, MAE 0.022, PSNR 21.982, SSIM 0.781.

• Regular inference was the weakest in the ablation. Regular single-sample inference gave MSE 0.010, MAE 0.022, PSNR 21.942, SSIM 0.777.

• Averaging improves fidelity but smooths structure. The authors report that averaging produces smoother reconstructions that may blur fine anatomical structures, which they frame as the distortion-perception trade-off.

• Multi-sample strategies cost more inference time. With K = 30, all multi-sample strategies required substantially greater inference time than single-sample inference.

• Official challenge results. The released checkpoint with K = 50 sample averaging achieved MSE 0.006, PSNR 24.008, and SSIM 0.832 on the official BraTS validation set; MAE is not reported for this submission.

• Distortion metrics may not track perceptual quality. The ablation showed that lower distortion did not necessarily imply better perceptual quality or anatomical plausibility, since blurrier reconstructions could still rank above sharper predictions.

Methodology in Plain English

The authors start from the BraTS Local Inpainting dataset: 1,251 training and 219 validation skull-stripped, co-registered T1-weighted MRI volumes, all shaped 240 × 240 × 155 voxels at 1 mm isotropic spacing. Each case comes with a healthy tissue mask, an unhealthy tissue mask, and their union, which defines the region to be filled in. The input image is voided over that union, and supervision during training is restricted to the healthy mask because tumor intensities are not valid targets. Ground-truth volumes for the validation set are not publicly available, so the authors monitor overfitting by placing synthetic masks over visible healthy tissue in the challenge validation cases. To augment training masks, they generate five mask variants per case, including the original, from transformed connected components of training-set lesions placed under constraints on lesion distance, brain coverage, overlap, and diversity.

Preprocessing crops each volume identically using only information available at inference: the foreground is the union of visible nonzero voxels in the voided image and the inpainting mask, and the crop is centered on the midpoint of its bounding box, then cropped or zero-padded to 166 × 196 × 152 voxels. Intensities are normalized per case by computing the 0.5th and 99.5th percentiles of the complete voided volume before cropping, clamping the lower bound to be nonnegative, clipping, and linearly mapping to [0, 1].

The generative model itself works by defining a path between Gaussian noise and the clean image, but only inside the mask. Visible voxels are simply copied from the voided input at every timestep, which is expressed through a spatial time map that assigns time t to masked voxels and fixes visible voxels at the clean endpoint. A 3D U-Net takes a two-channel input (the current state and the complete mask) plus an embedding of the scalar flow time, and predicts a single-channel velocity volume. The U-Net has three resolution levels with channel widths 32, 64, and 128, two residual blocks per level using group normalization, SiLU activations, and 3 × 3 × 3 convolutions, with stride-two downsampling and nearest-neighbor upsampling. The flow time is scaled by 1000 and sinusoidally encoded, then mapped by a two-layer multilayer perceptron to a 128-dimensional representation added to every residual block.

Training minimizes three terms: the masked flow-matching mean-squared error against the target velocity, a masked mean-absolute error on a length-normalized endpoint error with a clamp of 10^-3 and a time weight of 1 + 0.5t², and a masked structural similarity loss. Training times are sampled from a mixture: with probability 0.75 uniformly on [0, 1], otherwise as t = 1 − u² to emphasize examples near the clean endpoint. At inference, the masked region starts as Gaussian noise and the learned velocity is integrated from t = 0 to t = 1 using four integration time points with midpoint integration, restoring observed voxels after every update. Because different noise draws give different plausible completions, the authors also explore averaging K samples, and selecting a representative sample via closeness to the mean, kernel density steering, or minimum Bayes-risk selection.

Why This Matters

The work targets a practical obstacle in clinical neuroimaging pipelines: downstream tools for brain extraction, tissue segmentation, and anatomical parcellation can behave unreliably when a lesion is present, and synthesizing a subject-specific healthy-tissue proxy may reduce that pathology-induced bias. Beyond the specific challenge entry, the paper positions RARF as a general region-aware framework applicable to arbitrary inpainting masks, arguing that medical imaging can both benefit from and contribute to general computer vision advances.

Real-world applications:

• Lesion filling in clinical MRI. Replacing tumor-affected voxels with plausible healthy tissue so that automated segmentation and parcellation tools can run on anatomy-consistent inputs.
• Subject-specific anatomical reference generation. Producing a healthy-tissue proxy for an individual patient, rather than relying on atlas or population averages.
• Medical imaging benchmark development. Providing a standardized, high-dimensional testbed where masked regions correspond to a specific underlying anatomy rather than one of many equally valid completions.
• General masked-region image generation. Applying the region-aware rectified flow formulation to other inpainting problems where observed context must be preserved exactly.

Industry relevance: the framework's inference-time flexibility matters commercially, because the paper shows that a single trained checkpoint can be paired with different sample-selection strategies to trade off distortion accuracy against perceptual sharpness and runtime. The reported fact that all multi-sample strategies at K = 30 require substantially greater inference time is directly relevant to any deployment decision, and the released code and weights lower the barrier for groups building on this method.

Future Directions

• Extending the ablation beyond inference. The authors explicitly state that ablations of the broader training modes supported by RARF are left for future work, as they fall outside the scope of this challenge submission.
• Complementing voxel-wise metrics. The paper motivates adding perceptual, anatomy-aware, and downstream-task evaluation, since lower distortion did not necessarily mean better anatomical plausibility.
• Balancing the distortion-perception trade-off. Averaging 30 or 50 samples improved MSE and PSNR but blurred fine structures, leaving open how to obtain sharp reconstructions that also score well on distortion metrics.
• Testing RARF beyond BraTS. The framework is described as task-agnostic and applicable to arbitrary inpainting masks, but the paper only instantiates and evaluates it on the BraTS Local Inpainting dataset.

Target Audience

This paper is most useful to researchers and engineers working on medical image synthesis, generative models for volumetric data, and brain MRI analysis, particularly those interested in flow matching and diffusion alternatives. It also serves participants and organizers of inpainting and restoration challenges who want a concrete example of how region-aware generation and inference-time sample selection are evaluated under an official BraTS protocol. Readers without a background in flow-based generative models will find the methods section demanding, though the problem framing and the discussion of evaluation trade-offs are accessible to a broader imaging audience.

Authors’ abstract

Medical image inpainting has the potential to improve automated brain MRI analysis by reconstructing healthy tissue within pathological regions. We introduce RARF, a task-agnostic region-aware rectified flow framework for masked data generation. We instantiate the framework for 3D brain MRI inpainting as our submission to the BraTS Inpainting Challenge 2026. RARF restricts the stochastic interpolation process to the inpainting region, while the observed voxels remain fixed and provide patient-specific anatomical context. A three-dimensional neural network receives the partially voided image, with Gaussian noise filling the missing region, together with the inpainting mask and the corresponding timestep. The model is trained using masked flow-matching and reconstruction-consistency objectives, combined with mask-aware preprocessing and data augmentation. During inference, the learned velocity field transports the initial noise toward a plausible reconstruction of the missing tissue, which is then combined with the unchanged observed anatomy. Experiments under the BraTS evaluation protocol show that the proposed approach produces competitive reconstructions while maintaining anatomical consistency. Source code is available at: https://github.com/TomasGuija/rarf.