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

TL;DR

Brain-Semantoks introduces a self-supervised framework that learns robust, high-level brain dynamic representations from noisy fMRI data, improving downstream task performance and out-of-distribution generalization.

Contribution

The paper presents a novel semantic tokenizer and self-distillation training method tailored for brain dynamics, enabling more stable and meaningful representations from noisy fMRI signals.

Findings

Effective in downstream tasks with linear probes

Outperforms existing models in out-of-distribution scenarios

Scaling with more unlabeled data improves performance

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.