Astrostellar/UniBrain

UniBrain: Unified Multimodal Model for Brain MRI Imputation and Understanding

<p align="left"> <a href="https://medicalumm.github.io/unibrain.github.io/"><img src="https://img.shields.io/badge/UniBrain-Project_Page-0A66C2?logo=safari&logoColor=white" alt="UniBrain project page"></a> <a href="https://arxiv.org/abs/2606.16484"><img src="https://img.shields.io/badge/UniBrain-Paper-red?logo=arxiv&logoColor=white" alt="UniBrain paper"></a> <a href="https://github.com/zhiyuns/UniBrain"><img src="https://img.shields.io/badge/UniBrain-Code-536af5?logo=github&logoColor=white" alt="UniBrain code"></a> </p>

UniBrain is a unified multimodal model for brain MRI analysis. In one autoregressive context, it can impute missing MRI sequences, interpret the available and generated images, and produce a disease diagnosis. This repository hosts the UniBrain model checkpoints.

For installation, training, evaluation, and usage instructions, please visit the official GitHub repository.

<p align="center"> <img src="https://github.com/zhiyuns/UniBrain/raw/main/assets/main_figure.png" alt="Overview of the UniBrain framework" width="95%"> </p>

UniBrain is initialized from BAGEL-7B-MoT, a Mixture-of-Transformer-Experts (MoT) model for multimodal understanding and generation. It adapts BAGEL to brain MRI using an interleaved, description-enriched training flow and three main ideas:

• —Unified MRI generation and understanding: missing-sequence imputation and downstream interpretation share one autoregressive context.
• —Self-alignment: medical image reconstruction provides dense supervision for fine-grained anatomical representation learning without requiring detailed captions for every image.
• —Dynamic hidden states: training conditions the model on its own generated visual context to reduce exposure bias during long multimodal sequences.

Model details

Reported results

The following results are reported on the RadGenome-Brain MRI evaluation split in the paper and project page.

MRI diagnosis and report generation

MRI modality imputation

License

The UniBrain model weights are released under the Apache License 2.0. UniBrain builds on BAGEL and AutoRG-Brain; the code, base model, incorporated components, and datasets retain their respective licenses and terms.

Acknowledgements

The implementation is adapted from BAGEL, a unified multimodal foundation model for natural images. The training and evaluation data are based on RadGenome-Brain_MRI from the AutoRG-Brain project.

Citation

If you find UniBrain useful, please cite:

@article{unibrain2026,
  title   = {Unified Multimodal Model for Brain MRI Imputation and Understanding},
  author  = {Zhiyun Song, Che Liu, Tian Xia, Avinash Kori, Wenjia Bai},
  journal = {arXiv preprint arXiv:2606.16484},
  year    = {2026}
}