| --- |
| license: apache-2.0 |
| base_model: |
| - ByteDance-Seed/BAGEL-7B-MoT |
| pipeline_tag: any-to-any |
| library_name: bagel-mot |
| language: |
| - en |
| datasets: |
| - Astrostellar/RadGenome-Brain_MRI_parquet |
| tags: |
| - medical |
| - medical-imaging |
| - brain-mri |
| - multimodal |
| - image-to-image |
| - image-text-to-text |
| - modality-imputation |
| --- |
| |
| # 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](https://github.com/zhiyuns/UniBrain). |
|
|
| <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](https://huggingface.co/ByteDance-Seed/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 |
|
|
| | Item | Description | |
| | --- | --- | |
| | Base model | [ByteDance-Seed/BAGEL-7B-MoT](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT) | |
| | Architecture | Unified MoT architecture | |
| | Domain | 2D axial brain MRI slices | |
| | Tasks | MRI modality imputation, brain MRI understanding/diagnosis | |
| | Training data | [RadGenome-Brain_MRI](https://huggingface.co/datasets/JiayuLei/RadGenome-Brain_MRI), using the [UniBrain preprocessed release](https://huggingface.co/datasets/Astrostellar/RadGenome-Brain_MRI_parquet) | |
| | Inference precision | BF16 | |
|
|
|
|
| ## Reported results |
|
|
| The following results are reported on the RadGenome-Brain MRI evaluation split in the paper and [project page](https://medicalumm.github.io/unibrain.github.io/). |
|
|
| ### MRI diagnosis and report generation |
|
|
| | Available modalities | Top-1 Acc | ROUGE | |
| | --- | ---: | ---: | |
| | T1w only | 74.47 | 36.93 | |
| | T1w + T2w | 76.60 | 38.23 | |
| | T1w + T2w + T2-FLAIR | 78.01 | 38.68 | |
| | Complete data | 82.06 | 38.94 | |
|
|
| ### MRI modality imputation |
|
|
| | Imputation sequence | PSNR | Top-1 Acc | |
| | --- | ---: | ---: | |
| | T1w → T2w | 22.23 | 68.09 | |
| | T1w, T2w → T2-FLAIR | 22.58 | 67.38 | |
| | T1w, T2w, T2-FLAIR → T1c | 22.26 | 74.47 | |
|
|
|
|
| ## License |
|
|
| The UniBrain model weights are released under the Apache License 2.0. UniBrain builds on [BAGEL](https://github.com/ByteDance-Seed/Bagel) and [AutoRG-Brain](https://github.com/ljy19970415/AutoRG-Brain); the code, base model, incorporated components, and datasets retain their respective licenses and terms. |
|
|
| ## Acknowledgements |
|
|
| The implementation is adapted from [BAGEL](https://github.com/ByteDance-Seed/Bagel), a unified multimodal foundation model for natural images. The training and evaluation data are based on [RadGenome-Brain_MRI](https://huggingface.co/datasets/JiayuLei/RadGenome-Brain_MRI) from the [AutoRG-Brain](https://github.com/ljy19970415/AutoRG-Brain) project. |
|
|
| ## Citation |
|
|
| If you find UniBrain useful, please cite: |
|
|
| ```bibtex |
| @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} |
| } |
| ``` |
|
|
|
|
