---
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}
}
```