brain.spaceSTST-JEPA

STST-JEPA

Shallow-Target Spatio-Temporal Joint Embedding Prediction Architecture for EEG self-supervised learning

Self-supervised EEG foundation model. Pretrained on 47,703 sessions (ages 5 to 81). Brain age regression and multi task transfer from one frozen backbone.

Roy Segal  ·  Yoni Svechinsky  ·  Tomer Fekete

brain.space  ·  2026

Paper (arXiv) Abstract Method brain.space Contact
3.06 yr
Brain age MAE (r = 0.924)
47,703
Pretraining sessions, ages 5 to 81
SOTA × 3
Downstream tasks: sex, age, psychopathology

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 nonstationarity, 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 multichannel 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 fine tuning 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 × 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.

Mask, predict, reconstruct

STST-JEPA is a self-supervised transformer trained on 30 second, 256 Hz, 128 channel EEG windows. Each channel is patchified by a temporal convolution, and the channel axis is collapsed at every time index by a coordinate aware pooled attention block, yielding one token per time index and admitting arbitrary montages (115 and 128 channel layouts, missing channels masked). Roughly 24% of the space by time grid is masked with rectangular blocks. Visible tokens pass through a 24 layer, width 768 context encoder and a cross-attention predictor. Prediction targets are the outputs of an exponential moving average copy of the shallow tokenizer (stop-gradient); a per patch head reconstructs each masked patch to its 16 sample waveform. Loss = latent prediction (MSE, weight 1.0) + reconstruction (smooth-L1, weight 0.35). Age is decoded by a lightweight attentive probe on frozen embeddings.

Input window 128 ch · 30 s · 256 Hz Tokenizer patch + PMA pool coordinate aware Block mask ~24% space×time Context encoder 24 layers · width 768 visible tokens only Predictor fills masked positions EMA tokenizer stop-gradient target full window Recon head back to 16 samples latent prediction loss (λ = 1.0) predicted masked tokens vs EMA tokenizer targets reconstruction floor (λ = 0.35)
Latent prediction against the shallow EMA target Per patch waveform reconstruction

Predict age from a frozen backbone

A frozen backbone with an attentive probe reaches MAE 3.06 yr, Pearson r 0.924, on 3,367 held out validation sessions - a 69.7% reduction versus predicting the training mean.

Table 1 - Age prediction on the held out validation set (combined brain.space + HBN, N = 3,367 sessions).
MethodMAE (yr)RMSE (yr)R²
STST-JEPA + attentive probe (brain.space + HBN)3.065.110.85
Predict training mean (17.7 yr)≈ 10.0913.27≈ 0
Scatter of predicted versus true age on 3,367 validation sessions. The red OLS fit tracks the identity line at young ages and falls below it at older ages. Inset reports MAE 3.061 and correlation 0.924.
Figure 1. Predicted vs true age, validation set (N = 3,367); MAE 3.06 yr, r = 0.924.

SOTA on three NeuralBench downstream tasks: Age, Sex, Psychopathology

The same pretrained backbone, with final layer fine tuning, on the public NeuralBench × brain.space EEG leaderboard (fixed protocol test partition). Rank 1 on all three tasks with 30 second windows.

Table 2 - NeuralBench × brain.space leaderboard, single shared model (30 second windows).
TaskMetricOurs (30 s)Prior best (model)MarginRank
Sexbal. acc.0.9110.910 (ShallowFBCSPNet)+0.0011 / 18
AgePearson r0.7490.721 (REVE, Ouahidi et al. 2025)+0.0281 / 17
PsychopathologyPearson r0.2150.137 (CBraMod)+0.0781 / 15

Brain age gap correlates with reduced performance

The bias corrected brain age gap (predicted minus true age) correlated against 21 behavioural capacity targets on the union of validation and test partitions (N = 8,600 sessions).

Forest plot of Pearson r between the brain age gap and 21 behavioural capacity targets with 95% confidence intervals. Seven rows survive Benjamini-Hochberg FDR at q equals 0.05, all negative.
Figure 3. Brain age gap vs 21 behavioural targets; 7 pass FDR at q = 0.05, all negative.

Cite this work

@misc{segal2026ststjepashallowtargetspatiotemporaljoint,
  title={STST-JEPA: Shallow-Target Spatio-Temporal Joint Embedding Prediction Architecture For EEG Self-Supervised Learning},
  author={Roy Segal and Yoni Svechinsky and Tomer Fekete},
  year={2026},
  eprint={2607.06629},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2607.06629},
}
Contacttomer.fekete@brain.space
PaperarXiv:2607.06629