Brainμ-Med: Foundation Model Framework for Neuroscience and Neurological Medicine

A unified next-token prediction framework for multimodal brain activity and neurological disease modeling

Overview

Brainμ-Med is a foundational large-scale model framework for neuroscience and neurological medicine, developed by the Beijing Academy of Artificial Intelligence (BAAI).
It is built upon the “Next-token Prediction” multimodal EMU3 framework and fine-tuned from the Brainμ foundation model using large-scale open-source medical data.

At its core, Brainμ-Med provides a unified representation and modeling framework for diverse brain-related signals and medical information, enabling cross-modality, cross-task, cross-species, and cross-individual analysis within a single model.

The project aims to:

• Bridge neuroscience and clinical medicine through a unified token-based modeling paradigm.
• Support medical assessment of neurological disorders via multimodal neuro-data understanding.
• Provide an open, extensible framework for research in brain decoding, disease diagnosis, prognosis prediction, and treatment response modeling.

---

Key Features

• Foundation model for neuroscience & neuromedicine
• Tokenization-based multimodal brain representation
• Multimodal alignment on EMU3
• Medical-focused capabilities
• Cross-modality & cross-population generalization

---

Conceptual Architecture

The Brainμ-Med framework can be conceptually divided into three main components:

---

Medical and Clinical Applications

Brainμ-Med is designed as a general-purpose backbone for medical applications involving the brain, with particular emphasis on neurological and psychiatric disorders. Example application areas include:

• Disease Diagnosis & Classification
• Prognosis and Outcome Prediction
• Cognitive and Behavioral Assessment
• Cross-Dataset and Cross-Site Generalization

Brainμ-Med has demonstrated SOTA performance on multiple benchmark tasks involving brain disorders, suggesting strong potential as a foundation infrastructure for AI-assisted neuromedicine.

---

Example Use Cases

Below are conceptual examples of how Brainμ-Med can be used in research and clinical workflows:

• Researcher Workflow
• Clinical Decision Support (Research Setting)
• Data Integration Across Studies

---

Quickstart: Neurological Disease Assessment

This section provides a minimal, disease-focused quickstart showing how Brainμ-style tokenization and Brainμ-Med disease heads can be used to assess Alzheimer’s disease (AD), depression (DE), and Parkinson’s disease (PD) from preprocessed EEG signals.

Note: The paths and model names below are illustrative and should be adapted to your actual environment and checkpoint locations.

1. Prepare EEG and Electrode Coordinate Data

• A preprocessed EEG segment (filtered, normalized, and cropped) with shape ((C, T)) where (C) is the number of channels and (T) is the number of time steps.

Example:

import numpy as np

eeg_path = "/path/to/eeg_segment.npy"          # shape: (C, T)

eeg = np.load(eeg_path).astype("float32")        # (C, T)

2. Load the Brain Tokenizer and Disease-Specific Heads

We use a Brain Tokenizer (BRVQ) to encode EEG into discrete brain tokens, and then plug in disease-specific heads for AD / DE / PD classification.

import torch
from braintoken import get_brvq_model

device = "cuda" if torch.cuda.is_available() else "cpu"

Load Brain Tokenizer backbone
brain_tokenizer = get_brvq_model("stageII", device=device, model='eeg')

3. Run Disease Predictions (AD / Depression / PD)

After loading the checkpoints, you can obtain class probabilities and predicted labels for each disease:

# Alzheimer’s disease (AD)
probs_ad, pred_ad = forward_eeg(eeg, coordinates, disease="AD")
print("AD probabilities:", probs_ad, "pred class:", pred_ad)

# Depression (DE)
probs_de, pred_de = forward_eeg(eeg, coordinates, disease="DE")
print("Depression probabilities:", probs_de, "pred class:", pred_de)

# Parkinson’s disease (PD)
probs_pd, pred_pd = forward_eeg(eeg, coordinates, disease="PD")
print("PD probabilities:", probs_pd, "pred class:", pred_pd)

You can wrap this logic into a script such as:

python disease_inference.py \
  --eeg_path /path/to/eeg_segment.npy \
  --model_ckpt /path/to/model.ckpt \

to build a simple command-line tool for research-oriented neurological disease assessment from EEG.

---

Datasets and Performance

Brainμ-Med is fine-tuned with open-source medical data, which may include:

• Public neuroimaging datasets (e.g., brain disorder cohorts, developmental/aging cohorts).
• Open neurological and psychiatric datasets with both imaging and clinical labels.
• Publicly available clinical scales, cognitive assessments, and symptom questionnaires.

Brainμ-Med has achieved state-of-the-art (SOTA) performance on multiple brain disorder tasks, such as:

• Distinguishing patients with neurological or psychiatric disorders from healthy controls.
• Classifying disorder subtypes or stages.
• Predicting clinical scores from brain activity patterns.

---

License

The license and terms of use for Brainμ-Med are determined by the project maintainers.
If a LICENSE file is present in this repository, please refer to it for details regarding:

---

• Research usage
• Commercial usage
• Data and model redistribution

Acknowledgements

Brainμ-Med is developed by the Beijing Academy of Artificial Intelligence (BAAI), building on:

• The Brainμ foundation model for multimodal brain activity modeling.
• The EMU3 framework for next-token prediction in multimodal large language models.

We thank the open-source neuroscience and medical communities for providing datasets, tools, and benchmarks that make this line of research possible.

This work is supported by the National Science and Technology Major Project (No. 2022ZD0116314).

本项目受新一代人工智能国家科技重大专项（No. 2022ZD0116314）支持。