model_card.md

# CortexFM — Detailed Model Card

This document complements the top-level `README.md` with reproducibility-grade detail: full architecture specification, training hyperparameters, complete preprocessing pipeline, evaluation protocol, and an extended limitations discussion.

For a quick-start summary, see `README.md`. For licence terms, see `LICENSE`. For upload / download instructions, see `upload_instructions.md`.

---

## 1. Architecture details

### 1.1 End-to-end forward pass

```
spike_counts (B, T, N=64)          emg_envelope (B, T, M=16)
    |                                       |
    v                                       v
SpikeTokenizer                        EMGTokenizer
  per-unit W_unit (N x d)               per-muscle e_muscle (M x d)
  + log(1 + alpha * c)                  + scalar->vector MLP f_val
  + temporal pos. embed                 + temporal pos. embed
  -> (B, T, d)                          -> (B, T*M, d)
    |                                       |
    +---- modality flag g_0 ----+---- modality flag g_1 ----+
    |                                       |
    v                                       v
              concat along sequence (B, T + T*M, d)  =  (B, 1088, 192)
                                       |
                                       v
                          CortexFMBackbone
                  PreNorm Transformer encoder, 10 layers
                  6 heads, d_model = 192, FFN dim = 4 * d = 768
                  GELU activation, dropout = 0.1
                  SDPA kernel: [FLASH_ATTENTION, EFFICIENT_ATTENTION]
                  Final LayerNorm
                                       |
                            split sequence back
                                       |
        +-------------------------------+--------------------------+
        |                                                          |
        v                                                          v
spike hidden (B, T, d)                                emg hidden (B, T*M, d)
        |                                                          |
        v                                                          v
SpikeReconHead                                          EMGReconHead
  LN -> Linear -> GELU -> Linear                          LN -> Linear -> GELU -> Linear
  -> log_rate (B, T, N)                                   -> emg_pred (B, T*M)
        |                                                          |
        +--------- ContrastiveProjector (d -> d_p=128, L2 normalize) ---+
                                                                       |
                                                                       v
                                                          q, k in R^{B x d_p}
```

### 1.2 Per-component specifications

| Component | Class (`cortex_fm.*`) | Parameter count | Notes |
|---|---|---|---|
| SpikeTokenizer | `models.tokenizer.SpikeTokenizer` | 64 × 192 + temporal pos. embed = 12,288 + (64 × 192) | `count_scale = α` learned scalar; `log(1 + α · count)` activation |
| EMGTokenizer | `models.tokenizer.EMGTokenizer` | 16 × 192 + temporal pos. embed + value MLP (2-layer GELU) | per-muscle per-bin tokens, muscle-major order |
| Backbone | `models.CortexFMBackbone` | 4,449,024 | PreNorm; SDPA FLASH/EFFICIENT |
| SpikeReconHead | `models.SpikeReconHead` | LN + Linear(192→768) + GELU + Linear(768→64) | outputs `log_rate` per unit |
| EMGReconHead | `models.EMGReconHead` | LN + Linear(192→768) + GELU + Linear(768→1) | 37,633 params total (≈ 0.75 % of model) |
| ContrastiveProjector | `models.ContrastiveProjector` | 2-layer MLP, d = 192 → 128 | L2 normalized output |
| modality_embed | `nn.Embedding(2, 192)` | 384 | spike flag g₀ and EMG flag g₁ |
| **Total** | | **5,044,994** | |

### 1.3 Attention kernel choice

The backbone wraps `nn.TransformerEncoder` inside

```python
with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]):
    h = self.encoder(x, src_key_padding_mask=...)
```

so PyTorch 2.10 dispatches to FLASH attention when sequence length and dtype permit, falling back to the EFFICIENT kernel otherwise. On the RTX 5080 (sm_120), this gives ~2 × token throughput over the standard kernel and ~860 K tokens/s at the training context of T = 64 bins (sequence length 1,088).

### 1.4 Spike tokenizer math

$$
\mathbf{s}_t = \mathbf{W}_{\text{unit}}^\top \cdot \log(1 + \alpha \cdot \mathbf{c}_t) + \mathbf{p}_t^{\text{spike}}
$$

