DEFT: Disentanglement-Enhanced Fine-Tuning for EEG Foundation Models

AI summary

Overview

Research area: EEG foundation models (EFMs), representation learning, and cross-subject transfer learning for brain-computer interfaces and clinical EEG decoding.

Technical level: Advanced. The core idea is intuitive, but the method relies on precision-weighted Gaussian state updates adapted from Xi-vector and Kalman filtering, plus several jointly optimized losses.

One-sentence scope: The paper proposes DEFT, a plug-and-play adapter that splits a pretrained EEG foundation model's mixed token representation into a task-content branch and a subject-style branch before classification, and evaluates it on six downstream datasets across four EFM backbones.

What This Paper Is About

EEG foundation models are pretrained on large multi-subject corpora, and their embeddings turn out to encode subject identity very strongly — a linear probe recovers who a person is more easily from pretrained CBraMod embeddings than from raw EEG, while class-label probing improves only modestly. In cross-subject decoding, that subject structure is a nuisance the downstream classifier must first suppress before it can learn the actual task. DEFT addresses this by inserting a disentanglement step between the frozen-style pretrained backbone and the prediction head, so that subject-consistent structure is modeled explicitly in its own branch and prediction is made only from the task-content branch.

Key Contributions

1. Diagnosis of a specific failure mode. The authors show, via probing, that pretrained EFM embeddings make subject identity more linearly accessible than raw EEG and form subject-dependent clusters, arguing that "the strongest signal in a pretrained token may be the wrong signal for the downstream task."
2. A factorize-before-classify adapter (DEFT). A post-backbone module that maps mixed pretrained tokens into a content branch and a style branch, applies the classification head only to the content branch, and reinjects window-level factors back into the original token grid via residual connections.
3. A recursive Gaussian disentangler. An iterative two-branch inference procedure that adapts precision-weighted linear-Gaussian aggregation (Xi-vector) and Kalman-style prior–observation fusion to EEG tokens, using four alternating static/dynamic update steps with learned per-dimension precision and a learned counterpart-removal gate.
4. A combined intrinsic and external supervision scheme. Temporal-variance ordering plus style-stability losses shape the branches from within, while subject-aware contrastive learning on the style branch and a FiLM-conditioned reconstruction loss anchor them externally — with subject identities needed only during training, not at inference.

Main Findings

• Consistent average gain across backbones and datasets. Over the 24 matched backbone–dataset cells in the main table, DEFT improves mean Balanced Accuracy by 0.0460, with a two-way backbone/dataset clustered-bootstrap 95% confidence interval of [0.0240, 0.0745].
• Backbone-level average B-ACC gains differ in magnitude: +0.0628 for CBraMod, +0.0574 for CSBrain, +0.0425 for CodeBrain, and +0.0211 for LaBraM (averaged over the six representative datasets).
• Largest task families for improvement. The biggest mean B-ACC gains occur on MentalArith., BCIC-IV-2a, and FACED; gains also appear on sleep staging, speech decoding, and MDD classification.
• Some headline cell-level deltas are large. CodeBrain on BCIC-IV-2a moves B-ACC from 0.3678 to 0.4621 (+0.0943); CSBrain on FACED moves from 0.4169 to 0.5748 (+0.1579). The paper also retains every observed decrease rather than only reporting wins — for example CBraMod on ISRUC-S1 drops in Kappa (−0.0229) and F1 (−0.0145), and LaBraM on MentalArith. drops in PR-AUC (−0.0860).
• Component sensitivity is task dependent. Ablations on CBraMod with fixed seeds 0–2 show the largest reductions on BCIC-IV-2a, MentalArith., BCIC-Speech, and Mumtaz2016, while changes on FACED and ISRUC-S1 are smaller. Removing the invariance component on BCIC-IV-2a changes B-ACC by −0.1606 (Welch 95% CI [−0.2999, −0.0213]); removing the variance-ratio term drops that dataset's B-ACC to 0.3808 and removing the Gaussian component to 0.3782, versus 0.5579 for the full model.
• The content branch improves class geometry after subject centering. Using a within-subject class separation score after subject centering and PCA50, the content branch improves on five of six datasets: BCIC-IV-2a (0.490 → 0.637), FACED (0.520 → 0.783), ISRUC-S1 (1.679 → 2.556), MentalArith. (0.580 → 1.115), and Mumtaz2016 (2.291 → 6.352).
• One geometry mismatch. BCIC-Speech is the only case where the post-hoc score falls (0.405 → 0.307), despite metric gains for all four backbones; the authors suggest speech cues may be localized across channels and windows rather than visible as global PCA clusters.
• Branch specialization confirmed by frozen linear probes. Content yields higher task-probe B-ACC and style yields higher subject-probe B-ACC in all six evaluated settings.

Methodology in Plain English

The starting point is a pretrained EEG backbone that turns a trial into a grid of tokens. DEFT compresses each time window of that grid into a single summary vector, then hands the sequence of summaries to a disentangler that produces two parallel sequences: a content version and a style version.

