UniBrain: Unified Multimodal Model for Brain MRI Imputation and Understanding

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.

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

Item	Description
Base model	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, using the UniBrain preprocessed release
Inference precision	BF16

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

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