Neural State Prediction: Obstructing Shortcut Learning in EEG Foundation Models

Kieren Yu; Ziyang Liu; Chang Huang; Jintai Chen; Kaishun Wu

cs.AISep 25, 2026

Neural State Prediction: Obstructing Shortcut Learning in EEG Foundation Models

Authors: Kieren Yu, Ziyang Liu, Chang Huang, Jintai Chen, Kaishun Wu

Organizations: Hong Kong University of Science and Technology (Guangzhou)

Abstract

EEG foundation models increasingly use masked prediction to learn from unlabeled recordings, but optimizing this objective does not ensure transferable neural representations. A central challenge is that stable positional cues and local correlations can make masked regions predictable without integrating distributed neural context. To reduce this reliance on low-information prediction paths, we introduce Neural State Prediction (NSP), a latent-predictive framework that constrains both the prediction target and the available context. NSP uses a Target Encoder updated by an exponential moving average (EMA) to define latent supervision. Identity residualization removes additive effects associated with channel identity and relative time from the targets, while topology-separated context excludes their immediate spatial and temporal neighborhood from the visible input. We pretrain NSP on 2.2 million EEG segments from TUEG and evaluate it across 30 downstream datasets spanning clinical diagnosis, sleep staging, emotion recognition, motor imagery, event-related potentials, cognitive-state decoding, and language retrieval. Under full-parameter multi-task fine-tuning on EEG-FM-Bench, NSP achieves 63.94 macro balanced accuracy across 14 datasets, exceeding the strongest evaluated baseline by 2.35 percentage points. Controlled component ablations assess the contribution of each mechanism, while matched context controls and held-out interventions characterize the role of context geometry, signal content, and positional information. Jointly designing latent targets and their context offers a promising direction for EEG foundation models that learn from distributed signal structure.

Figures & tables

Figure 1: EEG-FM-Bench macro summary.

Figure 2: Comparison of EEG foundation-model pretraining paradigms. The panels contrast view-level agreement, signal reconstruction, discrete or autoregressive token prediction, and NSP’s Target–Context latent prediction.

Figure 3: Neural State Prediction architecture. Residualized, stop-gradient Target states supervise the masked Context branch through its predictor. A separate complete-input pass through the shared Context Encoder supplies moment-based state regularization. The Target Encoder follows the Context Encoder by EMA; prediction is scored only at target positions.

Task	Dataset	Channels	Window (s)	Samples	Subjects	Classes	Task	Dataset	Channels	Window (s)	Samples	Subjects	Classes
Alzheimer’s identification	ADFTD	19	10	7,013	88	3	Motor imagery	Mimul-11	60	5	41,347	11	3
Depression diagnosis	Mumtaz2018	19	5	14,683	64	2		PhysioMI	64	4	9,747	109	4
Parkinson’s diagnosis	Singh2021	60	5	39,543	129	2		BCIC-2a	22	4	5,088	9	4
Schizophrenia diagnosis	Albrecht2019	62	5	24,928	77	2		Scherer2015	30	4	3,550	9	5
Abnormal classification	TUAB	23	30	272,320	2,383	2	Motor execution	Srisrisawang2024	60	3	19,157	20	16
Anomalous event detection	TUEV	21	5	113,353	370	6	Visual target detection	Things-EEG-2	63	5	41,570	10	2

Table 1: Overview of 16 selected NeuralBench-EEG-Core datasets and 14 EEG-FM-Bench datasets ( Banville et al., 2026 ; Xiong et al., 2025 ) . Dataset–task pairs shared by the two benchmarks are excluded from the NeuralBench-EEG-Core selection. Sampling follows each benchmark’s dataset-specific preprocessing; NSP uses 200 Hz for EEG-FM-Bench. Benchmark details and complete results are provided in Appendix E .

Figure 4: EEG-FM-Bench radar summary. Dataset-wise BAC under four adaptation protocols across all 14 datasets; radial positions use per-dataset relative scaling and NSP is highlighted in red.

