⚡ electrical engineering

PRiSE-EEG: A Prior-Guided Foundation Model with Depth-Stratified Experts for Cross-Paradigm EEG Representation Learning

PRiSE-EEG is a prior-guided foundation model that leverages CKA-calibrated depth-stratified experts to resolve optimization conflicts and balance shared versus specialized learning across heterogeneous EEG paradigms, achieving superior cross-paradigm performance on 12 benchmarks.

Original authors: Wei Xiong, Jiangtong Li, Kun Zhu, Jie Li

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Wei Xiong, Jiangtong Li, Kun Zhu, Jie Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a single student to be an expert in six very different subjects: sleep analysis, emotion detection, motor control, and more. The problem is that these subjects often contradict each other. What helps you learn to recognize a sleeping brain might confuse you when you're trying to detect a seizure.

This is the challenge with EEG (brainwave) foundation models. Existing AI models try to learn all these tasks at once, but they often get stuck because the "rules" for one task mess up the "rules" for another.

The paper introduces PRiSE-EEG, a new AI model designed to solve this. Here is how it works, explained simply:

1. The Problem: The "One-Size-Fits-All" Mistake

The authors first looked at how current models learn. They found two big issues:

Conflicting Instructions: When the model tries to learn from different brainwave datasets, the math tells it to move in opposite directions. It's like a teacher telling a student to "run fast" for one test and "walk slowly" for the next, all at the same time. The student gets confused and learns nothing well.
Ignoring the Map: Current models often treat brain sensors (electrodes on the head) as just random numbers. They ignore the fact that the brain has a physical map (front, back, left, right) and that different parts of the brain work together in networks.

2. The Solution: A "Layered Team" of Experts

The authors discovered something interesting about how the brain (and AI) processes information.

The Early Layers are General: The first few layers of processing are good at seeing the "big picture" that applies to all brain tasks (like basic rhythms).
The Deep Layers are Specialized: The later layers need to get very specific to handle unique tasks (like distinguishing a seizure from a normal wave).

PRiSE-EEG uses this insight to build a "Depth-Stratified Expert" system. Think of it like a company with a team of specialists:

The Junior Team (Early Layers): These are "Shared Experts." They are generalists who handle the common rules that apply to everyone. They make sure the model doesn't forget the basics.
The Senior Team (Deep Layers): These are "Specialized Experts." They are hired specifically for the hard, unique details of each task.

Instead of forcing every layer to use the same mix of experts, PRiSE-EEG uses a smart calculator (based on something called CKA) to decide exactly how many generalists vs. specialists each layer needs. It's like a manager who knows, "For the first step of this project, we need a generalist; for the final step, we need a specialist."

3. The Input: Using a "Brain Map"

To feed data into this system, the model doesn't just look at raw numbers. It uses Prior-Guided Tokenization.

The Analogy: Imagine reading a book. Instead of just looking at individual letters, you group them into words and sentences based on grammar rules you already know.
The Application: PRiSE-EEG groups brain sensors based on their location (Frontal, Parietal, etc.) and how they naturally connect. It uses these "weak hints" (priors) to organize the data before the AI even starts learning. This helps the model understand the brain's physical structure from day one.

4. The Results: Better at Everything

The team tested this new model on 12 different brainwave datasets covering things like:

Detecting sleep stages.
Recognizing emotions (happy, sad, angry).
Identifying seizures.
Detecting Alzheimer's disease.

The Outcome: PRiSE-EEG beat all previous models. It was better at handling the "conflicting instructions" because it knew when to be general and when to be specific. It also performed better because it respected the physical map of the brain.

Summary

PRiSE-EEG is like a smart student who knows that learning to drive a car and learning to fly a plane require different skills. Instead of trying to use the exact same brainpower for both, this model has a flexible system:

It organizes brain data using a map of the head.
It uses a "generalist" team for the basics.
It switches to a "specialist" team for the hard, specific details.

This approach allows it to learn many different brain tasks simultaneously without getting confused, making it the most accurate model of its kind tested so far.

Technical Summary: PRiSE-EEG

Problem Statement

Electroencephalography (EEG) foundation models (EEG-FMs) aim to learn reusable representations across heterogeneous paradigms. However, existing approaches face two primary limitations:

Optimization Conflicts: Heterogeneous EEG tasks (e.g., sleep staging vs. emotion recognition) often induce conflicting gradient update directions. Dense Transformer models using uniform parameters struggle to reconcile these divergent optimization goals, leading to suboptimal cross-paradigm transfer.
Underutilized Priors: Current tokenization strategies often ignore structured spatial regularities inherent in sensor-space EEG. While electrode coordinates are known, the structured organization of brain networks and short-window channel interactions are frequently underused as inductive biases for patch formation.

Furthermore, existing MoE (Mixture-of-Experts) approaches in EEG often allocate shared and specialized experts uniformly across all layers, ignoring the observation that representation geometry changes with depth (shallow layers being more general, deeper layers more specific).

Methodology: PRiSE-EEG

The authors propose PRiSE-EEG (Prior-guided EEG foundation model with depth-Stratified Experts), a framework designed to balance shared EEG regularities with task-specific specialization. The architecture consists of three core components:

