eugenehp/ST-EEGFormer

# 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

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]

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()

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         │
└──────────────────────────────────────┘

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

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_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]

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")

python scripts/convert_to_safetensors.py --all

@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}
}