STST-JEPA predicts brain age from EEG with 3-year accuracy

STST-JEPA (Shallow-Target Spatio-Temporal Joint Embedding Prediction Architecture) is a new self-supervised transformer designed for EEG data, developed by Segal, Svechinsky, and Fekete. The model was pretrained on a massive dataset of 47,703 sessions from the Healthy Brain Network (HBN) and another public corpus, covering ages 5 to 81. Its training objective combines a latent-prediction task—predicting masked token representations against an exponential-moving-average of the tokenizer—with an auxiliary signal-reconstruction term, applied to 30-second multi-channel windows under spatiotemporal block masks. This approach allows the model to learn rich representations from raw EEG without requiring large labeled datasets.

The results are striking. When a lightweight attentive probe is trained on frozen pretrained embeddings, STST-JEPA achieves a held-out validation mean absolute error of just 3.06 years for age regression, compared to a baseline of roughly 10 years. On the public NeuralBench EEG leaderboard, the same encoder (with light fine-tuning) secures rank-1 placements for sex classification (balanced accuracy 0.911), age prediction (r=0.749), and psychopathology composite regression (r=0.215). Notably, the model's age-prediction residual is negatively correlated with cognitive efficiency across several tasks. This suggests the embeddings capture clinically relevant information beyond just chronological age, making STST-JEPA a promising foundation for brain-age biomarkers using low-cost, portable EEG.