lumamba-rs

LuMamba — EEG foundation-model inference in Rust, on the RLX runtime.

Pure-Rust port of PulpBio/LuMamba (BioFoundation). LuMamba keeps LUNA's topology-invariant front-end — variable-channel EEG is compressed into a fixed set of learned queries by cross-attention — and replaces LUNA's rotary Transformer encoder with a stack of bidirectional Mamba (FEMBA) blocks that run in linear time over the patch sequence.

This crate is a sibling of luna-rs; the LUNA front-end and the reconstruction head are shared, the temporal encoder is the part that differs.

Install

[dependencies]
lumamba = "0.0.1"        # RLX on CPU by default

Use the library via LuMambaEncoder (see API). For a GPU/accelerator backend, disable defaults and pick a backend feature, e.g. lumamba = { version = "0.0.1", default-features = false, features = ["metal"] }.

To get the CLI tools (infer, eval, safetensors_info, download_weights) instead:

cargo install lumamba                    # CPU
cargo install lumamba --features metal   # or another backend

Architecture

EEG (B, C, T)
  ├─ PatchEmbed (3× Conv2d + GroupNorm + GELU)        ┐  host (CPU) prepare
  └─ FreqEmbed  (FFT magnitude/phase → MLP)           ├─  (FFT, conv, MLP)
             → + NeRF channel-location embedding      ┘
  → CrossAttention channel unification (C → Q queries)  ┐
  → reshape [B, S, Q·E]                                 │  RLX compute graph
  → N × { LayerNorm → BiMamba(fwd ⊕ flip∘rev) → +res }  │
  → reconstruction head  →  (B, C, T)                   ┘

Each Mamba block is the stock mamba_ssm.Mamba recurrence, run in forward and reversed streams and combined by addition:

in_proj → [x | z]
x → causal depthwise conv1d → SiLU
x → x_proj → (dt, B, C);  delta = softplus(dt_proj(dt) + bias)
y = selective_scan(x, delta, A = -exp(A_log), B, C)   # RLX Op::SelectiveScan
y = y + D ⊙ x
y = y * SiLU(z)
out = out_proj(y)

The temporal recurrence uses RLX's first-class selective_scan op, the depthwise conv a grouped conv2d, and everything else the standard graph ops — so the whole model compiles to one CompiledGraph per input shape and runs on any RLX backend.

LuMamba-tiny (the released checkpoints)

Quick start

# default build = RLX on CPU
cargo build --release

# fetch a checkpoint (needs the hf-download feature)
cargo run --release --features hf-download --bin download_weights

# reconstruction on a synthetic epoch
cargo run --release --bin infer -- \
  --config config.json \
  --weights <path>/LuMamba_LeJEPA_reconstruction_128slices.safetensors \
  --output out.safetensors -v

# emit the encoder latent (foundation embedding [S, Q·E]) instead
cargo run --release --bin infer -- \
  --config config.json --weights <path>.safetensors \
  --output latent.safetensors --encode -v

Inspect a checkpoint's tensors:

cargo run --release --bin safetensors_info -- <path>.safetensors

Backends

Latest RLX (0.2.10). Pick a backend feature:

cargo run --release --no-default-features --features rlx-mlx --bin infer -- --device mlx …

EDF/FIF input loading is behind --features edf (pulls exg + exg-luna); without it, infer runs on a synthetic epoch.

HuggingFace weights (--features hf-download)

With the hf-download feature, --weights / --config (on infer and eval) accept an hf://<owner>/<repo>/<file> spec that is fetched from the Hub and cached under ~/.cache/huggingface — anything else is a local path. Without the feature, local paths still work and an hf:// spec errors with a message pointing at the flag.

# auto-download + run in one step
cargo run --release --features hf-download --bin infer -- \
  --config config.json \
  --weights hf://PulpBio/LuMamba/LuMamba_LeJEPA_reconstruction_128slices.safetensors \
  --output out.safetensors -v

# or just fetch (prints the cached path); --list shows the released checkpoints
cargo run --release --features hf-download --bin download_weights -- --list
cargo run --release --features hf-download --bin download_weights -- \
  --file LuMamba_ReconstructionOnly.safetensors

