Add model_card.md

model_card.md

@@ -0,0 +1,303 @@
+# 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`).