lumamba 0.0.1

# lumamba-rs

[![crates.io](https://img.shields.io/crates/v/lumamba.svg)](https://crates.io/crates/lumamba)
[![docs.rs](https://docs.rs/lumamba/badge.svg)](https://docs.rs/lumamba)
[![license: GPL-3.0](https://img.shields.io/crates/l/lumamba.svg)](./LICENSE)

**LuMamba** — EEG foundation-model inference in Rust, on the [RLX](https://github.com/MIT-RLX/rlx) runtime.

Pure-Rust port of [`PulpBio/LuMamba`](https://huggingface.co/PulpBio/LuMamba)
([BioFoundation](https://github.com/pulp-bio/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`](https://crates.io/crates/luna-rs); the LUNA front-end and the
reconstruction head are shared, the temporal encoder is the part that differs.

## Install

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

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

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

```sh
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)

| | value |
|---|---|
| params | 4.19 M |
| patch_size | 40 |
| queries `Q` / embed `E` | 6 / 64 → `d_model` 384 |
| Mamba blocks | 2 (bidirectional, `add`) |
| `d_inner` / `d_state` / `d_conv` / `dt_rank` | 768 / 16 / 4 / 24 |
| cross-attn heads | 2 |

## Quick start

```sh
# 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:

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

## Backends

Latest RLX (0.2.10). Pick a backend feature:

| feature | backend |
|---------|---------|
| `rlx-cpu` (default) | multi-threaded CPU |
| `rlx-metal` | Apple Metal |
| `rlx-mlx` | Apple MLX |
| `rlx-cuda` / `rlx-rocm` / `rlx-gpu` / `rlx-tpu` | platform GPU/accelerator |
| `rlx-blas-accelerate`, `rlx-apple-silicon` | BLAS / aggregate |

```sh
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.

```sh
# 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

| feature | default | enables |
|---|:--:|---|
| `rlx-cpu` | ✅ | multi-threaded CPU backend |
| `rlx-metal` / `rlx-mlx` | | Apple Metal / MLX backends |
| `rlx-gpu` | | portable GPU via wgpu |
| `rlx-cuda` / `rlx-rocm` / `rlx-tpu` | | NVIDIA / AMD / Google TPU backends |
| `rlx-blas-accelerate` / `rlx-apple-silicon` | | BLAS / Apple aggregate |
| `cpu` `metal` `mlx` `gpu` `cuda` `apple-silicon` | | short aliases for the above |
| `edf` | | EDF/FIF input loading (`exg` + `exg-luna`) |
| `hf-download` | | `hf://` weight download + `download_weights` bin |
| `validation` | | per-stage encoder debug graphs + `validate_encoder` bin |

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

| backend | host-prepare | graph-run | end-to-end (warm) |
|---|---:|---:|---:|
| Metal | ~20 ms | **~16 ms** | **~36 ms** |
| MLX | ~20 ms | ~20 ms | ~39 ms |
| CPU | ~20 ms | ~45 ms | ~64 ms |
| wgpu | ~20 ms | ~150 ms | ~170 ms |

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:

```sh
# 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

```rust
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:

| stage | max abs error |
|-------|---------------|
| inner Mamba (in_proj, conv, `selective_scan`, D-skip, gate, out_proj) | ~3e-6 |
| bidirectional wrapper (`fwd + flip∘rev∘flip`) | ~1e-6 |
| full 2-block encoder, S=128 | **1.5e-5** (rel 1.2e-6) |

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

```sh
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.)

```sh
# 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`](LICENSE).

The LuMamba model **weights** are **CC-BY-ND-4.0** from
[`PulpBio/LuMamba`](https://huggingface.co/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

```bibtex
@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},
}
```