Riemannian High-Order Pooling for Brain Foundation Models

This is the official PyTorch implementation of RHOP (ICLR 2026), a plug-and-play Riemannian classification head that augments EEG foundation backbones (BIOT/LaBraM) with stable first-order and second-order token statistics for robust and efficient EEG decoding across diverse downstream tasks.

Abstract

Electroencephalography (EEG) is a noninvasive technique for measuring brain electrical activity that supports a wide range of brain-computer interaction applications. Motivated by the breakthroughs of Large Language Models (LLMs), recent efforts have begun to explore Large EEG foundation Models trained on broad unlabeled corpora. However, most advances focus on improving the backbone while neglecting the classification head. Existing models often rely on a single class token, underutilizing the spatiotemporal structure and second-order statistics that are crucial for EEG decoding. We propose Riemannian High Order Pooling (RHOP), a plug-and-play module that injects principled Riemannian statistics into the classifier. RHOP maps each token to a quotient Gaussian jointly encoding mean and second-order information, yielding scale-invariant descriptors. Tokens are then aggregated by estimating a Riemannian Gaussian on the SPD manifold, where the Fréchet mean and covariance are embedded into an SPD descriptor. The resulting normalized vector is fused with the class token for prediction. RHOP is backbone-agnostic and integrates with modern EEG foundation models (e.g., BIOT and LaBraM). Across diverse EEG benchmarks, it improves accuracy and efficiency under full fine-tuning, linear probing, and from-scratch training settings.

Framework

Repository Layout

RHOP/
├── models/                       # BIOT, LaBraM, and other model modules
├── cov_pooling/                  # RHOP and covariance pooling methods
├── process/
│   ├── biot/                     # BIOT preprocessing scripts (TUAB/TUEV/CHB-MIT/SHHS/PTB-XL)
│   └── labram/                   # LaBraM preprocessing scripts and split metadata
├── checkpoints/
│   ├── biot/                     # Place downloaded BIOT checkpoints here
│   └── labram/                   # Place downloaded LaBraM checkpoints here
├── training_utils/               # Shared utilities for BIOT/LaBraM training
├── train_biot_binary.py          # BIOT binary classification training
├── train_biot_multiclass.py      # BIOT multiclass classification training
└── train_labram.py               # LaBraM fine-tuning with RHOP

Environment Setup

conda create -n labram python=3.10 -y
conda activate labram
pip install torch torchvision torchaudio
pip install pytorch-lightning timm pyhealth mne scipy einops pandas tqdm wfdb pyedflib openpyxl tensorboard

Checkpoint Download

Please download them from upstream sources, then place them under checkpoints/biot and checkpoints/labram.

mkdir -p checkpoints/biot checkpoints/labram

BIOT checkpoints

• Official BIOT pretrained folder:
  https://github.com/ycq091044/BIOT/tree/main/pretrained-models
• HuggingFace mirror (EEG-PREST-16-channels):
  https://huggingface.co/braindecode/biot-pretrained-prest-16chs
• HuggingFace mirror (EEG-SHHS+PREST-18-channels):
  https://huggingface.co/braindecode/biot-pretrained-shhs-prest-18chs
• HuggingFace mirror (EEG-six-datasets-18-channels):
  https://huggingface.co/braindecode/biot-pretrained-six-datasets-18chs

LaBraM checkpoints

• Official LaBraM checkpoints folder (labram-base.pth, vqnsp.pth):
  https://github.com/935963004/LaBraM/tree/main/checkpoints

Expected local paths for this repo:

• checkpoints/biot/EEG-PREST-16-channels.ckpt
• checkpoints/biot/EEG-SHHS+PREST-18-channels.ckpt
• checkpoints/biot/EEG-six-datasets-18-channels.ckpt
• checkpoints/labram/labram-base.pth
• checkpoints/labram/vqnsp.pth

Dataset Download

Please obtain datasets from their official providers before preprocessing:

• TUAB / TUEV (TUH EEG portal, request/access required):
  https://isip.piconepress.com/projects/nedc/html/tuh_eeg/

Data Preprocessing

Please edit the path variables in each script before running.

BIOT preprocessing

# TUAB
python process/biot/TUAB/process.py

# TUEV
python process/biot/TUEV/process.py

LaBraM preprocessing

# TUAB preprocessing for LaBraM
python process/labram/make_TUAB.py

# TUEV event extraction (generates processed_train/processed_eval)
python process/labram/make_TUEV.py

# TUEV subject-based resplit for LaBraM
python process/labram/make_TUEV2.py

Training

BIOT binary classification (TUAB)

python train_biot_binary.py \
  --dataset TUAB \
  --model BIOT \
  --pretrain_model_path checkpoints/biot/EEG-six-datasets-18-channels.ckpt \
  --device 0

BIOT multiclass classification (TUEV)

python train_biot_multiclass.py \
  --dataset TUEV \
  --model BIOT \
  --device 0

LaBraM fine-tuning with RHOP

python train_labram.py \
  --dataset TUEV \
  --finetune ./checkpoints/labram/labram-base.pth \
  --device cuda:0 \
  --re_type RGCP \
  --rgcp_type SICE \

Citation

If RHOP is useful in your work, please cite:

@inproceedings{hu2026riemannian,
  title={Riemannian High-Order Pooling for Brain Foundation Models},
  author={Hu, Chen and Chen, Ziheng and Wang, Rui and Zheng, Yefeng and Sebe, Nicu},
  booktitle={The Fourteenth International Conference on Learning Representations},
  year={2026},
  url={https://openreview.net/forum?id=66h1sCMm7F}
}

Acknowledgements

This codebase builds on open implementations from BIOT and LaBraM.