UniBrain: A Unified Multimodal Understanding and Completion Model for Brain MRI

UniBrain is a unified multimodal model for brain MRI that can simultaneously achieve modality completion and disease diagnosis understanding, addressing key challenges of data missingness and multimodal fusion in medical imaging. Proposed by Zhiyun Song et al., this model has been accepted by MICCAI 2026, and its PyTorch implementation is open-sourced, including training scripts, evaluation tools, and support for pre-trained models.

Multimodal Large Language Models (MLLMs) have great potential in the medical field, but they face challenges such as scarcity of high-quality training data and frequent missingness of clinical data. In brain MRI analysis, patients often fail to complete a full set of scans, leading to modality missingness. Traditional methods either discard samples or train completion models separately, making end-to-end diagnosis and understanding difficult.

Three Core Technical Innovations

1. Interleaved Description-Enhanced Data Flow: Autoregressive training enables deep integration of generation and medical reasoning, eliminating the need for completion before diagnosis.
2. Self-Alignment Strategy: Uses dense image embeddings for self-reconstruction to learn anatomical representations, reducing reliance on manual annotations.
3. Dynamic Hidden State Mechanism (DHS): Mitigates exposure bias in long-context reasoning and maintains anatomical structure consistency.

Three-Stage Training

1. Medical Reconstruction Self-Alignment: Self-supervised pre-training to learn the basics of anatomical knowledge.
2. Unified Multimodal Modeling: Uses complete datasets to learn inter-modal relationships.
3. Self-Enforced Fine-Tuning: Improves generation quality and diagnostic accuracy.

Supports standard MRI modalities such as T1n, T1c, T2w, and T2f.

UniBrain was validated on multi-disease datasets:

• Modality Completion Quality: PSNR and SSIM metrics show that generated images have accurate anatomy and distinguishable pathological features.
• Disease Diagnosis Accuracy: Remains robust when modalities are incomplete; single-modal input can also assist diagnosis through generation.
• Joint Evaluation: A positive correlation exists between generation quality and diagnostic accuracy, verifying the effectiveness of the unified strategy.

• Reduce Scanning Costs: Shortens patient scanning time and medical costs, suitable for patients with mobility issues or claustrophobia.
• Improve Diagnostic Efficiency: Assists radiologists in comprehensive evaluation and reduces information loss.
• Promote Field Development: Proves the feasibility of the unified generation-understanding framework and provides references for subsequent research.

The PyTorch implementation of UniBrain is fully open-sourced, including training scripts, evaluation tools, and support for loading pre-trained models. It is built based on BAGEL and follows the corresponding protocol. The project provides detailed documentation (environment configuration, data preparation, training process, evaluation scripts) to help research institutions reproduce and fine-tune, promoting the democratization of medical imaging AI.