File size: 4,490 Bytes
eae8616 55263a9 eae8616 55263a9 65cf414 55263a9 65cf414 55263a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | ---
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}
}
```
|