Pre-converted safetensors weights for the NeuroRVQ multi-scale biosignal tokenizer, ready for use with neurorvq-rs (pure-Rust inference on Burn 0.20) or any framework that supports safetensors.

Weights are converted from the official PyTorch .pt checkpoints published at ntinosbarmpas/NeuroRVQ.

Model Files

Tokenizers (encoder → RVQ → decoder)

File	Modality	Params	Size	Embed	Patch	RVQ
`NeuroRVQ_EEG_tokenizer_v1.safetensors`	EEG	76.0 M	304 MB	200	200	8 levels
`NeuroRVQ_ECG_tokenizer_v1.safetensors`	ECG	68.1 M	272 MB	40	40	8 levels
`NeuroRVQ_EMG_tokenizer_v1.safetensors`	EMG	143.6 M	574 MB	200	200	16 levels

Foundation Models (encoder only)

File	Modality	Params	Size	Depth
`NeuroRVQ_EEG_foundation_model_v1.safetensors`	EEG	58.6 M	234 MB	12 blocks
`NeuroRVQ_EMG_foundation_model_v1.safetensors`	EMG	111.2 M	445 MB	12 blocks

Config Flags

File	Description
`flags/NeuroRVQ_EEG_v1.yml`	EEG — 103 channels, patch=200, embed=200
`flags/NeuroRVQ_ECG_v1.yml`	ECG — 15 channels, patch=40, embed=40
`flags/NeuroRVQ_EMG_v1.yml`	EMG — 16 channels, patch=200, embed=200

Quick Start — Rust

bash

# Install
cargo add neurorvq-rs

# Download weights + config
huggingface-cli download eugenehp/NeuroRVQ \
    NeuroRVQ_EEG_tokenizer_v1.safetensors \
    flags/NeuroRVQ_EEG_v1.yml \
    --local-dir weights/

# Run tokenization
cargo run --release --bin infer -- \
    --config weights/flags/NeuroRVQ_EEG_v1.yml \
    --weights weights/NeuroRVQ_EEG_tokenizer_v1.safetensors

Library API

rust

use neurorvq_rs::{NeuroRVQEncoder, Modality, data, channels};
use std::path::Path;

let (model, _ms) = NeuroRVQEncoder::<B>::load_with_modality(
    Path::new("flags/NeuroRVQ_EEG_v1.yml"),
    Path::new("NeuroRVQ_EEG_tokenizer_v1.safetensors"),
    Modality::EEG,
    device,
)?;

// Tokenize → 4 branches × 8 RVQ levels of discrete indices
let tokens = model.tokenize(&batch)?;
for (br, levels) in tokens.branch_tokens.iter().enumerate() {
    for (lv, indices) in levels.iter().enumerate() {
        println!("Branch {} Level {}: {} tokens", br, lv, indices.len());
    }
}

Foundation Model API

rust

use neurorvq_rs::{NeuroRVQFoundationModel, Modality};

let (fm, _ms) = NeuroRVQFoundationModel::<B>::load(
    Path::new("flags/NeuroRVQ_EEG_v1.yml"),
    Path::new("NeuroRVQ_EEG_foundation_model_v1.safetensors"),
    Modality::EEG,
    device,
)?;

let features = fm.encode(&batch)?;       // 4 branch feature vectors
let pooled = fm.encode_pooled(&batch)?;   // Mean-pooled for classification

Quick Start — Python

python

from safetensors.torch import load_file

state_dict = load_file("NeuroRVQ_EEG_tokenizer_v1.safetensors")
model.load_state_dict(state_dict, strict=False)

Architecture

Raw Signal [B, N, T]
    │
    ▼
┌──────────────────────────────────────┐
│  Multi-Scale Temporal Conv           │  4 parallel branches
│  EEG/ECG: k=21,15,9,5               │  modality-specific kernels
│  EMG:     k=51,17,8,5               │
└──────────────────────────────────────┘
    │ ×4 branches
    ▼
┌──────────────────────────────────────┐
│  Transformer Encoder                 │  12 blocks, 10 heads
│  + spatial / temporal pos. embeds    │  shared weights across branches
└──────────────────────────────────────┘
    │ ×4 branches
    ▼
┌──────────────────────────────────────┐
│  Encode Heads                        │  Linear → Tanh → Linear
│  embed_dim → code_dim (128)          │
└──────────────────────────────────────┘
    │ ×4 branches
    ▼
┌──────────────────────────────────────┐
│  Residual Vector Quantization        │  8 levels (EEG/ECG)
│  L2-norm codebook lookup             │  16 levels (EMG)
│  codebook: 8192 × 128               │
└──────────────────────────────────────┘
    │ ×4 branches        ← discrete token indices
    ▼
┌──────────────────────────────────────┐
│  Transformer Decoder                 │  3 blocks
│  per-branch PatchEmbed (1×1 conv)    │
└──────────────────────────────────────┘
    │ concat 4 branches
    ▼
