braindecode/duin-pretrained

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). The authors' fine-tuning loop was run with both the original model and the braindecode port, starting from the same weights:

	paper (Table 2)	authors' code	braindecode port
Du-IN (mae)	62.70 ± 4.69	60.52 ± 4.58	60.16 ± 4.62
Du-IN (vqvae)	58.24 ± 4.83	55.10 ± 4.51	55.00 ± 4.51
scratch	56.29 ± 5.20	54.06 ± 5.12	53.73 ± 5.13

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.

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