# sleepfm-rs

Pure-Rust inference for the **SleepFM** (Sleep Foundation Model) polysomnography foundation model, built on [Burn 0.20](https://burn.dev).

SleepFM is a contrastive multi-modal sleep model that processes variable-channel PSG inputs (EEG/EOG, Respiratory, ECG/EKG, EMG) into fixed-size embeddings using a SetTransformer architecture.

**Reference:** [SleepFM: Multi-modal Multi-task Sleep Foundation Model](https://github.com/rthapa84/sleepfm-clinical)

## Architecture

| Component | Description |
|---|---|
| **Tokenizer** | 6-layer 1D CNN (1→4→8→16→32→64→128 channels, stride 2, kernel 5) converts 640-sample patches into 128-dim tokens |
| **Spatial Pooling** | Attention pooling across channels per time step |
| **Positional Encoding** | Sinusoidal encoding + LayerNorm |
| **Transformer Encoder** | 6 pre-norm transformer layers (8 heads, FFN dim 2048) |
| **Temporal Pooling** | Attention pooling across time → single 128-dim embedding |

## Quick start

```bash
# 1. Convert PyTorch weights to safetensors
python3 scripts/convert_to_safetensors.py \
    --input path/to/sleepfm/checkpoints/model_base/best.pt \
    --output data/model_base.safetensors

# 2. Copy config
cp path/to/sleepfm/checkpoints/model_base/config.json data/

# 3. Build and run inference
cargo build --release
cargo run --release --bin infer -- \
    --config data/config.json \
    --weights data/model_base.safetensors
```

## Build options

```bash
# CPU (default — Rayon multi-threading)
cargo build --release

# CPU + Apple Accelerate BLAS (macOS)
cargo build --release --features blas-accelerate

# GPU (wgpu — auto-detects Metal/Vulkan/DX12)
cargo build --release --no-default-features --features wgpu

# GPU (macOS Metal native shaders)
cargo build --release --no-default-features --features metal
```

## Modalities

SleepFM processes 4 PSG signal modalities through the same SetTransformer:

| Modality | Typical channels | Description |
|---|---|---|
| **BAS** | 10 | EEG + EOG (brain/eye) |
| **RESP** | 7 | Respiratory (chest, abdomen, SpO2, airflow, etc.) |
| **EKG** | 2 | Electrocardiogram |
| **EMG** | 4 | Electromyogram (chin, legs) |

Each modality is processed independently by the shared model, producing per-modality embeddings that are L2-normalised and compared via contrastive learning.

## Benchmarks

Benchmarks measured on Apple Silicon (M-series), single-batch inference over 5 minutes of PSG data at 128 Hz (38,400 samples), averaged over 10 runs after 3 warmup iterations.

### Inference time by modality

![Benchmark comparison](figures/benchmark_comparison.png)

| Backend | EKG (2ch) | EMG (4ch) | RESP (7ch) | BAS (10ch) |
|---|---|---|---|---|
| PyTorch CPU | 21 ms | 35 ms | 56 ms | 76 ms |
| Rust + Accelerate | 24 ms | 36 ms | 54 ms | 72 ms |
| Rust (NdArray) | 38 ms | 50 ms | 68 ms | 89 ms |
| **Rust + Metal GPU** | **13 ms** | **13 ms** | **13 ms** | **13 ms** |

The Metal GPU backend runs at a **constant ~13 ms regardless of channel count**, achieving up to **5.8× speedup** over PyTorch CPU on BAS (10 channels).

### Scaling with channel count

![Scaling with channels](figures/scaling_channels.png)

CPU backends scale linearly with channel count (dominated by the per-channel tokenizer CNN). The GPU backend parallelises all channels, resulting in near-flat scaling.

### Relative speed

![Speedup ratio](figures/speedup_ratio.png)

| Backend vs PyTorch CPU | EKG (2ch) | EMG (4ch) | RESP (7ch) | BAS (10ch) |
|---|---|---|---|---|
| Rust + Accelerate | 0.9× | 1.0× | 1.0× | 1.0× |
| Rust NdArray | 0.6× | 0.7× | 0.8× | 0.9× |
| **Rust + Metal GPU** | **1.6×** | **2.8×** | **4.4×** | **5.8×** |

### Embedding visualization

![Embeddings](figures/embeddings_visualization.png)

128-dimensional pooled embeddings for each modality from the same PSG recording, showing distinct modality-specific activation patterns.

### Numerical parity

Layer-by-layer comparison against PyTorch (2 channels × 640 samples):

| Layer | Max abs diff | Cosine similarity |
|---|---|---|
| Tokenizer | 1e-6 | 1.000000 |
| Spatial pooling | 1e-6 | 1.000000 |
| PosEnc + LayerNorm | 1e-6 | 1.000000 |
| Transformer encoder | 3e-6 | 1.000000 |
| Temporal pooling | 2e-6 | 1.000000 |

Full model on random inputs (38,400 samples):

| Modality | Cosine sim | Max abs diff |
|---|---|---|
| BAS (10ch) | 1.000000 | 0.0022 |
| RESP (7ch) | 0.999999 | 0.0024 |
| EKG (2ch) | 0.999999 | 0.0041 |
| EMG (4ch) | 0.999999 | 0.0044 |

> Residual differences are expected f32 accumulation noise over 60 patches × 6 transformer layers.

### Reproduce benchmarks

```bash
# Rust CPU (NdArray)
cargo run --release --example benchmark -- \
    --config data/config.json --weights data/model_base.safetensors

# Rust CPU + Apple Accelerate (macOS)
cargo run --release --features blas-accelerate --example benchmark -- \
    --config data/config.json --weights data/model_base.safetensors \
    --output data/benchmark_rust_accelerate.json

# Rust GPU (Metal — macOS)
cargo run --release --no-default-features --features metal --example benchmark -- \
    --config data/config.json --weights data/model_base.safetensors \
    --output data/benchmark_rust_metal.json

# PyTorch CPU baseline
python3 scripts/benchmark.py

# Generate charts
python3 scripts/plot_benchmarks.py
```

## License

This Rust implementation (sleepfm-rs) is licensed under the **GNU General Public License v3.0 (GPLv3)**.

### ⚠️ Model Weights — Non-Commercial Use Only

The pretrained SleepFM model weights are licensed by their original authors under the
[**Creative Commons Attribution-NonCommercial 4.0 International License (CC BY-NC 4.0)**](https://creativecommons.org/licenses/by-nc/4.0/).

> Copyright (c) 2025 Rahul Thapa

This means the weights **may not be used for commercial purposes**. You are free to share and adapt the weights for non-commercial use, provided you give appropriate credit and indicate any changes made.

If you use the SleepFM model or weights, please cite the original work:
- [SleepFM: Multi-modal Multi-task Sleep Foundation Model](https://github.com/rthapa84/sleepfm-clinical)
