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
metadata
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.

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

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

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

@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}
}