Stabilized Transformer EEGXF: Robust EEG Decoding

• The paper presents a stabilized Transformer architecture that leverages a shallow, norm-first encoder and the Stabilized Lottery Ticket hypothesis to enable robust EEG decoding.
• It employs multi-head self-attention with regularization techniques such as BatchNorm, dropout, EMA, and gradient clipping to improve training stability and mitigate overfitting.
• The model achieves a balance between moderate in-distribution accuracy and enhanced robustness to temporal shifts and out-of-distribution uncertainty, making it ideal for challenging EEG applications.

The Stabilized Transformer (EEGXF) is a neural architecture designed for robust and efficient sequence modeling, combining a lightweight, stabilized Transformer encoder—originally motivated by the Stabilized Lottery Ticket Hypothesis—with targeted architectural adaptations for challenging domains such as naturalistic EEG decoding. This model achieves a distinctive trade-off between moderate in-distribution accuracy and enhanced robustness to temporal distribution shifts and out-of-distribution (OOD) uncertainty, distinguishing it from both classical CNN/LSTM and state-space alternatives in recent benchmarks (Ergezer, 29 Jan 2026).

1. Architectural Foundations and Stabilization Techniques

EEGXF incorporates a norm-first, shallow Transformer encoder with modifications intended to improve convergence and regularization for sequence-modality tasks. The input to the model is an EEG tensor $X\in\mathbb{R}^{B\times C\times T}$, with $B$ batches, $C=64$ channels, and $T$ temporal steps (e.g., $T=16,000$ for 64 s at 250 Hz). Key architectural components include:

• Input projection and embedding: The input is linearly projected via $W_{\mathrm{in}}\in\mathbb{R}^{C\times d}$ (with $d=128$), followed by ReLU activation and BatchNorm. The formula is:

$$x_{\mathrm{emb}} = \mathrm{ReLU}\Bigl(\mathrm{BatchNorm}(X^T W_{\mathrm{in}} + b_{\mathrm{in}})\Bigr) + P$$

where $P$ is a learned positional embedding $P\in\mathbb{R}^{T\times d}$.

• Transformer encoder: Stacked $L=2$ layers of multi-head self-attention (MHSA) with $H=4$ heads ($d_h=32$), using pre-LayerNorm (“norm-first” convention). Each layer computes queries, keys, values via

$$Q = \mathrm{LN}(x^{(\ell-1)})W^Q, \quad K = \mathrm{LN}(x^{(\ell-1)})W^K, \quad V = \mathrm{LN}(x^{(\ell-1)})W^V$$

and then

$$\mathrm{head}_i = \mathrm{softmax}\Bigl(\frac{Q_i K_i^T}{\sqrt{d_h}}\Bigr)V_i$$

$$\mathrm{MHSA}(X) = [\mathrm{head}_1;\dots;\mathrm{head}_H]W^O$$

Each attention output is fed into a residual connection, followed by a position-wise feedforward block:

$$\mathrm{FFN}(z) = W_2 \bigl[\mathrm{ReLU}(W_1 z + b_1)\bigr] + b_2$$

• Pooling and classification: A learned attention pooling replaces global average, computing:

$$S = \mathrm{softmax}\Bigl(\frac{q^T (XW_k)}{\sqrt{d}}\Bigr), \quad u = S (X W_v)$$

$u$ is then mapped by a single linear classifier to logits; cross-entropy loss is applied.

• Stabilization: The architecture is explicitly regularized via:
  • BatchNorm on inputs
  • LayerNorm with norm_first=True
  • Dropout ($p=0.1$) after each MHSA and feedforward block
  • High-variance (scaled Glorot) weight initialization
  • Gradient clipping (global norm 1.0)
  • Weight exponential moving average (EMA, decay=0.999)

2. Training Protocols and Optimization

EEGXF models are trained with AdamW optimizer ($\beta_1=0.9$, $\beta_2=0.999$, weight decay 0.01) and monitored using a ReduceLROnPlateau schedule (patience=5 epochs, factor=0.5), with early stopping (patience=10) up to 100 epochs. Input segments are standardized per channel. Dropout and BatchNorm function as primary regularizers. Training and validation employ common cross-entropy loss:

$$\mathcal{L} = -\sum_{c=1}^4 y_c^{\mathrm{true}} \log y_c^{\mathrm{pred}}$$

Gradient clipping and EMA further enforce training stability.

3. Pruning and Sparsity: Stabilized Lottery Ticket Hypothesis

