Now You Have My Healthy Attention: A U-DiT for Brain-MRI Inpainting

2026-07-30

AI Summary

This study addresses the challenge of missing healthy tissue in brain T1-weighted MRI caused by tumor occlusion by proposing a deterministic regression model based on the U-DiT architecture to generate anatomically plausible tumor-free reference images. The method incorporates 3D rotary positional encoding and global self-attention on a downsampled token grid, synergistically combining convolutional operations, skip connections, and a mirror prior from the contralateral hemisphere. An attention constraint is specifically designed to restrict the model’s focus to known healthy regions within the same volume, thereby enhancing structural symmetry and reconstruction fidelity. Evaluated on the BraTS-2026 validation set (219 cases), the model achieves SSIM = 0.864, PSNR = 24.7 dB, and MSE = 4.6 × 10⁻³ in healthy regions, significantly improving the reliability of downstream analyses.

Abstract

The ASNR-MICCAI BraTS Local Synthesis (Inpainting) task asks for the anatomically plausible completion of healthy brain tissue within a masked region of a T1-weighted MRI, providing a tumor-free anatomical reference for downstream analysis. As the task is scored by distortion metrics (SSIM, PSNR, MSE), we build a deterministic regression model and focus on giving it inductive biases tailored to inpainting. Our network follows the U-DiT principle of performing self-attention on a downsampled token grid: a volumetric encoder-decoder imports long-range context through a downsampled global self-attention block with three-dimensional rotary position embeddings, while convolutions and skip connections preserve high-frequency detail. Two ideas drive our results. First, we constrain the attention so that occluded ("void") tokens attend only to known-healthy tokens of the same volume, with a learned bias toward each query's contralateral homologue, forcing the completion to be inferred from observed anatomy rather than from other unknown regions. Second, we add a contralateral-symmetry input that supplies the mirrored healthy hemisphere as a patient-specific prior; since the brain is approximately bilaterally symmetric and lesions are typically unilateral, this prior improves the distortion metrics at matched structural similarity. On the official BraTS-2026 validation leaderboard our submission reaches a mean healthy-region SSIM of $0.864$, PSNR of $24.7$\,dB and MSE of $4.6{\times}10^{-3}$ over $219$ cases. We further analyse the residual smoothness inherent to distortion-optimal regression and discuss its implications for anatomical realism.

Problem

• brain-MRI inpainting
• healthy tissue completion
• tumor-free reference
• anatomical plausibility
• medical image synthesis

Innovation

• U-DiT
• brain-MRI inpainting
• constrained self-attention
• contralateral symmetry prior
• anatomically plausible completion