where $\mathbf{W}_{\text{unit}} \in \mathbb{R}^{N \times d}$ is the per-unit learned embedding, $\alpha$ is a learned global scale (`count_scale`), $\mathbf{c}_t \in \mathbb{N}^N$ is the per-unit spike count vector at bin $t$, and $\mathbf{p}_t^{\text{spike}} \in \mathbb{R}^d$ is the temporal positional embedding. Each unit retains an independent embedding direction — this is the per-unit identity preservation contrasted with NDT-3's 32-unit patch tokenization.

### 1.5 EMG tokenizer math

$$
\mathbf{m}_{t,i} = \mathbf{e}_{\text{muscle}}(i) + \mathbf{p}_t^{\text{emg}} + f_{\text{val}}(E_{t,i})
$$

with $f_{\text{val}}(x) = \mathbf{W}_2 \cdot \text{GELU}(\mathbf{W}_1 x + \mathbf{b}_1) + \mathbf{b}_2$. Each (time, muscle) pair becomes a separate token; the sequence length is $T \times 16$ for EMG.

### 1.6 Fusion

$$
\mathbf{X} = \text{Concat}(\mathbf{S} + \mathbf{g}_0, \mathbf{M} + \mathbf{g}_1) \in \mathbb{R}^{L \times d}, \quad L = T + T \cdot 16 = T \cdot 17
$$

At $T = 64$ this yields $L = 1088$, which aligns with FLASH attention's most efficient regime (sequence length ≈ 1024).

---

## 2. Hyperparameters

### 2.1 Pretraining hyperparameters (`pretrain_v1`)

| Group | Hyperparameter | Value |
|---|---|---|
| Model | d_model | 192 |
| Model | n_heads | 6 |
| Model | n_layers | 10 |
| Model | ffn_mult | 4 (FFN dim = 768) |
| Model | dropout | 0.1 |
| Tokenizer | spike `n_units` | 64 |
| Tokenizer | emg `n_muscles` | 16 |
| Tokenizer | `max_T` (positional embed size) | 1024 |
| Contrastive | projection dim `d_proj` | 128 |
| Data | `context_T` | 64 bins (1.28 s) |
| Data | `bin_size_s` | 0.02 |
| Training | batch_size | 8 |
| Training | max_epochs | 50 (early-best at 28) |
| Training | optimizer | AdamW |
| Training | learning_rate | 3 × 10⁻⁴ |
| Training | weight_decay | 0.01 |
| Training | warmup_steps | 500 |
| Training | LR schedule | cosine decay after warmup |
| Training | precision | bf16-mixed |
| Training | gradient_clip | 1.0 |
| Loss | $w_{\text{spike}}$ | 1.0 |
| Loss | $w_{\text{emg}}$ | 1.0 |
| Loss | $w_{\text{contrastive}}$ | 0.5 |
| Loss | InfoNCE temperature τ | 0.1 |
| Masking | spike mask ratio | 0.50 |
| Masking | emg mask ratio | 0.50 |
| Attention | SDPA backends | `FLASH_ATTENTION`, `EFFICIENT_ATTENTION` |

### 2.2 Loss formulas

$$
\mathcal{L}_{\text{spike}} = \frac{1}{TN} \sum_{t,n} \left[\exp(\log \hat{\lambda}_{t,n}) - c_{t,n} \log \hat{\lambda}_{t,n} \right]
$$

$$
\mathcal{L}_{\text{emg}} = \frac{1}{T \cdot 16} \sum_{t,i} (\hat{E}_{t,i} - E_{t,i})^2
$$

$$
\mathcal{L}_{\text{cont}} = \frac{1}{2B} \sum_{b=1}^{B} \left[ \ell_{\text{CE}}(\mathbf{q}_b \mathbf{K}^\top / \tau, b) + \ell_{\text{CE}}(\mathbf{k}_b \mathbf{Q}^\top / \tau, b) \right]
$$

InfoNCE similarity logits and softmax are computed in FP32 inside the BF16 autocast scope to avoid dynamic-range instability.

---

## 3. Training data details

### 3.1 Source

