STST-JEPA: Shallow-Target Spatio-Temporal Joint Embedding Prediction Architecture For EEG Self-Supervised Learning

Abstract: Brain age - the age inferred from a physiological recording - is an emerging biomarker whose deviation from chronological age tracks neurological and psychiatric burden, and EEG is an attractive substrate for it because it is cheap, portable, and temporally rich. Yet EEG brain-age models must contend with cross-site montage heterogeneity, small labelled cohorts, and dominant subject-level non-stationarity, and few EEG foundation models have been shown to deliver competitive age regression across the full pediatric to older adult range in which such a biomarker would actually be deployed. We introduce STST-JEPA, a self-supervised transformer for resting-state and task EEG, pretrained on 47,703 sessions spanning ages 5-81 from the brain.space and Healthy Brain Network (HBN) corpora. The model combines a latent-prediction objective - predicting masked-token representations against an EMA-of-tokenizer target - with an auxiliary signal-reconstruction term, applied to 30-second multi-channel windows under spatiotemporal block masks. A lightweight attentive probe trained on frozen pretrained embeddings achieves a best held-out-validation mean absolute error of 3.06 years (r = 0.924) for age regression on 3,367 sessions, against a predict-the-mean baseline of approximately 10 years MAE. With light task-specific finetuning of the model's final layers, the same pretrained encoder achieves rank-1 placements - with the model's native 30-second windows - on the public NeuralBench x brain.space EEG leaderboard for sex classification (balanced accuracy 0.911), age prediction (r = 0.749), and psychopathology composite regression (r = 0.215). We further show that the model's age-prediction residual is negatively correlated with cognitive efficiency over several tasks we examined.

• The paper introduces a self-supervised transformer architecture that uses masked latent-prediction and reconstruction to robustly learn EEG representations.
• It integrates montage geometry and channel-wise pooling to handle heterogeneous EEG layouts, achieving a near 70% reduction in age prediction MAE.
• Downstream probes demonstrate strong transfer with rank-1 performance in sex classification and psychopathology, validating its biomarker potential.

Summary of "STST-JEPA: Shallow-Target Spatio-Temporal Joint Embedding Prediction Architecture For EEG Self-Supervised Learning"

Architectural Overview and Motivation

STST-JEPA proposes a self-supervised transformer-based foundation model for EEG representation learning, targeting robust and generalizable biomarker prediction, with brain-age as a central application. The design motivation arises from the challenge of montage heterogeneity, limited labeled data, and dominant subject-level non-stationarity in EEG acquisition. The architecture combines a masked latent-prediction loss—supervising the outputs of a predictor against an EMA-of-tokenizer target—with an auxiliary patch-level signal reconstruction term, both operating under spatiotemporal block masking. The model ingests 30-second, 128-channel windows (256 Hz) and integrates montage geometry via coordinate-aware channel pooling through a PMA module.

Dataset and Pretraining Protocol

The corpus integrates 47,703 sessions from brain.space and HBN, spanning ages 5–81. Stratified subject-level partitioning eliminates leakage from repeated recordings across splits. EEG preprocessing involves robust artifact rejection, channel normalization, and coordinate-based embedding, unifying 115- and 128-channel layouts. Pretraining employs four random spatiotemporal masking blocks (approximately 24% mask ratio) and combines a mean squared latent-prediction loss ($\lambda_{lat}=1.0$) with a smooth-L1 reconstruction loss ($\lambda_{rec}=0.35$). Model parameters are optimized in mixed precision on distributed GPUs.

Tokenization, Embedding, and Encoder

Each input window is patchified into 16-sample temporal segments per channel. PMA module performs channel-wise pooling at each time index, integrating electrode coordinates for montage-awareness and handling missing channels robustly. The tokenizer creates a sequence of 480 tokens of width 768, and the context encoder is a pre-norm transformer stack (24 layers), followed by a two-layer predictor transformer with cross-attention. EMA-tokenizer targets are computed for the masked positions via exponential averaging; the loss is computed only at masked tokens.

Downstream Probes and Evaluation

