luna-rs

LUNA (Latent Unified Network Architecture) — EEG foundation model inference in Rust.

Pure-Rust port of LUNA (BioFoundation). Two inference engines are available:

LUNA compresses variable-channel EEG into fixed learned queries via cross-attention, then runs a rotary Transformer encoder.

Weights: thorir/LUNA on HuggingFace (auto-resolved from cache when present).

Quick start (RLX)

# Default build: RLX on CPU
cargo build --release

# Inference on synthetic EEG (set paths or use HF cache)
export LUNA_WEIGHTS=~/.cache/huggingface/hub/models--thorir--LUNA/snapshots/<hash>/LUNA_base.safetensors
export LUNA_CONFIG=tests/vectors/config.json

cargo run --release --bin infer -- \
  --weights "$LUNA_WEIGHTS" --config "$LUNA_CONFIG" --output out.safetensors -v

Apple Silicon GPU:

cargo run --release --no-default-features \
  --features rlx,rlx-mlx --bin infer -- \
  --device mlx --weights "$LUNA_WEIGHTS" --config "$LUNA_CONFIG" --output out.safetensors

Build

# RLX CPU (default)
cargo build --release

# RLX + Apple backends
cargo build --release --no-default-features --features rlx,rlx-cpu,rlx-metal,rlx-mlx

# Burn reference (NdArray CPU)
cargo build --release --no-default-features --features burn,ndarray

# Burn GPU (wgpu / Metal / Vulkan)
cargo build --release --no-default-features --features burn,wgpu
cargo build --release --no-default-features --features burn,metal   # macOS MSL
cargo build --release --no-default-features --features burn,vulkan # Linux SPIR-V

# Both engines (parity tests, backend_compare)
cargo build --release --no-default-features \
  --features burn,rlx,ndarray,rlx-cpu,rlx-metal,rlx-mlx

Benchmarks

Burn vs RLX (all compiled backends, optional parity vs Burn CPU):

cargo run --example backend_compare --release \
  --no-default-features \
  --features burn,rlx,ndarray,rlx-cpu,rlx-metal,rlx-mlx -- --parity

Burn-only latency across model variants (base / large / huge):

cargo run --example benchmark --release --features burn,hf-download
# or: ./bench.sh base,large,huge

Historical multi-platform charts live under figures/ (from bench.sh).

Architecture

EEG (B, C, T)
    ├─ PatchEmbedNetwork (3-layer CNN)
    └─ FrequencyFeatureEmbedder (FFT + MLP)
              → sum → + NeRF positions + channel MLP
              → CrossAttentionBlock (Q queries × C channels)
              → N × RotaryTransformerBlock
              → Reconstruction head or classification head

Tests

# RLX smoke + parity (needs cached LUNA-Base weights)
cargo test --release --no-default-features \
  --features burn,rlx,ndarray,rlx-cpu \
  --test rlx_graph_compile --test parity_rlx_vs_burn \
  --test rlx_prepare_burn --test rlx_graph_parity

# Burn ↔ Python vectors (export first)
python scripts/export_parity_vectors.py
cargo test --release --no-default-features --features burn,ndarray \
  --test python_parity --test intermediate_parity --test block_parity \
  --test decoder_parity --test f64_parity --test forward_pass

Set LUNA_WEIGHTS or cache thorir/LUNA locally. RLX device override: LUNA_RLX_DEVICE=cpu|metal|mlx.

Shared helpers: tests/common/mod.rs.

Examples (Burn)

Requires --features burn (and usually hf-download):

Use --variant large or huge where supported.

Binaries

API

RLX (default crate root when burn is off):

use luna_rs::{LunaEncoder, init_threads};
use std::path::Path;

let _ = init_threads(None);
let (mut enc, _) = LunaEncoder::load(
    Path::new("config.json"),
    Path::new("LUNA_base.safetensors"),
    rlx::Device::Cpu,
)?;
let out = enc.run_epoch(&signal, &locations, Some(&names), n_ch, n_t)?;

Burn (--features burn):

use luna_rs::{LunaEncoder, build_batch_named, TUEG_CHANNELS};
// LunaEncoder<B>, model::luna::Luna, load_from_csv, …

Project layout

src/
  lib.rs, config.rs, channel_*.rs
  rlx/          # RLX graph, prepare, encoder (default path)
  model/        # Burn modules
  bin/          # infer, download_weights, …
examples/
  backend_compare.rs   # Burn vs RLX benchmark
  reconstruct.rs, …    # Burn demos
tests/
  common/              # Shared weight paths, synthetic data
  *_parity.rs          # Numerical tests
scripts/               # Python vector export

Parity notes

• Burn ↔ Python: RMSE ≈ 2e-6 with exported vectors (mask=None inference).
• RLX ↔ Burn: max_abs ≲ 3e-6 on LUNA-Base, 22 ch × 1280 samples.
• RLX ↔ Python: requires scripts/export_parity_vectors.py (BioFoundation).
• MLX: attention uses [B,H,S,D] layout in the RLX graph (required by MLX SDPA).

Citation

@inproceedings{
  doner2025luna,
  title={{LUNA}: Efficient and Topology-Agnostic Foundation Model for {EEG} Signal Analysis},
  author={Berkay D{\"o}ner and Thorir Mar Ingolfsson and Luca Benini and Yawei Li},
  booktitle={NeurIPS},
  year={2025},
  url={https://openreview.net/forum?id=uazfjnFL0G}
}

License

Apache-2.0