Cargo features

Performance

LuMamba-tiny is small (4.2M params), so for a single epoch the GPU cold-start (one-time graph compilation) dominates — which made GPUs look slower than CPU at first glance. The compiled graph is cached per input shape, so that cost is paid once and amortizes over a recording's epochs. Measured warm steady-state per epoch (Apple M-series, 128-patch epoch):

So Metal and MLX beat the CPU once warm. Two things make this work:

• Parallelized host token-prep (rayon): the patch-embed conv + FFT + channel-MLP that run on the host every epoch dropped from ~50 ms → ~20 ms, bit-identical to the serial version. This is backend-independent, so it lifts every backend and lets the fast GPU graph dominate.
• Per-shape compiled-graph cache so warm epochs skip recompilation.

Tools to measure it yourself:

# cold (first) vs warm (cached-graph) per-epoch time
cargo run --release --features rlx-mlx --bin infer -- --device mlx \
  --config config.json --weights CKPT.safetensors --encode --repeat 20

# split host-prepare vs graph-run
LUMAMBA_TIMING=1 cargo run --release --bin infer -- … --repeat 6

Backend guidance: on Apple Silicon prefer metal or mlx; the portable wgpu path works and is numerically correct but has higher per-dispatch overhead for a model this small. For one-shot CLI use, the cold compile is unavoidable; for serving many epochs it's a one-time cost.

API

use lumamba::LuMambaEncoder;
use std::path::Path;

let (mut enc, _ms) = LuMambaEncoder::load(
    Path::new("config.json"),
    Path::new("LuMamba_LeJEPA_reconstruction_128slices.safetensors"),
    rlx::Device::Cpu,
)?;

// reconstruction → [C, T]
let out = enc.run_epoch(&signal, &chan_pos, Some(&channel_idx), n_ch, n_t)?;

// foundation embedding → [S, Q·E]
let (latent, shape) = enc.encode(&signal, &chan_pos, Some(&channel_idx), n_ch, n_t)?;

Validation

The bidirectional-Mamba encoder is numerically validated against an independent PyTorch reference (scripts/lumamba_ref.py) that loads the raw checkpoint and implements the mamba_ssm selective_scan_ref recurrence directly — no code shared with the RLX path. Feeding a byte-identical latent into both (validate_encoder, built with --features validation) and comparing:

i.e. agreement at the f32 numerical-noise floor. Reproduce with:

python scripts/lumamba_ref.py gen CKPT.safetensors 128
cargo run --release --features validation --bin validate_encoder -- \
  --config config.json --weights CKPT.safetensors \
  --hin /tmp/h_in.safetensors --out /tmp/h_out.safetensors
python scripts/lumamba_ref.py check CKPT.safetensors   # also: debug / debug2 for per-stage

Downstream evaluation harness

The eval binary scores a fine-tuned LuMamba classifier on a labeled EEG eval set and compares against the paper's published numbers. (The released HF checkpoints are pretrained encoders, num_classes: 0 — they have no classifier head, so eval reports that and asks for a fine-tuned checkpoint.)

# 1. build an eval set from preprocessed signals + labels (see the helper)
python scripts/make_eval_set.py /tmp/demo_eval        # synthetic demo
# (real use: call scripts/make_eval_set.build(...) with your epochs/labels)

# 2. score a fine-tuned checkpoint
cargo run --release --bin eval -- \
  --config finetune_config.json \
  --weights LuMamba_TUAB_finetuned.safetensors \
  --manifest tuab_eval/manifest.json

task        : TUAB
scored units: 276 (recording-level (mean prob))
balanced_accuracy = 0.8104
expected          = 0.8099  (published)
Δ vs paper        = +0.0005

What it does: per-epoch classification → optional recording-level mean-prob aggregation → metric. Supported metrics match the paper: balanced_accuracy (TUAB), auroc (binary + macro-OvR), aupr (APAVA). Classifier heads are auto-detected from the checkpoint keys: LUNA-style (learned-agg cross-attention, run in the RLX graph) and linear (BasicLinearClassifier, mean-pool → Linear → GELU, host-side) are supported; the Mamba classifier head errors with a clear message.