1. Prior-Guided Continuous Tokenization

Instead of discrete tokenization, PRiSE-EEG forms continuous multi-channel EEG patches by integrating weak sensor-space priors with dynamic interactions:

Static Path: Utilizes learnable queries and attention biases based on predefined anatomical regions (e.g., Frontal, Parietal) and functional networks (e.g., Default Mode Network, Salience Network). A learnable relaxation term ( $\Delta_{pri}$ ) allows the model to adjust these weak priors during training rather than enforcing rigid anatomical assignments.
Dynamic Path: Captures sample-dependent, sub-second channel interactions. It subdivides patches into micro-frames and computes lag-limited channel affinities to form dynamic group queries. A Gumbel selector identifies the most informative frame-pair groups for each sample.
Output: These paths generate continuous group tokens that handle montage heterogeneity while injecting neuroscientific structure.

2. CKA-Calibrated Depth-Stratified Experts

To address optimization conflicts and layer-wise representation differences, the model employs a Mixture-of-Experts (MoE) architecture where the allocation of shared vs. specialized experts is determined by the layer depth:

Pilot Analysis: Before training the main model, a dense "pilot" Transformer is trained on mixed-paradigm batches. Centered Kernel Alignment (CKA) is computed across layers to measure the ratio of between-paradigm similarity to within-paradigm similarity.
Calibration Rule: Layers with high cross-paradigm similarity (shallow layers) are assigned a higher proportion of shared experts. Layers with lower similarity (deeper layers) are assigned more specialized experts.
Implementation: This allocation is mapped via a sigmoid function based on the pilot's CKA metrics and fixed during the main training phase. This ensures that early blocks preserve general EEG regularities while later blocks specialize in task-specific transformations.

3. Training Strategy

Pre-training: The model is pre-trained on 17 diverse public datasets (over 28,000 hours) using three self-supervised objectives: Masked Temporal Patch reconstruction (MAE-TP), Masked Channel reconstruction (MAE-CH), and Future Patch prediction (GPT).
Fine-tuning: PRiSE-EEG supports unified multi-dataset fine-tuning with task-specific dual-stream MLP heads (averaging over temporal and group dimensions separately) to preserve structural information without parameter explosion.

Key Contributions

Characterization of Cross-Paradigm Conflicts: The authors empirically demonstrate that heterogeneous EEG tasks induce conflicting gradient updates and that representation similarity follows a depth-wise transition (general to specific), motivating non-uniform expert allocation.
CKA-Calibrated Depth-Stratified Experts: A novel mechanism that converts block-wise CKA sharedness into shared expert proportions using a simple sigmoid rule, dynamically adjusting capacity allocation based on representation geometry rather than uniform heuristics.
Prior-Guided Continuous Tokenization: A tokenization scheme that combines weak static region/network priors with dynamic short-window interactions, avoiding rigid anatomical constraints while leveraging structured sensor-space information.
Comprehensive Validation: The model is validated on 12 public benchmarks covering 10 task categories (e.g., seizure detection, emotion recognition, sleep staging), demonstrating robust cross-paradigm performance.

Experimental Results

Evaluated on 12 EEG datasets, PRiSE-EEG (specifically the B variant with ~41M parameters) achieved state-of-the-art or competitive results under matched protocols:

Performance Gains: It showed significant improvements (6–8% on primary metrics) on disease-related and cognitive tasks (e.g., ADFTD, SEED-V, Workload) and consistent 2–5% gains on tasks like TUEV and PhysioMI.
Ablation Studies:
- Removing either the static or dynamic path in tokenization caused >3% performance degradation.
- Reversing the expert ratio sequence (assigning more specialized experts to shallow layers) caused the largest drop (>8%), confirming the necessity of depth-stratified allocation.
- Using fixed shared expert ratios was inferior to the CKA-calibrated approach.
Scaling: The model exhibits predictable scaling behavior with both model size (up to 125M parameters) and pre-training data volume (up to 28,000 hours), with high $R^2$ values indicating performance continues to improve with scale.
Interpretability: Grad-CAM visualizations confirmed the model attends to neuroscientifically valid regions (e.g., frontal/temporal for emotion, parietal for sleep stages) and correctly distinguishes between AD and FTD based on known pathological patterns.

Significance and Claims

The paper claims that PRiSE-EEG provides a compact framework for balancing shared EEG regularities with task-specific specialization. Its significance lies in:

Moving beyond uniform adaptation: Demonstrating that expert specialization in foundation models should be depth-stratified according to the representation structure learned by dense pilot models, rather than being input-adaptive only.
Leveraging weak priors: Showing that incorporating weak sensor-space priors (regions and networks) and dynamic interactions improves robustness without requiring hard anatomical constraints.
Unified Learning: Validating that a single pre-trained backbone can effectively adapt to diverse clinical and cognitive tasks through unified fine-tuning, reducing the need for task-specific architectures.

The authors modestly note that while the model scales effectively, it remains smaller than contemporary language/vision foundation models, and future work is needed to extend the framework to multimodal scenarios (e.g., Video-EEG) and achieve zero-shot generalization to novel tasks.