The broader stabilization paradigm is based on the Stabilized Lottery Ticket (SLT) Hypothesis, which postulates that sparse subnetworks of a neural model, identified after a short “warmup” and pruned by magnitude from a converged solution, can be rewound and retrained to match dense performance. The process follows:

1. Given $\theta_0$ (initial parameters), train to early checkpoint $\theta_t$ ($t=0.05T$).
2. Train to convergence ($\theta_T$), then determine binary mask $M$ by keeping the $(1-s)n$ parameters of largest $|\theta_T|$.
3. Rewind: set $\theta_0^{\rm sparse} = \theta_t \odot M$, then retrain with full learning-rate schedule.

Formally, for elementwise masking:

$$(\theta \odot M)_i = \theta_i \cdot M_i, \;\; M_i\in\{0,1\}$$

This approach, with no fine-tuning, matches iterative magnitude pruning (MP) up to $s \approx 85\%$ sparsity, and is further enhanced at extreme sparsity ($>85\%$) by coupling SLT pruning with subsequent iterative MP (the SLT–MP protocol), yielding superior BLEU scores in translation settings and substantial FLOPs reduction (Brix et al., 2020).

4. In-Distribution Performance and Efficiency

In EEG decoding benchmarks on the 4-class HBN movie-watching dataset, EEGXF achieves the following mean accuracies over increasing temporal context:

Segment length (s)	Accuracy (%)
8	60.3
16	63.4
32	80.1 ± 2.0
64	81.8 ± 1.8
128	76.7 ± 1.5

With ≈0.5 million parameters, EEGXF is highly lightweight but lags S5 (98.7 ± 0.6%) and CNN baselines (98.3 ± 0.3%) at 64s, though models like CNN require nearly 20x the parameters (Ergezer, 29 Jan 2026).

5. Robustness and Generalization Properties

EEGXF exhibits distinct robustness profiles in several settings:

• Zero-shot cross-frequency shift: Trained at 250 Hz, tested at 128/64 Hz (after anti-alias filtering), EEGXF shows minimal (<2 pp) accuracy drop (e.g., 40.9% at 250 Hz, 39.7% at 128 Hz, and 40.7% at 64 Hz), demonstrating strong temporal-resolution invariance.
• Leave-one-subject-out (LOSO) generalization: Mean accuracy is 48.4% (Std = 12.1%, $n=20$ runs), 7.5 pp lower than S5 (55.9% ± 15.9%), with statistical significance ($t_{19}=3.13$, $p=0.0055$).
• Zero-shot cross-task OOD: On 245 unseen tasks, EEGXF defaults to neutral class (“Resting State”, 26.0% confidence), in contrast to S5 (overconfident “Movie 3”, 60.0%). EEGXF thus demonstrates “conservative collapse,” reflecting greater uncertainty on OOD data.
• In-distribution calibration: On 64s segments, EEGXF yields NLL = 0.999 ± 0.023 (vs. S5 = 0.056, CNN = 0.079), Brier score = 0.077 ± 0.005, and ECE = 13.4% ± 1.7, indicating poor in-distribution calibration but advantageous OOD uncertainty estimation.

6. Comparative Assessment and Use Cases

EEGXF’s trade-offs are summarized in terms of validation accuracy, model size, and robustness:

Model	Params (M)	Peak Acc (%)	Cross-Freq. Shift	OOD Uncertainty	Cross-Subj. Gen.	In-Dist. Calibration
S5	0.18	98.7	Moderate	Overconfident	Best	Best
EEGXF	0.5	81.8	Best	Conservative	Good	Poor
CNN	4.4	98.3	Weak	Poor	Weak	Best

EEGXF is favored when resilience to temporal resampling and OOD uncertainty are critical, such as in clinical or edge-deployment scenarios.

7. Theoretical and Practical Insights

Empirical evidence supports the SLT hypothesis: sparse subnetworks, if pruned using the converged parameter magnitudes and appropriately rewound, retain the functional capacity of dense models given sufficient retraining. The key parameter is the sign, not magnitude, of the initial weights—a “Constant Lottery Ticket” (CLT) with fixed sign and constant per-layer magnitude performs on par with SLT (Brix et al., 2020). In the EEG context, EEGXF leverages these insights, combining parameter efficiency and robustness, albeit at a cost to in-distribution accuracy and calibration relative to S5 and CNN. This suggests applications demanding reliable uncertainty estimates or deployment in shifting or unpredictable environments may benefit from EEGXF’s distinctive architecture and training regimen.