Dataset	Metric	BENDR	BIOT	LaBraM	CBraMod	LUNA	REVE	NSP
Clinical and neurological disease
Mumtaz2018	BAC	86.22 $\pm$ 00.47	79.11 $\pm$ 03.39	85.61 $\pm$ 03.43	86.04 $\pm$ 02.48	85.42 $\pm$ 00.98	87.33 $\pm$ 01.44	90.43 $\pm$ 06.40
Singh2021	BAC	65.82 $\pm$ 05.62	74.18 $\pm$ 04.16	76.64 $\pm$ 04.06	74.37 $\pm$ 02.38	64.95 $\pm$ 00.49	72.58 $\pm$ 01.72	76.67 $\pm$ 01.70
Albrecht2019	BAC	57.17 $\pm$ 00.87	56.18 $\pm$ 01.50	63.62 $\pm$ 01.93	57.05 $\pm$ 03.12	66.75 $\pm$ 02.18	49.21 $\pm$ 10.36	68.43 $\pm$ 02.71
CHB-MIT	BAC	78.97 $\pm$ 01.04	73.58 $\pm$ 05.20	82.09 $\pm$ 00.75	82.89 $\pm$ 01.58	70.46 $\pm$ 02.44	76.51 $\pm$ 03.73	81.10 $\pm$ 01.39
BCI and sensorimotor decoding

Table 2: Primary-metric performance on 16 selected NeuralBench-EEG-Core datasets ( Banville et al., 2026 ) . BAC (%) is shown for 14 classification tasks and Top-5 accuracy (%) for two retrieval tasks. Dataset-level values are mean

\pm

sample SD; the macro row averages the 14 BAC means. Red and blue mark the best and second-best mean for each task and for the macro row.

Pretraining design						Frozen Single-Task BAC
Variant	Block mask	Target	ResID	Exclusion	Reg.	TUAB	TUEV	SEED	PhysioMI	Macro
Target construction
Shared-target baseline	–	Shared	–	None	$\checkmark$	73.27 $\pm$ 00.35	38.37 $\pm$ 02.41	40.53 $\pm$ 01.78	25.57 $\pm$ 01.42	44.44
+ Block masking	$\checkmark$	Shared	–	None	$\checkmark$	74.27 $\pm$ 00.23	44.17 $\pm$ 01.00	38.37 $\pm$ 04.12	26.20 $\pm$ 00.78	45.75
+ EMA target	$\checkmark$	EMA	–	None	$\checkmark$	78.17 $\pm$ 00.23	51.40 $\pm$ 00.98	51.97 $\pm$ 00.32	25.97 $\pm$ 01.12	51.88
+ Identity residualization	$\checkmark$	EMA	$\checkmark$	None	$\checkmark$	78.73 $\pm$ 00.15	59.77 $\pm$ 02.17	51.97 $\pm$ 01.27	26.50 $\pm$ 00.35	54.24

Table 3: Component ablations: frozen single-task transfer with 10% labels. ResID: identity residualization; Reg.: state regularization. Block mask “–” denotes temporal masking, not unmasked input. Exclusion means additional context removal; Random is size matched. BAC is reported as % (mean

\pm

sample SD); Macro averages four tasks. Red and blue mark the best and second-best mean per column.

Figure 5: Task-wise transfer across pretraining duration. BAC (0–1 scale) across checkpoints under frozen single-task transfer (10% labels) and single-task full fine-tuning.

Figure 6: Label efficiency across downstream tasks. BAC (0–1 scale) versus labeled-data fraction on TUAB, TUEV, SEED, and PhysioMI; bands show variability across evaluations.

Figure 7: Representation geometry and spatial evidence. Top: raw, frozen, and fine-tuned HMC sleep-stage representations visualized by t-SNE. Bottom: TUEV event-class Grad-CAM maps and interchannel correlations. The rows depict different datasets.

Appendix figures & tables36 assets

Supplementary material from the paper’s appendix.

Appendix

Input & Tokenization
Hyperparameter	Configuration
Segment duration	6 s
Sampling rate	200 Hz
EEG channels	19
Patch duration	0.2 s
Context Encoder & Predictor

Appendix

Table 4: Default architecture and pretraining configuration of NSP.