DANDI Archive Dandiset **000941** (Rouse & Schieber 2018, Univ. of Kansas) — paired M1 single-unit spikes and intramuscular EMG recorded from MonkeyL performing an 8-direction × 5-grasp = 40-cue center-out reach-grasp task.

- 64 single units recorded from left M1 via Utah array.
- 16 muscles instrumented with intramuscular fine-wire electrodes.
- 16-muscle FALCON official order: `APL, BCPs, DLTa, DLTp, ECRB, ECU, EDC, FCR, FCU, FDI, FDPr, FDPu, Hypoth, PECmaj, TCPlat, Thenar`.
- License: CC-BY-4.0.

### 3.2 Split (FALCON M1 official)

| Split | Sessions | Total duration | Use in CortexFM |
|---|---|---|---|
| Held-in calibration | 4 (20120924, 20120926, 20120927, 20120928) | **3 h 38 min** | Pretraining |
| Held-in minival | 4 | 1 min 4 s | Main FALCON M1 evaluation (Chapter 6) |
| Held-out calibration | 3 (20121004, 20121017, 20121024) | 4 min 29 s | OOD session-1 adaptation (Chapter 7) |
| **Total** | **11** | **3 h 48 min** | |

### 3.3 NWB → Zarr conversion

Per-session sizes after the FALCON preprocessing chain and Blosc-zstd Zarr compression:

| Split | Sessions | NWB size | Zarr size | Compression |
|---|---|---|---|---|
| Held-in calibration | 4 | 300.2 MB | 51.3 MB | 17.1 % |
| Held-in minival | 4 | 2.5 MB | 0.3 MB | 11.0 % |
| Held-out calibration | 3 | 9.1 MB | 1.1 MB | 12.0 % |
| **Total** | **11** | **311.8 MB** | **52.6 MB** | **16.9 %** |

### 3.4 Preprocessing pipeline (FALCON-aligned)

EMG raw signal $x_{\text{raw}}(t) \in \mathbb{R}^{16}$ at 1 kHz is processed through eight stages reproduced from the FALCON official `stability-benchmark` repository:

1. Notch filter at 60, 180, 200, 300, 400 Hz (width 2 Hz).
2. 4th-order Butterworth high-pass, cutoff 65 Hz.
3. Rectification (absolute value).
4. 99 % quantile clipping (outlier rejection).
5. 95 % quantile per-session normalization.
6. Polyphase resampling 1 kHz → 50 Hz (20 ms bin), `scipy.signal.resample_poly`.
7. Re-rectification (removes residual ringing).
8. 10 Hz low-pass Butterworth filter → final envelope.

Spike counts are computed on the same 20 ms bin grid:

$$
c_{t,n} = |\{k : t \cdot \Delta \le s_{n,k} < (t+1) \cdot \Delta\}|, \quad \Delta = 0.02 \text{ s}.
$$

Output: spike count matrix $\mathbf{C} \in \mathbb{N}^{T \times N}$ and EMG envelope matrix $\mathbf{E} \in \mathbb{R}^{T \times 16}$, sharing a common time axis. The script `cortex_fm.data.preprocess_m1` enforces axis alignment, muscle count, and timestamp consistency via assertions.

---

## 4. Evaluation

### 4.1 FALCON M1 protocol

- Benchmark: FALCON M1 task, `falcon-challenge` 1.0.2.
- Decoder wrapper: `cortex_fm.eval.falcon_m1_decoder.CortexFMFalconDecoder` (implements `falcon_challenge.interface.BCIDecoder`).
- Streaming inference: rolling 64-bin (1.28 s) spike buffer, last-bin EMG prediction per step.
- Per-session reset via `reset(dataset_tags)` zeroes the rolling buffer to honor the no-session-leakage contract.
- Batch size: 4 (FALCON M1 recommended).
- Metric: variance-weighted R² over 16 muscles, computed only on bins where `eval_mask == True`.

### 4.2 Auxiliary co-bps

CortexFM's `SpikeReconHead` outputs Poisson rates that yield bits-per-spike above per-unit mean-rate baseline (in-house definition, not FALCON's held-unit co-smoothing). Mean 0.756 ± 0.128 bits/spike on the four held-in calibration files.