The disentangler works by alternating estimation. Each branch peeks at the other branch's current guess, subtracts a learned fraction of it from the observation, and then updates its own estimate using a Gaussian fusion rule where learned per-dimension precision values decide how much to trust the new observation versus the previous estimate. The style branch uses a static transition (identity), the content branch uses a learned mixture of transition bases selected by a small gating network. Four alternating steps are run. The resulting window-level factors are then re-expanded into the full token grid by adding back the per-channel-group residual, so channel information survives the compression.

Supervision is layered on top. A hinge loss requires the content branch to vary more across windows than the style branch, with a margin, and a separate penalty pulls style variation down — encoding the prior that task evidence changes over time while subject identity stays stable within a trial. A subject-aware contrastive loss pulls style summaries from the same subject together and pushes different subjects apart, using a subject-balanced sampler that guarantees at least one same-subject positive per anchor on the six main datasets. A reconstruction loss rebuilds the original mixed window representations from the content sequence, modulated by FiLM parameters computed from a pooled style summary, so the split cannot throw away information. Cross-entropy on the content branch drives the task. Subject labels are used only for the contrastive term during training.

Evaluation keeps everything matched: public checkpoints for CodeBrain, CBraMod, LaBraM, and CSBrain are not re-pretrained, the backbone and head stay trainable, and only the adapter differs between the vanilla and DEFT conditions. All experiments ran on NVIDIA RTX 4090 GPUs, with BCIC-IV-2a, FACED, ISRUC-S1, MentalArith., BCIC-Speech, and Mumtaz2016 under subject-disjoint splits; chance B-ACC is 0.250, 0.111, 0.200, 0.500, 0.333, and 0.500 respectively.

Why This Matters

Impact on research. The paper reframes a common assumption in EEG foundation-model work: that better pretrained embeddings automatically mean better downstream decoding. Its probing result — that subject identity becomes more linearly accessible after EFM encoding while class-label probing improves only modestly — suggests scaling pretraining can amplify the wrong signal for cross-subject tasks. Placing disentanglement at the adaptation stage, rather than inside pretraining, gives the field a cheap, backbone-agnostic intervention that works with released checkpoints.

Real-world applications:

• Assistive and rehabilitative BCIs using motor imagery, where the decoder must generalize to a new user whose neural signature it has never seen.
• Sleep staging in clinical or home monitoring, where per-patient variability in electrode placement and anatomy degrades automated scoring.
• Mental workload assessment in safety-critical settings such as aviation or driving, where alertness estimates must transfer across operators.
• Clinical screening, including MDD classification and speech decoding, where subject-level confounds are a persistent obstacle to deployment.

Industry relevance. Because DEFT is a post-backbone adapter with only a small parameter overhead relative to the original fine-tuning model, it can be layered onto existing EFM deployments without retraining the foundation model. The fact that subject labels are unnecessary at inference matters commercially: enrollment data is often unavailable for a new end user.

Future Directions

• Extending to decoder-based and language-connected EFMs. The authors position DEFT for encoder-based EFMs that expose token embeddings, and explicitly note that models such as NeuroLM and ECHO form a complementary direction the adapter has not been applied to.
• Handling task dependence in the losses. Ablations show component sensitivity varies substantially by dataset, and the appendix objective-combination control shows task-dependent interactions among the variance, contrastive, and reconstruction terms. How to set or learn the four loss weights and the variance margin per task is left open.
• Understanding the BCIC-Speech geometry mismatch. The only dataset where the class-separation score falls despite metric gains suggests the current global-PCA diagnostic misses localized, channel- and window-specific structure; better representation diagnostics and possibly channel-aware objectives are natural next steps.
• Broadening the evidence base. The conclusion notes appendix checks on one additional EFM (REVE), one independent MDD dataset, and one MOMENT/UCI HAR setting — a signal that larger multi-corpus validation, and extension beyond EEG to other biosignals, remains future work.

Target Audience

Researchers and engineers working on EEG foundation models, cross-subject transfer learning, and brain-computer interfaces who are already comfortable with token-based architectures and want a drop-in adaptation method. It also suits representation-learning researchers interested in disentanglement applied to non-stationary physiological time series, and practitioners deploying EFMs who need robustness to unseen subjects without retraining the backbone. Readers without a background in probabilistic state-space filtering will find the disentangler's update equations the steepest part; the motivation and empirical results are accessible to a broader machine-learning audience.

Authors’ abstract

EEG foundation models (EFMs) have improved EEG representation learning by transferring knowledge from large-scale multi-subject pretraining to downstream tasks. However, their pretrained embeddings often retain substantial subject-related variation, placing an additional burden on the downstream classifier, especially in cross-subject settings. To address this issue, we propose Disentangled EFM Fine-Tuning (DEFT), a plug-and-play adapter for EFM fine-tuning. DEFT encourages branch disentanglement of pretrained tokens into task-oriented content and subject-consistent style, then performs downstream prediction from the content branch. It instantiates this factorize-before-classify design with recursive Gaussian refinement, temporal-variance regularization, subject-aware contrastive learning, and reconstruction. Experiments across multiple EEG datasets and four EFM backbones show that DEFT improves matched fine-tuning baselines. Further ablations and analyses support the intended content/style branch organization.