Hyperparameter	Configuration
Pretraining corpus	TUEG ( Obeid and Picone, 2016 )
Pretraining examples	2.2M segmented examples
Optimizer	AdamW ( Loshchilov and Hutter, 2017 )
Learning rate	$1\times 10^{-4}$
Adam $(\beta_{1},\beta_{2})$	$(0.9,0.999)$
Weight decay	0.05

Appendix

Table 5: Pretraining and optimization hyperparameters of NSP.

Figure 8: NSP pretraining diagnostics. Lines show the across-run mean and shaded regions show sample SD for optimization losses and representation-health statistics. The final panel displays the learning-rate and EMA schedules.

Mumtaz2018				Singh2021
Model	BAC	AUROC	AUCPR	Model	BAC	AUROC	AUCPR
BENDR	86.22 $\pm$ 00.47	92.20 $\pm$ 00.26	91.17 $\pm$ 00.27	BENDR	65.82 $\pm$ 05.62	70.83 $\pm$ 06.33	85.33 $\pm$ 03.73
BIOT	79.11 $\pm$ 03.39	86.20 $\pm$ 01.13	82.88 $\pm$ 05.00	BIOT	74.18 $\pm$ 04.16	81.60 $\pm$ 04.45	88.86 $\pm$ 02.51
LaBraM	85.61 $\pm$ 03.43	88.33 $\pm$ 05.14	87.57 $\pm$ 03.82	LaBraM	76.64 $\pm$ 04.06	87.32 $\pm$ 04.69	94.19 $\pm$ 02.42
CBraMod	86.04 $\pm$ 02.48	92.57 $\pm$ 02.82	93.63 $\pm$ 01.06	CBraMod	74.37 $\pm$ 02.38	83.18 $\pm$ 04.65	91.94 $\pm$ 02.19
LUNA	85.42 $\pm$ 00.98	94.04 $\pm$ 00.69	94.43 $\pm$ 01.02	LUNA	64.95 $\pm$ 00.49	73.50 $\pm$ 01.11	87.02 $\pm$ 00.59

Appendix

Table 6: NeuralBench clinical and neurological-disease results (%, mean

\pm

sample SD). Red and blue mark the best and second-best mean for each metric.

Thielen2021				Scherer2015
Model	BAC	F1	$\kappa$	Model	BAC	F1	$\kappa$
BENDR	4.96 $\pm$ 00.24	1.64 $\pm$ 00.37	-0.04 $\pm$ 00.26	BENDR	21.62 $\pm$ 00.66	19.11 $\pm$ 01.61	2.02 $\pm$ 00.83
BIOT	4.80 $\pm$ 00.25	3.84 $\pm$ 00.04	-0.21 $\pm$ 00.26	BIOT	26.48 $\pm$ 02.48	25.69 $\pm$ 03.03	8.10 $\pm$ 03.10
LaBraM	86.39 $\pm$ 00.74	86.48 $\pm$ 00.78	85.67 $\pm$ 00.78	LaBraM	21.83 $\pm$ 00.17	20.34 $\pm$ 00.83	2.29 $\pm$ 00.21
CBraMod	86.37 $\pm$ 01.33	86.38 $\pm$ 01.31	85.66 $\pm$ 01.40	CBraMod	21.22 $\pm$ 01.80	15.89 $\pm$ 05.21	1.52 $\pm$ 02.25
LUNA	87.06 $\pm$ 00.28	87.05 $\pm$ 00.28	86.38 $\pm$ 00.30	LUNA	20.95 $\pm$ 01.09	14.21 $\pm$ 06.01	1.19 $\pm$ 01.37

Appendix

Table 7: NeuralBench BCI and sensorimotor results (%, mean

\pm

sample SD). Red and blue mark the best and second-best mean for each metric.

