Brain-Semantoks: Learning Semantic Tokens of Brain Dynamics with a Self-Distilled Foundation Model

AI summary

Overview

Research area: Self-supervised learning for neuroimaging, specifically foundation models for fMRI brain dynamics. Technical level: Intermediate — the core ideas are accessible, but the paper assumes some familiarity with self-supervised learning (SSL), transformers, and fMRI analysis. Scope: A new pretraining framework called Brain-Semantoks that learns abstract, temporally stable representations of brain dynamics using semantic network tokens and self-distillation, showing strong linear-probe performance and the first scaling analysis of its kind.

What This Paper Is About

Existing fMRI foundation models (e.g., BrainLM, Brain-JEPA) train by masking and reconstructing the BOLD signal at the level of individual brain regions — an objective that forces the model to encode noise and low-level temporal fluctuations, producing representations that require heavy fine-tuning downstream. The authors argue that predicting stable phenotypes (like age, sex, or disease) actually requires abstraction away from the raw signal, not faithful reconstruction. Brain-Semantoks proposes to fix this by restructuring both the input (aggregating ROIs into functional-network tokens) and the training objective (self-distillation across time rather than reconstruction).

Key Contributions

1. A semantic tokenizer grounded in a neuroscientific prior: instead of treating each ROI as a token, it aggregates ROIs within each functional brain network into a compact, robust token, producing a shorter and more meaningful sequence for the transformer.
2. A self-distillation pretraining objective that aligns global representations across long temporal views of the same scan, replacing mask-and-reconstruct learning.
3. The Teacher-guided Temporal Regularizer (TTR) — a novel curriculum that guides the student to first learn the time-averaged signature of each network before modeling temporal variation, resolving the training instability that arises from applying distillation to low signal-to-noise fMRI data.
4. The first detailed scaling analysis for fMRI foundation models, showing predictable power-law gains both in- and out-of-distribution as pretraining data grows, without any domain adaptation.

Main Findings

• Strong linear-probe performance: Using only a frozen encoder and a single linear layer, Brain-Semantoks achieves the best accuracy on 8 of 9 downstream tasks, significantly outperforming BrainLM and Brain-JEPA — particularly on out-of-distribution clinical datasets (e.g., Schizophrenia 69.26% vs ~56–58%, MDD 62.60% vs ~53–58%, ASD 65.13% vs ~53–54%).
• Beats fully supervised baselines: Without any fine-tuning, the model surpasses end-to-end supervised models (e.g., BolT, BrainMass) on 8 diverse tasks, and fine-tuning yields only marginal improvements — evidence of well-disentangled representations.
• Strong scaling laws: Performance on nearly all tasks improves with the logarithm of pretraining data size, including large OOD gains (e.g., on HBN, whose age range differs by over 20 years from the UKB pretraining cohort).
• Task-based fMRI transfer: On the Hariri emotion task, the model predicts block type from short segments with 93.8–97.9% balanced accuracy versus 81.5–94.7% for Brain-JEPA.
• Built-in interpretability: Because the model is trained with slice-masking (hiding entire networks or time slices), network importance can be probed in-distribution. Findings align with neuroscience — e.g., default-mode network for ASD, and surprisingly, cerebellar activity as more predictive for MDD than the default-mode network.
• Computationally feasible: Joint end-to-end pretraining of tokenizer and encoder takes under two hours on a single GPU with <20 GB memory.

Methodology in Plain English

The authors treat an fMRI scan as a sequence of tokens, much like words in a sentence. Rather than using each noisy ROI as a token, they group ROIs into 9 functional brain networks (Yeo-7 plus subcortical and cerebellar) and use small convolutional modules, each specialized for one network, to compress each network's time series into a handful of rich patch tokens. These feed a transformer encoder.

Training uses a student–teacher setup: the same scan is cut into two long, overlapping temporal views, lightly corrupted (random zeroing, Gaussian noise, amplitude scaling), and the student — which sees a masked version of its view — is trained to match the teacher's outputs. The teacher's weights are a momentum-updated average of the student's, providing stable targets. Three losses drive learning:

1. A CLS loss that aligns global summary representations across views (with a coding-rate regularizer to prevent collapse).
2. A token loss that asks the student to predict masked token embeddings, using slice masking (hiding entire networks or time blocks) to discourage trivial interpolation.
3. The TTR loss, active only in the first 5% of training, which asks the student to match the teacher's time-averaged network tokens — forcing a stable initial solution before temporal detail is learned.

After pretraining, the teacher encoder is frozen and evaluated via simple linear probing on downstream tasks, with an ensemble of 8 temporal crops at test time. Pretraining used ~39,000 UKB resting-state scans; evaluation spanned demographic, cognitive, mood, and clinical diagnosis tasks across five external datasets.

Why This Matters

The paper challenges the dominant reconstruction-centric paradigm in neuroimaging foundation models, arguing that abstraction — not signal fidelity — is what enables transfer to stable phenotypes. This has practical consequences: models become more robust to scanner differences and cohort shifts, and usable with a simple linear classifier rather than task-specific fine-tuning.

Real-world applications:

• Clinical screening: Predicting disorders such as schizophrenia, depression, or autism from resting-state fMRI without site-specific retraining.
• Biomarker discovery: Identifying which functional networks carry predictive signal for a given condition, aided by the model's in-distribution interpretability.
• Cognitive and developmental phenotyping: Predicting language and cognitive scores from brain scans to support research on development and aging.
• Cross-cohort research: Enabling analyses across datasets with heterogeneous acquisition protocols and preprocessing, reducing the need for harmonic or domain adaptation.

Industry relevance: Brain-imaging device makers, neurotech companies, and clinical trial platforms (especially in psychiatry and neurology) could use this style of model as a general-purpose encoder for downstream diagnostic or drug-response prediction, lowering the data burden for new task-specific classifiers.

Future Directions

• Extending the self-distillation framework to other modalities (e.g., EEG, MEG, task-based fMRI with more diverse paradigms).
• Scaling pretraining further to test whether performance continues to grow log-linearly, or plateaus, and whether even larger multi-site corpora accelerate OOD transfer.
• Improving temporal resolution for fast task dynamics, since the current approach downsamples to 2 s and is designed for slow macro-scale dynamics.
• Translating the interpretability probes into validated clinical biomarkers, particularly for the cerebellar signals identified in MDD.
• Investigating whether the semantic tokenizer can be adapted or learned end-to-end to discover networks rather than relying on predefined atlases.

Target Audience

Researchers and graduate students in computational neuroscience, neuroimaging, and machine learning who are interested in self-supervised learning, foundation models, or biomarker discovery from fMRI. It is also relevant to clinical researchers evaluating machine-learning tools for psychiatric and cognitive phenotyping, and to engineers at health-tech or neurotech companies seeking transferable representations of brain dynamics.

Authors’ abstract

The development of foundation models for functional magnetic resonance imaging (fMRI) time series holds significant promise for predicting phenotypes related to disease and cognition. Current models, however, are often trained using a mask-and-reconstruct objective on small brain regions. This focus on low-level information leads to representations that are sensitive to noise and temporal fluctuations, necessitating extensive fine-tuning for downstream tasks. We introduce Brain-Semantoks, a self-supervised framework designed specifically to learn abstract representations of brain dynamics. Its architecture is built on two core innovations: a semantic tokenizer that aggregates noisy regional signals into robust tokens representing functional networks, and a self-distillation objective that enforces representational stability across time. We show that this objective is stabilized through a novel training curriculum, ensuring the model robustly learns meaningful features from low signal-to-noise time series. We demonstrate that learned representations enable strong performance on a variety of downstream tasks even when only using a linear probe. Furthermore, we provide comprehensive scaling analyses indicating more unlabeled data reliably results in out-of-distribution performance gains without domain adaptation.