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

arXiv: https://arxiv.org/abs/2607.27974 · HTML: https://mlsanity.com/paper/2607.27974/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 26 references that are in the map.
- First posted on arXiv 30 Jul 2026.

## Builds on (most cited first, 17 of 17)

- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time](https://mlsanity.com/paper/2203.05482/) (2022)
- [Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network](https://mlsanity.com/paper/1609.05158/) (2016)
- [The RSNA-ASNR-MICCAI BraTS 2021 Benchmark on Brain Tumor Segmentation and Radiogenomic Classification](https://mlsanity.com/paper/2107.02314/) (2021)
- [Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks](https://mlsanity.com/paper/1807.07356/) (2018)
- [Scalable High-Resolution Pixel-Space Image Synthesis with Hourglass Diffusion Transformers](https://mlsanity.com/paper/2401.11605/) (2024)
- [MedPerf: Open Benchmarking Platform for Medical Artificial Intelligence using Federated Evaluation](https://mlsanity.com/paper/2110.01406/) (2021)
- [U-DiTs: Downsample Tokens in U-Shaped Diffusion Transformers](https://mlsanity.com/paper/2405.02730/) (2024)
- [Denoising Diffusion Models for 3D Healthy Brain Tissue Inpainting](https://mlsanity.com/paper/2403.14499/) (2024)
- [fastWDM3D: Fast and Accurate 3D Healthy Tissue Inpainting](https://mlsanity.com/paper/2507.13146/) (2025)
- [Brain Tumour Removing and Missing Modality Generation using 3D WDM](https://mlsanity.com/paper/2411.04630/) (2024)
- [Symmetry Awareness Encoded Deep Learning Framework for Brain Imaging Analysis](https://mlsanity.com/paper/2407.08948/) (2024)
- [Robust 3D Brain MRI Inpainting with Random Masking Augmentation](https://mlsanity.com/paper/2511.20202/) (2025)
- [WaveDiT: Distribution-Aware Wavelet Flow Matching for Efficient 3D Brain MRI Synthesis](https://mlsanity.com/paper/2606.08670/) (2026)

## Abstract (arXiv)

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×10-3 over 219 cases. We further analyse the residual smoothness inherent to distortion-optimal regression and discuss its implications for anatomical realism.