ERP CORE-ERN				ERP CORE-N170
Model	BAC	AUROC	AUCPR	Model	BAC	AUROC	AUCPR
BENDR	78.06 $\pm$ 04.59	86.11 $\pm$ 04.33	96.37 $\pm$ 01.35	BENDR	67.03 $\pm$ 00.61	71.94 $\pm$ 01.25	47.25 $\pm$ 02.01
BIOT	49.40 $\pm$ 01.12	50.54 $\pm$ 06.31	82.92 $\pm$ 04.01	BIOT	49.84 $\pm$ 00.36	50.20 $\pm$ 00.46	25.55 $\pm$ 00.46
LaBraM	81.96 $\pm$ 01.52	90.10 $\pm$ 01.67	97.18 $\pm$ 00.75	LaBraM	72.14 $\pm$ 00.55	78.75 $\pm$ 01.01	57.33 $\pm$ 02.71
CBraMod	81.32 $\pm$ 01.04	89.43 $\pm$ 00.77	97.31 $\pm$ 00.29	CBraMod	71.14 $\pm$ 00.55	78.72 $\pm$ 01.37	57.09 $\pm$ 03.52
LUNA	81.57 $\pm$ 03.26	89.74 $\pm$ 02.40	97.09 $\pm$ 01.05	LUNA	68.74 $\pm$ 02.04	76.63 $\pm$ 01.34	54.25 $\pm$ 01.05

Appendix

Table 8: NeuralBench event-related and evoked-response results (%, mean

\pm

sample SD). Red and blue mark the best and second-best mean for each metric.

Zyma2019				Hinss2023
Model	BAC	AUROC	AUCPR	Model	BAC	F1	$\kappa$
BENDR	65.28 $\pm$ 02.34	74.69 $\pm$ 05.12	89.33 $\pm$ 03.31	BENDR	45.89 $\pm$ 00.19	45.01 $\pm$ 01.05	18.83 $\pm$ 00.28
BIOT	66.40 $\pm$ 03.04	72.30 $\pm$ 05.30	89.05 $\pm$ 01.27	BIOT	43.18 $\pm$ 04.73	42.35 $\pm$ 05.03	14.77 $\pm$ 07.09
LaBraM	68.92 $\pm$ 02.42	80.90 $\pm$ 03.78	92.18 $\pm$ 02.89	LaBraM	65.46 $\pm$ 03.17	63.16 $\pm$ 04.69	48.20 $\pm$ 04.76
CBraMod	72.02 $\pm$ 01.30	77.04 $\pm$ 00.39	91.57 $\pm$ 00.49	CBraMod	61.37 $\pm$ 01.70	59.32 $\pm$ 02.71	42.06 $\pm$ 02.54
LUNA	72.35 $\pm$ 00.41	76.85 $\pm$ 00.69	90.16 $\pm$ 02.05	LUNA	66.01 $\pm$ 00.94	64.59 $\pm$ 01.16	49.01 $\pm$ 01.40

Appendix

Table 9: NeuralBench cognitive-state and language results (mean

\pm

sample SD). All metrics except median rank (MedR) are percentages. Retrieval uses Top-1, Top-5, and MedR; MedR is an unscaled rank and lower is better. Red and blue mark the best and second-best mean for each metric.

Full Fine-Tuning				Parameter-Efficient Multi-Task
Model	Single-task			Model	Frozen
	BAC	F1	$\kappa$		BAC	F1	$\kappa$
BENDR	59.50 $\pm$ 00.42	58.88 $\pm$ 00.98	39.42 $\pm$ 00.64	BENDR	34.12 $\pm$ 00.24	24.48 $\pm$ 02.82	01.10 $\pm$ 00.32
BIOT	63.87 $\pm$ 01.77	63.12 $\pm$ 01.94	45.99 $\pm$ 02.67	BIOT	57.85 $\pm$ 00.19	58.17 $\pm$ 00.21	37.17 $\pm$ 00.29
LaBraM	61.59 $\pm$ 01.71	60.30 $\pm$ 01.60	43.28 $\pm$ 02.48	LaBraM	52.13 $\pm$ 00.05	51.23 $\pm$ 00.34	28.33 $\pm$ 00.12
EEGPT	70.83 $\pm$ 00.28	70.48 $\pm$ 00.16	56.41 $\pm$ 00.42	EEGPT	61.80 $\pm$ 00.65	60.41 $\pm$ 01.07	42.89 $\pm$ 00.97

Appendix

Table 10: Complete SEED results across four adaptation settings (%, mean

\pm