---
license: mit
tags:
  - eeg
  - bci
  - brain-computer-interface
  - foundation-model
  - vit
  - masked-autoencoder
  - mae
  - neuroscience
  - safetensors
  - burn
  - rust
language:
  - en
library_name: steegformer-rs
pipeline_tag: feature-extraction
---

# ST-EEGFormer — Safetensors Weights

Pre-converted [safetensors](https://github.com/huggingface/safetensors) weights for the [ST-EEGFormer](https://github.com/LiuyinYang1101/STEEGFormer) EEG foundation model, ready for use with **[steegformer-rs](https://github.com/eugenehp/steegformer-rs)** (pure-Rust inference on [Burn 0.20](https://burn.dev)) or any framework that supports safetensors.

Weights are converted from the official PyTorch `.pth` checkpoints published at [LiuyinYang1101/STEEGFormer](https://github.com/LiuyinYang1101/STEEGFormer/releases).

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)

| File | Variant | Params | Size | Layers | Heads | embed_dim |
|------|---------|--------|------|--------|-------|-----------|
| [`ST-EEGFormer_small_encoder.safetensors`](ST-EEGFormer_small_encoder.safetensors) | **Small** | 25.6 M | 102 MB | 8 | 8 | 512 |
| [`ST-EEGFormer_base_encoder.safetensors`](ST-EEGFormer_base_encoder.safetensors) | **Base** | 85.6 M | 342 MB | 12 | 12 | 768 |
| [`ST-EEGFormer_large_encoder.safetensors`](ST-EEGFormer_large_encoder.safetensors) | **Large** | 303.0 M | 1,212 MB | 24 | 16 | 1024 |
| [`ST-EEGFormer_largeV2_encoder.safetensors`](ST-EEGFormer_largeV2_encoder.safetensors) | **Large V2** | 303.1 M | 1,212 MB | 24 | 16 | 1024 |

### Full MAE (encoder + decoder, for reconstruction / fine-tuning)

| File | Variant | Params | Size | Decoder dim | Decoder depth |
|------|---------|--------|------|-------------|---------------|
| [`ST-EEGFormer_small_mae.safetensors`](ST-EEGFormer_small_mae.safetensors) | **Small** | 33.1 M | 132 MB | 384 | 4 |
| [`ST-EEGFormer_base_mae.safetensors`](ST-EEGFormer_base_mae.safetensors) | **Base** | 111.5 M | 446 MB | 512 | 8 |
| [`ST-EEGFormer_large_mae.safetensors`](ST-EEGFormer_large_mae.safetensors) | **Large** | 329.1 M | 1,316 MB | 512 | 8 |
| [`ST-EEGFormer_largeV2_mae.safetensors`](ST-EEGFormer_largeV2_mae.safetensors) | **Large V2** | 329.3 M | 1,317 MB | 512 | 8 |

### Config

| File | Description |
|------|-------------|
| [`config.json`](config.json) | Model 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

```bash
# Install
cargo add steegformer-rs

# Download weights
huggingface-cli download eugenehp/ST-EEGFormer \
    ST-EEGFormer_small_encoder.safetensors \
    config.json \
    --local-dir weights/

# Run inference
cargo run --release --bin infer -- \
    --config weights/config.json \
    --weights weights/ST-EEGFormer_small_encoder.safetensors
```

### Library API

```rust
use steegformer_rs::{STEEGFormerEncoder, ModelConfig, data};
use std::path::Path;

// Load model
let cfg = ModelConfig::small();
let (encoder, _ms) = STEEGFormerEncoder::<B>::load_from_config(
    cfg,
    Path::new("ST-EEGFormer_small_encoder.safetensors"),
    device,
)?;

// Build input: 4 channels × 6 seconds @ 128 Hz
let channels = &["Fz", "C3", "C4", "Pz"];
let signal = vec![0.0f32; channels.len() * 768];
let batch = data::build_batch_named::<B>(signal, channels, 768, &device);

// Extract embeddings
let result = encoder.run_batch(&batch)?;
println!("Embedding shape: {:?}", result.shape);  // [512]
```

## Quick Start — Python

```python
from safetensors.torch import load_file

# Load encoder weights
state_dict = load_file("ST-EEGFormer_small_encoder.safetensors")

# Build model and load
from models_mae_eeg import mae_vit_small_patch16
model = mae_vit_small_patch16()
model.load_state_dict(state_dict, strict=False)
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 patches
```

## Numerical 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:

```python
import torch
from safetensors.torch import save_file

ckpt = torch.load("checkpoint.pth", map_location="cpu", weights_only=False)
state_dict = ckpt["model"]

# Encoder only
encoder = {k: v.float().contiguous() for k, v in state_dict.items()
           if any(k.startswith(p) for p in
                  ["patch_embed.", "cls_token", "enc_", "blocks.", "norm."])}
save_file(encoder, "encoder.safetensors")
```

Or use the included conversion script:

```bash
python scripts/convert_to_safetensors.py --all
```

## Citation

```bibtex
@inproceedings{yang2026_steegformer,
  title={Are {EEG} Foundation Models Worth It? Comparative Evaluation
         with Traditional Decoders in Diverse {BCI} Tasks},
  author={Liuyin Yang and Qiang Sun and Ang Li and Marc M. Van Hulle},
  booktitle={The Fourteenth International Conference on Learning Representations},
  year={2026},
  url={https://openreview.net/forum?id=5Xwm8e6vbh}
}
```

## License

MIT — same as the original ST-EEGFormer release.

## Links

| | |
|---|---|
| **Rust crate** | [github.com/eugenehp/steegformer-rs](https://github.com/eugenehp/steegformer-rs) |
| **Original code** | [github.com/LiuyinYang1101/STEEGFormer](https://github.com/LiuyinYang1101/STEEGFormer) |
| **Original weights** | [GitHub Releases](https://github.com/LiuyinYang1101/STEEGFormer/releases) |
| **Paper** | [OpenReview (ICLR 2026)](https://openreview.net/forum?id=5Xwm8e6vbh) |
| **Burn framework** | [burn.dev](https://burn.dev) |