File size: 14,608 Bytes
2e836d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 | # 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`).
|