| name | stst-jepa-eeg-foundation |
| description | STST-JEPA: Shallow-Target Spatio-Temporal Joint Embedding Predictive Architecture for EEG self-supervised learning. Largest EEG foundation model (47,703 sessions, ages 5-81) using JEPA-style latent prediction with EMA tokenizer + auxiliary signal reconstruction. Rank 1 on NeuralBench for sex, age, psychopathology. Brain age gap correlates with cognitive efficiency. Activation: stst-jepa, eeg foundation model, brain age, self-supervised eeg, eeg self-supervised learning, JEPA eeg, EEG2Rep, brain space, neuralbench, brain age gap, cognitive efficiency |
| tags | ["eeg","self-supervised-learning","foundation-model","brain-age","transformer","JEPA"] |
STST-JEPA — Shallow-Target Spatio-Temporal JEPA for EEG
Paper: STST-JEPA: Shallow-Target Spatio-Temporal Joint Embedding Prediction Architecture For EEG Self-Supervised Learning
arXiv: 2607.06629v2
Authors: Roy Segal, Yoni Svechinsky, Tomer Fekete (brain.space)
Published: 2026-07-07 (updated 2026-07-09)
Categories: cs.LG, q-bio.NC
Core Contribution
Largest EEG foundation model pretraining to date: 47,703 sessions spanning ages 5-81 from brain.space and Healthy Brain Network (HBN) corpora. Uses a JEPA-style joint embedding predictive architecture with shallow (tokenizer-level) EMA targets plus auxiliary signal reconstruction, achieving rank-1 on NeuralBench across 3 downstream tasks from a single shared backbone.
Key Innovations
1. Joint Latent-Prediction + Reconstruction Objective
- Latent prediction (λ=1.0): Predict masked token representations against EMA-of-tokenizer targets (MSE)
- Auxiliary reconstruction (λ=0.35): Smooth-L1 loss reconstructing each masked patch back to raw 16-sample waveform
- Maps loosely to predictive processing accounts of cortex (Rao & Ballard 1999; Friston 2010): latent generative model disciplined by signal-level consequences
- Reconstruction deliberately down-weighted — acts as soft floor preventing representation collapse, not co-equal supervision
2. PMA (Pooled Multihead Attention) Channel Pooling
- Collapses C=128 channels into single token per temporal index via Set Transformer PMA
- 16 learned inducing queries attend over channel dimension via 16-head cross-attention
- Coordinate-aware: each patch embedding combined with learned projection of 3D electrode coordinates
- Mask-aware: invalid/missing channels suppressed — handles montage heterogeneity (115-channel brain.space vs 128-channel HBN)
- Gated coordinate learning: gate initialized to zero, trained at 5× LR
3. Shallow-Target Architecture
- EMA target stream is tokenizer-level (not encoder-level) — unlike canonical JEPA
- Context encoder: 24-layer transformer (d=768, 16 heads, FFN=3072) — no EMA twin
- Predictor: 2-layer transformer with cross-attention, sinusoidal position queries at masked positions
- Closest published analogue: multilayered MAE with EMA-targeted tokens
4. Spatiotemporal Block Masking
- 4 rectangular blocks sampled per example
- Each block: random channels [0.2, 0.5] × random time patches [0.1, 0.3]
- Average mask ratio: ~24% of grid positions
5. Lightweight Attentive Probe Protocol
- Self-attention (CLS token) + single query cross-attention + linear head
- Single-query design: O(d_p) params vs O(T · d_p) for standard attention
- Length-agnostic — reusable across window lengths
Architecture
Input: 30s × 256Hz × 128 channels = 7,680 samples/channel
↓
Tokenization: 1D temporal conv (P=16, stride=16) → 480 temporal patches × 128 channels
↓
PMA Channel Pooling: 128 channel patches → 1 token (d=768) per temporal index
↓
Token Sequence: T=480 tokens, d=768 (with RoPE positional encoding)
↓
[SPATIOTEMPORAL BLOCK MASKING — ~24% masked]
↓
Context Encoder (24-layer transformer, visible tokens only)
↓
Predictor (2-layer cross-attention, predicts at masked positions)
↓
┌─────────────────────┬──────────────────────┐
│ L_lat (MSE) │ L_rec (Smooth-L1) │
│ vs EMA tokenizer │ vs raw 16-sample │
│ target tokens │ patch waveform │
│ λ=1.0 │ λ=0.35 │
└─────────────────────┴──────────────────────┘
Hyperparameters
| Component | Parameter | Value |
|---|
| Input | Window | 30s @ 256Hz (7,680 samples/ch) |
| Input | Channels | 128 (unified budget, mask handles missing) |
| Tokenizer | Patch length | 16 samples (480 patches/ch) |
| PMA | Inducing queries | 16 |
| Encoder | Layers / width / heads | 24 / 768 / 16 |
| Encoder | FFN dim | 3,072 |
| Encoder | Dropout | 0.1 |
| Positional | Encoding | RoPE (base 10,000) |
| Predictor | Layers | 2 (cross-attention) |
| EMA | Momentum | 0.9996 → 0.9999 (warmup 12,000 steps) |
| Masking | Blocks per example | 4 rectangular |
| Loss | λ_lat / λ_rec | 1.0 / 0.35 |
| Optimizer | AdamW peak LR | 5e-5 |
| Optimizer | Weight decay | 0.05 |
| Training | Devices | 7 GPUs (DDP, 16-mixed) |
| Training | Global batch | 245 windows (35/device) |
| Training | Epochs | 15 (~155,500 steps) |
EEG Preprocessing Pipeline
- Upstream conditioning (done before SSL pipeline):
- High-pass 1Hz + transient spike interpolation
- Notch filter at line frequency + harmonics
- Bad channel removal (spectral quality model or RMS threshold)
- ASR (κ=20 robust SD threshold)
- ICA (Picard, extended) — remove ocular + cardiac components, protect alpha
- Resample to 256Hz via polyphase resampling
- Window extraction: 30s with 6s stride (train) / 2s stride (val)
- Robust normalization: channel-wise median/IQR z-scoring
- Validity mask: missing channels + nonfinite stats + amplitude > 150
Results
Age Regression (held-out validation, N=3,367 sessions)
| Method | MAE (yr) | RMSE (yr) | R² |
|---|
| STST-JEPA + probe (combined) | 3.06 | 5.11 | 0.85 |
| STST-JEPA + probe (brain.space only) | 4.82 | 7.06 | 0.654 |
| Predict training mean (17.7yr) | ~10.09 | 13.27 | ~0 |
NeuralBench Leaderboard (single shared backbone)
| Task | Metric | Ours (30s) | Prior Best | Rank |
|---|
| Sex | bal. acc. | 0.911 | 0.910 (ShallowFBCSPNet) | 1/18 |
| Age | Pearson r | 0.749 | 0.721 (REVE) | 1/17 |
| Psychopathology | Pearson r | 0.215 | 0.137 (CBraMod) | 1/15 |
Brain Age Gap × Behavioural Capacity
- 7 of 21 targets survive BH-FDR correction (q=0.05)
- All significant targets show negative correlation: older-looking brain → worse cognitive efficiency
- Strongest: Stroop efficiency (r=-0.089), n-back (r=-0.057)
- Effects small (|r|<0.10) but directionally consistent
Loss Dynamics
- L_lat: drops ~4 orders of magnitude (0.0575 → 1.93e-4, 99.7% reduction)
- L_rec: modest decrease (0.317 → 0.169, -47%), acts as soft floor
- Late training plateau: L_lat ∈ [1.5e-2, 2.2e-2], L_rec ∈ [0.13, 0.18]
Brain-Inspired Framing
Joint objective maps to predictive processing theory:
- Latent prediction ≈ cortex's generative model predicting hidden causes
- Reconstruction ≈ sensory prediction error disciplining the model
- Not a claim about cortical fidelity — architectural intuition only
Limitations & Caveats
- No controlled ablation of joint objective (beyond compute budget)
- No fixed-protocol test on internal brain.space partition
- No clinical biomarker claim
- Pediatric-heavy corpus (HBN) — MAE not directly comparable to adult benchmarks
- 30s windows give advantage over 2s standard NeuralBench protocol
- Mixed resting/task paradigms in pretraining — BAG may carry task state information
Projected Improvements (from paper)
- 3× corpus scaling (~143K sessions): projected MAE 2.6-2.9yr
- Auxiliary losses: age regression head (λ≈0.05), cohort token, contrastive-by-age-band
- Fine-tuning: end-to-end fine-tuning projected 0.2-0.5yr MAE reduction; LoRA adapters (rank 8-16)
Related Work
- JEPA family: I-JEPA (Assran et al. 2023), V-JEPA (Bardes et al. 2024), data2vec (Baevski et al. 2022)
- EEG foundation models: BENDR, BrainBERT, BIOT, LaBraM, Neuro-GPT, EEG2Rep, REVE
- Self-supervised SSL: BYOL, MAE, VICReg, LeJEPA
Related Skills
- [[reve-eeg-foundation]] - REVE EEG foundation model
- [[laya-eeg-foundation]] - Laya LeJEPA approach to EEG
- [[eeg-fm-audit-systematic-evaluation]] - EEG foundation model evaluation framework
- [[identity-trap-eeg-foundation-models]] - EEG foundation model identity trap
