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

2026-07-09 · Source: cs.LG updates on arXiv.org · Field: Science & Research — Artificial Intelligence & Machine Learning, Health & Medical Research, Life Sciences & Biology · Depth: Expert, extended

Summary

STST-JEPA, a self-supervised transformer, is introduced for resting-state and task EEG analysis, pretrained on 47,703 sessions from the brain.space and Healthy Brain Network (HBN) corpora, spanning ages 5–81. This 24-layer model combines a latent-prediction objective, which predicts 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, achieved a held-out-validation mean absolute error of 3.06 years (r = 0.924) for age regression on 3,367 sessions, representing a 69.7% reduction from a 10.09 years baseline. Furthermore, with minimal task-specific fine-tuning, the pretrained encoder secured rank-1 positions 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) using its native 30-second windows. The model's age-prediction residual also showed a negative correlation with cognitive efficiency across examined tasks.

Key takeaway

For machine learning engineers developing EEG-based biomarkers, you should consider STST-JEPA's joint self-supervised learning approach. Its ability to achieve a 3.06-year MAE for age prediction and rank-1 performance on multiple NeuralBench tasks, even with heterogeneous data, suggests a robust foundation model. You can adapt this architecture to build generalizable EEG representations for diverse clinical or research applications.

Key insights

STST-JEPA effectively predicts brain age and other traits from EEG using a novel self-supervised learning approach.

Principles

• Joint latent-prediction and signal-reconstruction improves EEG representation.
• Coordinate-aware channel pooling handles montage heterogeneity.
• Brain-age gap correlates with cognitive efficiency.

Method

STST-JEPA pretrains a 24-layer transformer on 30-second EEG windows using a joint objective: masked-token latent prediction (λ_lat=1.0) and smooth-L1 signal reconstruction (λ_rec=0.35) with spatiotemporal block masks.

In practice

• Use 30-second EEG windows for robust age prediction.
• Employ joint latent-prediction and reconstruction for self-supervised EEG.
• Apply attentive probes on frozen embeddings for downstream tasks.

Topics

• EEG Self-Supervised Learning
• Brain Age Prediction
• Transformer Architecture
• Joint Embedding Predictive Architecture
• NeuralBench Benchmark
• Cognitive Biomarkers

Best for: AI Scientist, Machine Learning Engineer, Research Scientist

Related on AIssential

• STST-JEPA: Shallow-Target Spatio-Temporal Joint Embedding Prediction Architecture For EEG Self-Supervised Learning — STST-JEPA uses self-supervised learning on EEG data to create a robust foundation model for brain age and other neurological biomarkers.
• How AI Learned to Teach Itself [JEPA] — Self-supervised learning evolves from contrastive methods to JEPA, predicting latent embeddings for robust world models.
• IQ-JEPA: A Joint-Embedding Predictive Architecture with a Hermitian Vision Transformer for Sound Speed and Attenuation Estimation from Ultrasound IQ Data — IQ-JEPA uses a Hermitian Vision Transformer and self-supervision to efficiently estimate ultrasound sound speed and attenuation from IQ data.
• Contrastive Joint-Embedding Prediction for Representation Learning in Structural MRI — COJEPA combines JEPA and contrastive learning for robust, self-supervised 3D MRI representation learning.
• MSBraM: A Multi-scale Self-supervised Brain Foundation Model for Hierarchical EEG Dynamics Learning — Effective EEG foundation models critically depend on explicitly modeling multi-scale temporal dynamics for hierarchical representation learning.
• LatentWave: JEPA Pretraining for Wireless Foundation Models — JEPA pretraining for wireless foundation models enhances representation transferability by predicting masked regions in latent space.

See Counsel's argued verdicts on the open AI decisions leaders are weighing →

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by cs.LG updates on arXiv.org.