ST-EEGFormer — Safetensors Weights
Pre-converted
safetensors weights for the
ST-EEGFormer EEG foundation model, ready for use with
steegformer-rs (pure-Rust inference on
Burn 0.20 ) or any framework that supports safetensors.
Weights are converted from the official PyTorch
.pth checkpoints published at
LiuyinYang1101/STEEGFormer .
ST-EEGFormer won 1st Place in the NeurIPS 2025 EEG Foundation Challenge and was accepted at ICLR 2026 .
Model Files
Encoder Only (for inference / embedding extraction)
Full MAE (encoder + decoder, for reconstruction / fine-tuning)
Config
File Description config.jsonModel hyperparameters for all variants
Large V2 has undergone further pre-training on the HBN dataset for the NeurIPS 2025 EEG Foundation Challenge.
Quick Start — Rust
1 # Install
2 cargo add steegformer-rs
3
4 # Download weights
5 huggingface-cli download eugenehp/ST-EEGFormer \
6 ST-EEGFormer_small_encoder.safetensors \
7 config.json \
8 --local-dir weights/
9
10 # Run inference
11 cargo run --release --bin infer -- \
12 --config weights/config.json \
13 --weights weights/ST-EEGFormer_small_encoder.safetensors
Library API
1 use steegformer_rs :: { STEEGFormerEncoder , ModelConfig , data } ;
2 use std :: path :: Path ;
3
4 // Load model
5 let cfg = ModelConfig :: small ( ) ;
6 let ( encoder , _ms ) = STEEGFormerEncoder :: < B > :: load_from_config (
7 cfg ,
8 Path :: new ( "ST-EEGFormer_small_encoder.safetensors" ) ,
9 device ,
10 ) ? ;
11
12 // Build input: 4 channels × 6 seconds @ 128 Hz
13 let channels = & [ "Fz" , "C3" , "C4" , "Pz" ] ;
14 let signal = vec! [ 0.0f32 ; channels . len ( ) * 768 ] ;
15 let batch = data :: build_batch_named :: < B > ( signal , channels , 768 , & device ) ;
16
17 // Extract embeddings
18 let result = encoder . run_batch ( & batch ) ? ;
19 println! ( "Embedding shape: {:?}" , result . shape ) ; // [512]
Quick Start — Python
1 from safetensors . torch import load_file
2
3 # Load encoder weights
4 state_dict = load_file ( "ST-EEGFormer_small_encoder.safetensors" )
5
6 # Build model and load
7 from models_mae_eeg import mae_vit_small_patch16
8 model = mae_vit_small_patch16 ( )
9 model . load_state_dict ( state_dict , strict = False )
10 model . eval ( )
Architecture
EEG signal (B, C, T) — up to 142 channels, 128 Hz, ≤ 6s
│
▼
┌──────────────────────────────────────┐
│ PatchEmbedEEG │
│ Unfold → 16-sample patches │
│ Linear(16, embed_dim) │
│ → (B, num_patches × C, D) │
└──────────────────────────────────────┘
│
+ Sinusoidal Temporal PE (fixed)
+ Learned Channel Embedding (nn.Embedding(145, D))
│
▼
┌──────────────────────────────────────┐
│ [CLS] token prepend │
└──────────────────────────────────────┘
│
▼
┌──────────────────────────────────────┐
│ N × Transformer Encoder Block │
│ Pre-norm: LN → MHSA → residual │
│ LN → FFN → residual │
│ (qkv_bias=True, GELU activation) │
└──────────────────────────────────────┘
│
▼
┌──────────────────────────────────────┐
│ LayerNorm → CLS token │
│ → (B, embed_dim) embedding │
└──────────────────────────────────────┘MAE Pre-training (decoder, included in *_mae.safetensors)
Encoder output (25% of tokens)
│
▼
Linear(embed_dim → decoder_dim)
+ Insert mask tokens at masked positions
+ Decoder temporal/channel PE
│
▼
M × Decoder Transformer Blocks
│
▼
Linear(decoder_dim → patch_size)
→ Reconstructed EEG patchesNumerical Parity (Rust vs Python)
Verified at every stage against the official PyTorch implementation:
Stage RMSE Pearson r Patch embedding 0.000000 1.000000 Channel embedding 0.000000 1.000000 Temporal encoding 0.000000 1.000000 After positional encoding 0.000000 1.000000 After transformer block 0 0.000004 1.000000 Full encoder (8 blocks) 0.000001 1.000000
Benchmarks
Platform: Apple M4 Pro, 64 GB RAM, macOS (arm64)
Inference Latency — ST-EEGFormer-Small (22ch × 768 samples)
Backend Mean Min Rust CPU (NdArray + Accelerate) 608.4 ms 601.4 ms Python CPU (PyTorch 2.6) 78.1 ms 77.2 ms Rust GPU (Burn wgpu + Metal) 38.1 ms 7.9 ms Python MPS (PyTorch + Metal) 19.2 ms 19.0 ms
Channel Scaling (T=768)
Channels Rust CPU Python CPU Rust GPU Python MPS 4 75.5 ms 21.8 ms 11.5 ms 4.0 ms 22 596.0 ms 77.9 ms 32.7 ms 19.3 ms 64 3853.2 ms 301.9 ms 119.4 ms 90.1 ms
Weight Key Format
Encoder keys
patch_embed.proj.weight [embed_dim, 16]
patch_embed.proj.bias [embed_dim]
cls_token [1, 1, embed_dim]
enc_channel_emd.channel_transformation.weight [145, embed_dim]
enc_temporal_emd.pe [1, 512, embed_dim]
blocks.{i}.norm1.weight [embed_dim]
blocks.{i}.norm1.bias [embed_dim]
blocks.{i}.attn.qkv.weight [3*embed_dim, embed_dim]
blocks.{i}.attn.qkv.bias [3*embed_dim]
blocks.{i}.attn.proj.weight [embed_dim, embed_dim]
blocks.{i}.attn.proj.bias [embed_dim]
blocks.{i}.norm2.weight [embed_dim]
blocks.{i}.norm2.bias [embed_dim]
blocks.{i}.mlp.fc1.weight [4*embed_dim, embed_dim]
blocks.{i}.mlp.fc1.bias [4*embed_dim]
blocks.{i}.mlp.fc2.weight [embed_dim, 4*embed_dim]
blocks.{i}.mlp.fc2.bias [embed_dim]
norm.weight [embed_dim]
norm.bias [embed_dim]Decoder keys (MAE only)
decoder_embed.weight [dec_dim, embed_dim]
decoder_embed.bias [dec_dim]
mask_token [1, 1, dec_dim]
dec_channel_emd.channel_transformation.weight [145, dec_dim]
dec_temporal_emd.pe [1, 512, dec_dim]
decoder_blocks.{i}.* (same structure as encoder)
decoder_norm.weight [dec_dim]
decoder_norm.bias [dec_dim]
decoder_pred.weight [16, dec_dim]
decoder_pred.bias [16]Conversion
These weights were converted from the official .pth files:
1 import torch
2 from safetensors . torch import save_file
3
4 ckpt = torch . load ( "checkpoint.pth" , map_location = "cpu" , weights_only = False )
5 state_dict = ckpt [ "model" ]
6
7 # Encoder only
8 encoder = { k : v . float ( ) . contiguous ( ) for k , v in state_dict . items ( )
9 if any ( k . startswith ( p ) for p in
10 [ "patch_embed." , "cls_token" , "enc_" , "blocks." , "norm." ] ) }
11 save_file ( encoder , "encoder.safetensors" )
Or use the included conversion script:
python scripts/convert_to_safetensors.py --allCitation
1 @inproceedings{yang2026_steegformer,
2 title={Are {EEG} Foundation Models Worth It? Comparative Evaluation
3 with Traditional Decoders in Diverse {BCI} Tasks},
4 author={Liuyin Yang and Qiang Sun and Ang Li and Marc M. Van Hulle},
5 booktitle={The Fourteenth International Conference on Learning Representations},
6 year={2026},
7 url={https://openreview.net/forum?id=5Xwm8e6vbh}
8 }
License
MIT — same as the original ST-EEGFormer release.
Links
