duin-pretrained / README.md
adammnrr's picture
Model card: add the independent CUDA replication (random run order)
e87c8a0 verified
|
Raw History Blame Contribute Delete
5.59 kB
---
library_name: braindecode
license: cc-by-4.0
tags:
- braindecode
- ieeg
- seeg
- speech-decoding
- foundation-model
- du-in
- pytorch_model_hub_mixin
- model_hub_mixin
pipeline_tag: feature-extraction
---
# Du-IN — braindecode weights
braindecode-native re-host of the official **Du-IN** pretrained checkpoints
(Zheng et al., NeurIPS 2024), converted so they load directly with
[`braindecode.models.DuIN`](https://braindecode.org/stable/generated/braindecode.models.DuIN.html).
Du-IN decodes speech (61 Chinese words) from sEEG. Each subject's signal (10
selected bipolar channels, 1000 Hz) goes through a subject-specific spatial
projection, a patch-wise convolutional tokenizer (one token per 100 ms) and an
8-layer Transformer encoder. The encoder is pretrained per subject by masked
modelling of discrete codes learned by a VQ-VAE (`mae`); the VQ-VAE encoder
itself (`vqvae`) is the other released stage.
## Provenance
| | |
|---|---|
| Original code | https://github.com/liulab-repository/Du-IN (MIT) |
| Original weights | https://huggingface.co/datasets/liulab-repository/Du-IN (`pretrains/duin/<subject>/<stage>/model/checkpoint-399.pth`, CC BY 4.0) |
| Paper | Zheng et al. (2024), *Du-IN: Discrete units-guided mask modeling for decoding speech from intracranial neural signals*, NeurIPS 2024, [arXiv:2405.11459](https://arxiv.org/abs/2405.11459) |
These weights are a **format conversion only** of the authors' checkpoints,
with no re-training. For each of the 24 checkpoints:
- the 183 encoder tensors are carried over unchanged (182 bit-for-bit, plus the
subject layer `W (160, 1)`, reshaped to `spatial_projection.weight (16, 10)`);
- the ported encoder reproduces the upstream `duin_cls` encoder output on the
same input to **< 1e-5** (max abs diff 7.2e-6 over all 24);
- `DuIN.from_pretrained` on the uploaded file gives the same state dict and
outputs as loading the original `.pth`.
The pretraining-only parts (`mask_emb`, VQ codebook, MAE token head) are
dropped. The classification head (`head_hidden`, `final_layer`, 61 outputs) is
**randomly initialised** and must be fine-tuned. Per-file source `sha256` are
in `conversion_report.json`.
## Files
| path | content |
|---|---|
| `config.json` | shared `DuIN` config (`n_chans=10`, `n_times=3000`, `sfreq=1000`, `n_outputs=61`) |
| `subj-XXX/mae/model.safetensors` | MAE-pretrained encoder of subject `XXX` (001 to 012), the paper's Du-IN |
| `subj-XXX/vqvae/model.safetensors` | VQ-VAE-stage encoder of subject `XXX` (the paper's "Du-IN (vqvae)") |
| `model.safetensors` | default = `subj-001/mae` |
Every checkpoint is **subject-specific**: it was trained on that subject's 10
channels and only fits that montage, in the channel order of the authors'
dataset. They are not a cross-subject foundation model.
## Usage
```python
from braindecode.models import DuIN
# MAE-pretrained Du-IN of subject 002; the 61-word head is re-initialised.
model = DuIN.from_pretrained(
"braindecode/duin-pretrained",
filename="subj-002/mae/model.safetensors",
)
# Pass n_outputs=... for another task (the head is re-created).
# input: (batch, 10, 3000) — 3 s at 1000 Hz, z-scored per channel
logits = model(x)
```
The reference head ends with a sigmoid and is trained with
`F.cross_entropy(torch.sigmoid(logits), target)`; braindecode returns logits.
Use that loss to match the reference training.
## Replication (braindecode PR #1298)
Fine-tuning on the 61-word task with the authors' `run_cls.py` protocol (200
epochs, AdamW, sigmoid cross-entropy, 6 seeds × 12 subjects, paper-time data
revision `056a87b7`, mean ± SE over subjects). There are two independent runs:
- **A**: the authors' `run_cls.train()` loop run with both the original model
and the port, starting from the same weights (Gaudi2 / CPU);
- **B**: a standalone braindecode re-implementation of the same loop with the
port, on CUDA, with the run order drawn at random per seed (see below).
| | paper (Table 2) | A: authors' code | A: port | B: port |
|---|---|---|---|---|
| Du-IN (mae) | 62.70 ± 4.69 | 60.52 ± 4.58 | 60.16 ± 4.62 | 59.37 ± 4.69 |
| Du-IN (vqvae) | 58.24 ± 4.83 | 55.10 ± 4.51 | 55.00 ± 4.51 | 54.02 ± 4.57 |
| scratch | 56.29 ± 5.20 | 54.06 ± 5.12 | 53.73 ± 5.13 | 53.77 ± 5.10 |
The port matches the authors' code to within 0.4 points. Training the port on
plain logits (`cross_entropy(logits)`) instead of the sigmoid loss gives 55.6
for mae in A and 54.5 in B.
**Split caveat.** The authors' split concatenates a subject's runs in
`os.listdir` order and puts the first 80 % of each word's trials in the
training set. Validation and test therefore come almost entirely from the
**last run listed**, and that run depends on the filesystem. If it is the
chronologically last session (`run6`), accuracy drops sharply. On subject 002
(scratch, seed 42), holding out `run6` gives 43–45 %, while holding out any
other run gives 71–78 %. A sorted run order underestimates the paper by about
8 points on average.
## License
Weights: **CC BY 4.0**, as released by the authors; cite the original work.
braindecode code: BSD-3-Clause; original Du-IN code: MIT.
## Citation
```bibtex
@inproceedings{zheng2024duin,
title={Du-IN: Discrete units-guided mask modeling for decoding speech from intracranial neural signals},
author={Zheng, Hui and Wang, Hai-Teng and Jiang, Wei-Bang and Chen, Zhong-Tao and He, Li and Lin, Pei-Yang and Wei, Peng-Hu and Zhao, Guo-Guang and Liu, Yun-Zhe},
booktitle={Advances in Neural Information Processing Systems},
volume={37},
year={2024}
}
```