┌──────────────────────────────────────┐
│  Decode Heads                        │  Amplitude (GELU)
│  4×embed_dim → decoder_out_dim       │  Sin/Cos phase (Tanh)
└──────────────────────────────────────┘
    │
    ▼
  Inverse FFT → Reconstructed Signal

Numerical Parity (Rust vs Python)

Verified against the official PyTorch reference implementation:

Layer	Max Abs Error	Notes
Encoder features	< 8 × 10⁻³	12 transformer layers, f32 accumulation
Encode heads	< 2 × 10⁻³	After Tanh squashing
RVQ quantized vectors	≈ 0 ¹	Exact with random-init codebooks
Token indices	99.3% exact ²	Pretrained weights
Decode outputs	< 8 × 10⁻¹ ¹	Dominated by ≤0.7% boundary tokens

¹ Differences stem from the ≤0.7% of tokens near codebook decision boundaries — a natural consequence of f32 arithmetic differences between frameworks.

² With random-init weights: 100% match (all "mismatches" resolve to identical codebook vectors, i.e., ties).

Benchmarks

Platform: Apple M4 Pro, 64 GB RAM, macOS 15 (arm64)

Tokenize Latency — All Backends

Configuration	Modality	PyTorch CPU	Rust NdArray	Rust wgpu (GPU)
EEG 4ch × 64t	EEG	179 ms	661 ms	51 ms
EEG 8ch × 32t	EEG	180 ms	662 ms	60 ms
EEG 16ch × 16t	EEG	180 ms	664 ms	62 ms
EEG 32ch × 8t	EEG	178 ms	664 ms	65 ms
EEG 64ch × 4t	EEG	179 ms	664 ms	68 ms
ECG 4ch × 150t	ECG	272 ms	1881 ms	92 ms
ECG 8ch × 75t	ECG	273 ms	1874 ms	92 ms
ECG 12ch × 50t	ECG	272 ms	1877 ms	93 ms
ECG 15ch × 40t	ECG	272 ms	1878 ms	93 ms
EMG 4ch × 64t	EMG	255 ms	998 ms	90 ms
EMG 8ch × 32t	EMG	255 ms	998 ms	88 ms
EMG 16ch × 16t	EMG	254 ms	1001 ms	90 ms

Tokenize Latency: NdArray vs wgpu vs PyTorch

Encode Latency: NdArray vs wgpu vs PyTorch

Rust — Tokenize Latency by Configuration

Rust — EEG Scaling by Channel Count

Rust — Model Construction Time

Backend Comparison Summary

Comparison	Result
wgpu vs NdArray	wgpu is ~12× faster (GPU acceleration)
wgpu vs PyTorch CPU	wgpu is ~3× faster for EEG/EMG/ECG
NdArray vs PyTorch CPU	PyTorch is ~3.7× faster (optimized BLAS)

Key Observations

wgpu (GPU) is the fastest backend — 51–93 ms across all configurations
PyTorch CPU uses Apple Accelerate/AMX BLAS and fused operators, making it faster than Rust NdArray on CPU
Latency scales with total patch count, not the channel/time decomposition — EEG (256 patches) < EMG (256 patches, 16 RVQ) < ECG (600 patches)
Construction time is ~13 ms (warm) / ~54 ms (cold start for EMG with larger kernels)
Standard deviation < 1% — highly stable inference latency

Why Rust?

	Python + PyTorch	Rust + Burn
Dependencies	pip, torch, numpy, einops, ...	Zero (single static binary)
GPU support	CUDA, MPS	wgpu (Metal, Vulkan, WebGPU)
Deployment	Interpreter + venv	Single binary, WASM, embedded
Memory	GC pauses	Deterministic, no GC
Latency (GPU)	—	51–93 ms (wgpu Metal)

Conversion

These weights were converted from the official .pt files:

python

import torch
from safetensors.torch import save_file

state_dict = torch.load("model.pt", map_location="cpu")
converted = {k: v.float().contiguous() for k, v in state_dict.items()
             if isinstance(v, torch.Tensor) and v.is_floating_point()}
save_file(converted, "model.safetensors")

Or use the included script:

bash

python scripts/convert_pt_to_safetensors.py \
    --input NeuroRVQ_EEG_tokenizer_v1.pt \
    --output NeuroRVQ_EEG_tokenizer_v1.safetensors

Citation

bibtex

@article{barmpas2024neurorvq,
  title={NeuroRVQ: Joint Neurophysiological Multi-Scale Temporal Tokenization and Reconstruction},
  author={Barmpas, Konstantinos and others},
  year={2024}
}

License

Apache-2.0 — same as the original NeuroRVQ release.

Links


Rust crate	github.com/eugenehp/neurorvq-rs
Original weights	huggingface.co/ntinosbarmpas/NeuroRVQ
Paper / Code	github.com/KonstantinosBarmpas/NeuroRVQ
Burn framework	burn.dev