### 4.3 Held-in evaluation results

See `README.md` for the full table. Key numbers:

- Zero-shot: per-session R² = −1.035 ± 0.234, NL = 0.131
- Ridge linear probe: pooled R² = +0.125 (positive regime)
- EMG-head FT 200 step: per-session R² = −0.038 ± 0.063
- Per-session affine: per-session R² = +0.484 ± 0.102, pooled R² = **+0.529**

### 4.4 Held-out OOD evaluation

Three sessions (20121004, 20121017, 20121024) recorded 6 – 30 days after pretraining:

- CortexFM + affine pooled R² = **+0.387**
- POYO-1 + affine pooled R² = −0.008
- **Δ = +0.395 pooled R²**, attributable to backbone representation quality (identical affine recipe applied to both backbones).

### 4.5 Saturation observation

EMG-head fine-tuning (37 K params), backbone LayerNorm unfreezing (8 K params), and per-session output-space affine (3 K params/session) are **substitutes**, not additives. All three reach the same pooled R² ≈ +0.529 on held-in and ≈ +0.387 on held-out. The thesis (Chapter 6, §6.5.4) interprets this as a single linear-decodability ceiling in the pretrained latent space.

---

## 5. Limitations

### 5.1 Data scope

The pretraining corpus is **single subject, single session block** (MonkeyL, 4 days). The model has not been validated for:

- Cross-subject transfer (e.g., MonkeyN, MC_Maze, human cortex).
- Cross-task transfer (the task is center-out reach-grasp; other behaviors are out of distribution).
- Cross-array transfer (Utah array M1, left hemisphere; other electrode types or brain regions untested).

### 5.2 Statistical power

Held-out evaluation uses **n = 3 sessions**. Effect sizes are large (per-session Δ vs POYO-1 averages +0.398, ~ 2.4× the standard deviation), but a formal Holm-corrected sign test on n = 3 yields p ≈ 0.25. The thesis labels this as "preliminary external validity" rather than a strong generalization claim.

### 5.3 Calibration dependence

On held-out sessions, the per-session affine requires **≥ ~400 calibration bins (~8 s)** to enter the positive-R² regime stably. Below ~100 bins (~2 s) R² oscillates. Real-time deployment must include a brief per-session calibration cycle.

### 5.4 Zero-shot regression quality

Pure zero-shot inference (no per-session correction) yields negative pooled R² because (a) pretraining mixes Poisson-NLL, MSE, and InfoNCE while FALCON measures EMG-MSE only; (b) the EMG tokenizer receives zeros at inference (out-of-distribution input); (c) no per-session linear correction is applied. The README explains the three resolution paths.

### 5.5 No clinical or assistive validation

CortexFM is a research checkpoint. It has not been tested for:

- Safety, robustness, or efficacy in a clinical BCI.
- Real-time closed-loop control in patient-facing systems.
- Regulatory compliance (e.g., FDA, MFDS, MDR).

Any downstream user planning clinical or assistive use must conduct full domain-specific validation and obtain appropriate regulatory clearance.

---

## 6. Reproducibility

- Configuration file: `src/cortex_fm/configs/pretrain_joint.yaml` in the source repository.
- Random seed: NumPy `np.random.seed(0)` (FALCON evaluator default) for evaluation reproducibility.
- Training duration: ~6 minutes on a single RTX 5080.
- Compute budget total (incl. evaluation): well under one GPU-hour for the full Chapter 5–6 pipeline.

---

## 7. Versioning

| Version | Date | Notes |
|---|---|---|
| pretrain_v1 (epoch28-0.2599.ckpt) | 2026-04-20 | Initial public release. 30 epochs, val_loss = 0.2599, 60.7 MB. |

---

## 8. Author and contact

- **Jaeguk Shin (신재국)** — M.S. candidate, Department of Artificial Intelligence, Dong-eui University, Busan, Republic of Korea.
- Thesis: *CortexFM: A Lightweight Multimodal Foundation Model for Spike–EMG Decoding on Public Brain–Computer Interface Data*, June 2026.
- License: MIT (see `LICENSE`).