Eval-set format (lumamba::eval): a JSON manifest (task, num_classes, metric, positive_class, aggregate, labels, recording_ids, expected) + a signals.safetensors (signals [N,C,T], chan_pos [C,3], optional channel_idx [C]). Build it with scripts/make_eval_set.py.

The metrics (src/metrics.rs) are unit-tested against scikit-learn reference values (AUROC, average-precision/AUPR, balanced accuracy); the scoring + aggregation path is unit-tested in src/eval.rs.

Reproducibility note. The paper's TUAB/APAVA/TDBrain/TUSL/TUAR numbers come from fine-tuning a head on each labeled dataset (credentialed clinical EEG) and scoring its test split. The fine-tuned checkpoints are not released and the datasets are access-controlled, so the numbers can't be reproduced from the public artifacts alone. This harness is the ready-to-run scorer for when you have a fine-tuned checkpoint + the matching test split. The one published quantity reproducible from the released weights — the 4.1M parameter count — matches exactly (encoder params = 4,100,473).

Project layout

src/
  lib.rs            crate root + re-exports
  config.rs         ModelConfig / DataConfig
  channel_*.rs      channel vocabulary + montage positions (montages/*.elc)
  metrics.rs        balanced-accuracy / AUROC / AUPR (sklearn-matched)
  eval.rs           eval-set loader, aggregation, scoring
  hf.rs             hf:// weight resolution
  rlx/
    graph.rs        the LuMamba compute graph (front-end → Mamba → heads)
    weights.rs      safetensors → parameter map
    prepare.rs      host-side patch / FFT / channel embedding (rayon)
    encoder.rs      LuMambaEncoder: compile-cache + run / encode / classify
    io.rs           epoch I/O (EDF/FIF behind `edf`)
  bin/              infer, eval, safetensors_info, download_weights, validate_encoder
scripts/            lumamba_ref.py (parity oracle), make_eval_set.py
tests/              smoke.rs

Parity notes

The architecture is mapped exactly against ground truth:

• Weights: the safetensors state-dict keys/shapes drive weights.rs (mamba_blocks.{i}.{mamba_fwd,mamba_rev}.*, norm_layers.{i}, the LUNA front-end and decoder head). A = -exp(A_log) is materialised at load time; the depthwise conv weight is reshaped to grouped-conv2d NCHW.
• SSM op: RLX selective_scan(x, delta, A, B, C) takes [b,s,h] layout and applies neither the softplus, the D skip, the z gate, nor A negation internally — all four are done explicitly here, matching mamba_ssm.
• Bidirectional combine: mamba_fwd(x) + flip(mamba_rev(flip(x))) (BioFoundation MambaWrapper, bidirectional_strategy="add"); the sequence flip uses RLX Op::Reverse along the patch axis (batch-general).
• Front-end: ported verbatim from luna-rs (Python parity ≈ 2e-6). The cross-attention temparature parameter is a fixed, unused 1.0 in the reference and is correctly ignored.
• softplus is composed as log(1+exp(x)) — exact in f32 across Mamba's dt range (confirmed by the delta stage matching to ~3e-7).

tests/smoke.rs checks graph compilation and finite, correctly-shaped outputs; the numerical parity above is the authoritative check (see Validation).

License

GPL-3.0-only. This crate links the RLX runtime (GPL-3.0-only), so the combined work is GPL-3.0. See LICENSE.

The LuMamba model weights are CC-BY-ND-4.0 from PulpBio/LuMamba — see that repo for their terms (no redistribution of modified weights). The bundled MNE-Python ASA montage .elc files retain their original (BSD-3-Clause) license.

Citation

@article{broustail2026lumamba,
  title  = {LuMamba: Latent Unified Mamba for Electrode Topology-Invariant and
            Efficient EEG Modeling},
  author = {Broustail, Dana{\'e} and Ingolfsson, Thorir Mar and others},
  year   = {2026},
  eprint = {2603.19100},
  archivePrefix = {arXiv},
}