A lightweight attentive probe protocol is deployed on frozen pretrained embeddings, supporting age regression, sex classification, and paradigm identification without retraining the backbone. Age prediction uses a self-attention/Cross-attention/Linear head probe, trained on session-window embeddings. Performance is measured on a held-out validation partition (N=3,367 sessions), with mean absolute error (MAE) and root mean squared error (RMSE) reported at the session level.

Numerical Results: Age Prediction and Auxiliary Tasks

Age Regression:

• Combined brain.space + HBN validation MAE: 3.06 years (RMSE: 5.11, $R^2=0.85$)
• Brain.space-only validation MAE: 4.82 years (RMSE: 7.06, $R^2=0.654$)
• Predict-the-mean baseline: $\approx 10.09$ years MAE

The model achieves nearly 70% reduction in MAE relative to the baseline. Pediatric-heavy cohort lowers conditional variance, mechanically reducing achievable MAE. The predicted vs. true age relationship exhibits classic regression-to-the-mean bias, motivating subsequent bias correction for the brain-age gap analysis.

Figure 1: Scatter plot of predicted vs. true age on the validation set; OLS fit and identity line, MAE = 3.06 years.

Figure 2: Validation age MAE vs. pretraining step showing steady decline and convergence at 3.06 years.

(Figure 1b)

Figure 1b: True age distribution in the held-out validation cohort, emphasizing pediatric skew.

Auxiliary Downstream Probes:

• Sex classification: 0.891 balanced accuracy (chance: 0.500)
• Paradigm classification (EO-EC): 0.882 balanced accuracy
• Paradigm multiclass (n-back, MOT, MMST): above chance, moderate signal

Benchmark Performance and Transfer

On the NeuralBench × brain.space public EEG leaderboard, STST-JEPA achieves rank-1 placements for:

• Sex (bal. acc. 0.911, tied best)
• Age (Pearson $r=0.749$, +0.028 over prior best)
• Psychopathology composite (Pearson $r=0.215$, +0.078 over prior best)

Sex and psychopathology correlations remain rank-1 under both 30s and 2s input windows; age correlation is rank-4 under the standard 2s protocol.

Brain-Age Gap (BAG) and Behavioral Performance

Bias-corrected BAG (predicted age minus bias-adjusted chronological age) is analyzed across 21 behavioral-efficiency targets. Seven targets survive FDR correction (q=0.05), all showing negative correlation: older-looking brain-age gap tracks poorer cognitive efficiency. Effect sizes are modest ($|r|<0.10$), but direction-of-effect consistency is high. Age-residualized general efficiency maintains negative correlation after double regression, confirming EEG signal beyond chronological age.

Figure 3: Forest plot of BAG vs. behavioral targets, with significant negative correlations surviving FDR correction.

Theoretical Implications and Future Directions

STST-JEPA demonstrates effective representation learning on heterogeneous EEG data, robust to montage discrepancies and limited labels, supporting the foundation model paradigm for neurophysiological signals. The joint latent-prediction + reconstruction objective maintains latent space signal relevancy without overfitting to surface artifacts. The reuse of frozen embeddings for downstream adaptive probes supports transfer and efficient benchmark deployment.

Planned future work includes:

• Corpus scaling (143,000 sessions projected), expecting log-axis incremental gains (MAE projected to 2.6-2.9 years).
• Auxiliary supervision (age regression head, montage token, cohort-balanced contrastive loss) to bias representation toward clinical-relevant axes.
• End-to-end and multi-task fine-tuning (subject-level distillation, LoRA-style adapters) for further reduction in MAE and augmentation of transfer.

Current limits include lack of controlled ablations for objective choice, fixed-protocol evaluation only on external partitions, and no claim to clinical biomarker validity.

Conclusion

STST-JEPA introduces a scalable, montage-aware transformer foundation model for EEG self-supervised representation learning. Its joint masked prediction and lightweight reconstruction objective yields robust embeddings across tasks, achieving substantial reduction in age regression MAE and rank-1 benchmark performance for diverse traits. Correlation of brain-age gap with behavioral efficiency validates the biomarker potential at cohort scale, with theoretical and architectural implications for future foundation model development in the neurophysiological domain. The modular probe protocol and cross-task transfer further position STST-JEPA as a versatile platform for EEG classification and regression tasks, with future improvement anticipated from corpus scaling, auxiliary losses, and fine-tuning strategies.