STST-JEPA: Spatio-Temporal EEG Embedding

• The paper introduces a shallow-target JEPA that leverages latent predictive learning with auxiliary reconstruction, achieving a 69.7% MAE reduction in brain-age prediction.
• STST-JEPA is a self-supervised transformer that processes 30-second EEG windows with patch convolution, multi-head attention pooling, and coordinate gating to manage montage heterogeneity.
• Multi-task transfer is demonstrated by rank-1 leaderboard performance in sex classification, age regression, and psychopathology composite tasks, confirming its robustness as an EEG foundation model.

Searching arXiv for the specified paper and closely related work to ground the article in current literature. STST-JEPA, short for Shallow-Target Spatio-Temporal Joint Embedding Prediction Architecture, is a self-supervised transformer for resting-state and task EEG introduced for EEG self-supervised learning and evaluated primarily in brain-age prediction (Segal et al., 7 Jul 2026). It is designed to address three constraints emphasized in EEG biomarker modeling: cross-site montage heterogeneity, small labelled cohorts, and dominant subject-level non-stationarity. The model is pretrained on 47,703 sessions spanning ages 5–81 from the brain.space and Healthy Brain Network (HBN) corpora, and combines a latent-prediction objective with an auxiliary signal-reconstruction term under spatiotemporal block masking (Segal et al., 7 Jul 2026). In the reported downstream evaluations, a lightweight attentive probe trained on frozen pretrained embeddings achieves a best held-out-validation mean absolute error of 3.06 years with $r=0.924$ for age regression on 3,367 sessions, and light task-specific fine-tuning yields rank-1 placements on the public NeuralBench x brain.space EEG leaderboard for sex classification, age prediction, and psychopathology composite regression (Segal et al., 7 Jul 2026).

1. Architectural formulation

STST-JEPA operates on 30 s EEG windows $X\in\mathbb R^{C\times S}$ with $C=128$ channels and $S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$ timepoints (Segal et al., 7 Jul 2026). A shared 1-D convolution with kernel length $P=16$ and stride $\Delta=16$ splits each channel into $T=\lfloor(S-P)/\Delta\rfloor+1=480$ non-overlapping temporal patches of length 16 samples. At each temporal index, the 128 per-channel patch embeddings are pooled into a single $d$-dimensional token through a Pooled Multihead Attention block with $Q=16$ learned inducing queries and $H=16$ heads. The tokenizer, consisting of the patch convolution, coordinate gating, and PMA, outputs a sequence of $X\in\mathbb R^{C\times S}$0 tokens in $X\in\mathbb R^{C\times S}$1 with $X\in\mathbb R^{C\times S}$2 (Segal et al., 7 Jul 2026).

A learnable scalar gate $X\in\mathbb R^{C\times S}$3 scales and adds a per-channel embedding of the 3D electrode coordinate before PMA. This gate is initialized to zero and trained at $X\in\mathbb R^{C\times S}$4 the base learning rate so that montage geometry is incorporated gradually rather than imposed from the start (Segal et al., 7 Jul 2026). Temporal position is encoded implicitly via rotary position embeddings in all self-attention layers, with no additive positional vectors.

The context encoder is a Pre-Norm transformer of depth $X\in\mathbb R^{C\times S}$5, width $X\in\mathbb R^{C\times S}$6, with 16 attention heads and FFN inner dimension 3 072, using GeLU non-linearities and dropout 0.1 (Segal et al., 7 Jul 2026). It ingests only visible, unmasked tokens. A two-layer predictor with cross-attention then attends from masked token positions, queried by fixed sinusoidal embeddings, to the encoder’s visible key/value pairs; RoPE is disabled in the predictor, and masked positions receive only the sinusoidal queries. The predictor output at each masked position is projected into the $X\in\mathbb R^{C\times S}$7 loss space by

$X\in\mathbb R^{C\times S}$8

A distinctive aspect of the model is the target tokenizer, defined as an exponential moving average copy of the live tokenizer, with no dropout and stop-gradient, producing reference tokens $X\in\mathbb R^{C\times S}$9 at every patch (Segal et al., 7 Jul 2026). Because the predictor is trained to match tokenizer-level representations rather than deep encoder features, the architecture is described as a “shallow-target JEPA”. This suggests that the pretraining signal is intentionally anchored to comparatively low-level latent structure rather than to higher-layer contextual representations.

2. Spatio-temporal masking and self-supervised objectives

Masking is applied on a 2D grid of channel-by-time patches. Four rectangular blocks $C=128$0 are sampled independently for each example, with channel-span fraction drawn from $C=128$1 of the 128 channels and time-span fraction drawn from $C=128$2 of the 480 time patches (Segal et al., 7 Jul 2026). Because blocks may overlap, the realized mask ratio concentrates near 24% of the approximately 61,440 patch positions. Masked positions are removed from the encoder input and predicted later by the predictor. They are also the only positions at which raw signal reconstruction is applied.

The latent-prediction objective is mean squared error between the predictor output $C=128$3 and the EMA-tokenizer target $C=128$4, with stop-gradient on the target branch:

$C=128$5

This is the JEPA component of the training criterion (Segal et al., 7 Jul 2026).

An auxiliary reconstruction pathway maps each predicted token $C=128$6 back to a raw temporal patch of length 16 using a lightweight decoder $C=128$7, instantiated as LN→Linear(768→1 536)→GeLU→Linear(1 536→16) (Segal et al., 7 Jul 2026). A Huber loss with $C=128$8 compares the reconstructed patch to the true preprocessed patch:

$C=128$9

Reconstruction is applied only at masked and valid, non-artifact patches.

The full pretraining objective is

$S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$0

with $S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$1 and $S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$2 (Segal et al., 7 Jul 2026). The stated interpretation is that the stronger emphasis on $S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$3 enforces a compact latent geometry, while the reconstruction term acts as a soft regularizer that keeps the latent embeddings grounded in waveform statistics. A plausible implication is that STST-JEPA is positioned between purely latent predictive learning and direct waveform reconstruction, rather than reducing pretraining to either extreme.

3. Cross-montage handling and preprocessing regime

The pretraining corpus combines 47,703 sessions spanning ages 5–81 years: 22,588 sessions from brain.space, with a montage of 115 channels, and 25,115 sessions from HBN, with up to 128 channels (Segal et al., 7 Jul 2026). Subject-level splits are stratified by age bins, session count, and a source-specific covariate—handedness for brain.space and sex for HBN—yielding 32,246 training sessions, 5,802 validation sessions, and 9,655 held-out test sessions. Downstream age-regression evaluation uses a separate brain.space finetune cohort of 9,599 sessions, ages 18–81.

Cross-dataset heterogeneity is addressed through a unified 128-channel budget and mask-aware PMA with coordinate gating (Segal et al., 7 Jul 2026). Invalid or missing channels are zero-padded but suppressed inside the PMA pooling by a validity mask, so that montages of different sizes map to the same 128-channel representation. The paper characterizes this as a mechanism for absorbing heterogeneous montages without retraining. This suggests that the architecture treats channel absence as structured missingness rather than forcing ad hoc montage-specific parameterization.

The preprocessing pipeline consists of a high-pass filter at 1 Hz, notch filtering at line-frequency harmonics, bad-channel removal using source-specific rules, Artifact Subspace Reconstruction with $S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$4, and extended Picard ICA with up to 42 components removing ocular and cardiac artifacts only (Segal et al., 7 Jul 2026). Signals are resampled to 256 Hz and robustly normalized channel-wise by median and interquartile range. Samples with $S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$5 are amplitude-rejected with a 1-sample buffer, and windows with more than 45% invalid samples are resampled. No further re-referencing or band-pass is applied in the SSL pipeline.

Optimization uses AdamW with peak learning rate $S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$6, preceded by a 3,000-step linear warmup from $S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$7, weight decay 0.05, and global-norm clipping at 1.0 (Segal et al., 7 Jul 2026). EMA momentum warms linearly from 0.9996 to 0.9999 over 12,000 steps and is then held constant. Training uses a batch size of 35 windows per GPU across 7 GPUs, corresponding to 245 windows per step, for 15 epochs, approximately 155,500 steps, in mixed precision with torch.compile and DDP. The random seed is 42 for data splits, masking, and initialization.

4. Age regression protocol and performance

The principal downstream evaluation is age regression using a frozen-backbone attentive probe (Segal et al., 7 Jul 2026). The probe prepends a learned CLS token to the 480-token embedding sequence, applies one self-attention block over all tokens, then a single-query cross-attention step in which CLS attends to the 480 tokens, followed by a linear regression head LN→Dropout(0.1)→Linear(768→1). This design keeps the pretrained encoder frozen while introducing a lightweight task-specific head.

Training for the probe caps input at at most 5 windows per session; HBN is also capped at 5 to balance sources (Segal et al., 7 Jul 2026). An internal train/validation split of 80%/20% over sessions is used with early-stop patience 3. Optimization uses AdamW with learning rate $S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$8 for up to 15 epochs. Session-level prediction is computed by averaging window-level outputs $S=30\,\mathrm s\times256\,\mathrm{Hz}=7{,}680$9 over the windows in each session.

The held-out validation results on 3,367 sessions are reported as MAE $P=16$0 years, RMSE $P=16$1 years, and Pearson $P=16$2, implying $P=16$3 (Segal et al., 7 Jul 2026). The baseline is a predict-the-mean estimator at approximately 10.09 years MAE. The contribution summary characterizes this as a 69.7% MAE reduction across ages 5–81. Since the age-range claim and the reduction figure are both explicitly stated, the result is presented as evidence that the learned embeddings retain age-discriminative structure even when only a lightweight frozen probe is trained.

A common misconception in EEG representation learning is that strong downstream regression necessarily requires end-to-end task-specific tuning. The reported protocol does not support that general claim for this setting: the strongest detailed age-regression result is obtained with frozen pretrained embeddings and an attentive probe rather than full-model retraining (Segal et al., 7 Jul 2026). At the same time, this should not be conflated with a claim that fine-tuning is unnecessary in general, because a separate evaluation uses light fine-tuning on leaderboard tasks.

5. Transfer on public benchmark tasks

STST-JEPA is also evaluated on the public NeuralBench x brain.space EEG leaderboard using light fine-tuning of only the upper layers of the pretrained model (Segal et al., 7 Jul 2026). The test partitions are the leaderboard’s held-out splits, while training uses its train splits. The model is evaluated with its native 30-second windows in the summary description, while the specific leaderboard task descriptions report 2 s windows.

The reported results place the model at rank 1 on three tasks:

Task	Reported metric	Rank
Sex classification (binary, 2 s windows)	balanced accuracy $P=16$4	1/18
Age regression (2 s windows)	Pearson $P=16$5	1/17
Psychopathology composite (2 s windows)	Pearson $P=16$6	1/15

For context, the paper also reports the prior best entries as 0.910 balanced accuracy for sex classification, $P=16$7 for age regression, and $P=16$8 for psychopathology composite regression (Segal et al., 7 Jul 2026). The contribution summary interprets this as multi-task transfer from a single backbone, and further describes the outcome as supporting the model’s status as an EEG foundation model. A plausible implication is that the pretrained representation is not narrowly specialized to age regression, even though age prediction is the primary detailed downstream analysis.

These leaderboard results should be distinguished from the frozen-probe age-regression experiment. The former involves light task-specific fine-tuning of the upper layers; the latter keeps the pretrained backbone frozen and alters only the probe. The two evaluations therefore test different transfer regimes rather than providing interchangeable estimates of the same capability.

6. Brain-age gap and cognitive efficiency

The paper further examines the behavioural validity of the model’s age-prediction residual through brain-age gap analysis (Segal et al., 7 Jul 2026). Bias-corrected BAG is computed using 5-fold cross-fit OLS,

$P=16$9

with $\Delta=16$0 and $\Delta=16$1. This residual is then correlated against 21 behavioural targets, comprising 14 subject-level composites and 7 per-recording task efficiencies, over 8,600 sessions from 2,109 subjects.

Seven targets survive BH-FDR $\Delta=16$2, and all are in the expected negative direction: older-looking brain corresponds to worse speed–accuracy efficiency (Segal et al., 7 Jul 2026). Two examples are reported explicitly. For Stroop efficiency at the per-recording level, with $\Delta=16$3, the correlation is $\Delta=16$4 with two-sided $\Delta=16$5. For the subject-level Stroop composite, with $\Delta=16$6, the correlation is $\Delta=16$7 with $\Delta=16$8. All effects are described as small, with $\Delta=16$9, but directionally consistent.

The interpretation given is that BAG residuals track cognitive efficiency consistently across tasks, pointing toward utility in normative modeling of brain function (Segal et al., 7 Jul 2026). The effect sizes, however, are explicitly small. This limits overly strong causal or clinical readings. A more conservative interpretation is that the residual appears behaviourally aligned in aggregate, but the reported associations do not by themselves establish large individual-level predictive utility.

7. Reported contributions, scope, and implications

The paper enumerates several contributions and implications (Segal et al., 7 Jul 2026). First, cross-montage generalization is attributed to the unified 128-channel budget and mask-aware PMA with coordinate gating, allowing heterogeneous montages such as 115 and 128 channels to be absorbed without retraining. Second, small labeled cohorts and non-stationarity are addressed by self-supervision on 47,703 sessions, producing embeddings with rich subject-level structure so that a lightweight frozen probe can support strong downstream performance. Third, superior age regression is summarized by the 69.7% MAE reduction relative to the 10-year baseline across the pediatric-to-older-adult age range. Fourth, multi-task transfer is expressed through rank-1 leaderboard performance on three diverse EEG tasks from a single backbone. Fifth, behavioural validity is linked to the BAG findings. Sixth, proposed clinical and research applications include a portable EEG foundation model for brain-age biomarker development in neurology and psychiatry, a pre-trained encoder for rapid adaptation to BCI tasks, cognitive-state decoding, and large-scale pediatric EEG analysis, and a framework for harmonizing multi-site EEG studies under a shared representation.

The article’s concluding summary states that shallow-target spatio-temporal JEPA training on large, heterogeneous resting-state and task EEG corpora yields robust age-discriminative representations that transfer across sex, psychopathology, and behavioural-capacity tasks (Segal et al., 7 Jul 2026). Within the scope of the reported experiments, this positions STST-JEPA as a foundation-style EEG pretraining framework centered on masked latent prediction with auxiliary reconstruction.

A potential misconception is that “shallow-target” implies architectural simplicity in the encoder itself. The reported model is not shallow in depth: the context encoder has 24 transformer layers (Segal et al., 7 Jul 2026). Rather, “shallow-target” refers specifically to the prediction target being tokenizer-level representations from the EMA branch instead of deep encoder features. This distinction is central to understanding both the naming and the pretraining objective.