Add MEG-XL meg-xl-med-v2 weights converted for braindecode.models.MEGXL (pnpl/MEG-XL@79838987, braindecode PR #1554 @ bf95a14d)

README.md

@@ -0,0 +1,121 @@
+---
+license: mit
+library_name: braindecode
+tags:
+- braindecode
+- pytorch
+- safetensors
+- meg
+- megxl
+- brain-to-text
+- foundation-model
+---
+
+# MEG-XL (pretrained, `meg-xl-med-v2`)
+
+Weights of MEG-XL (Jayalath and Parker Jones, ICML 2026) for
+[`braindecode.models.MEGXL`](https://braindecode.org/stable/generated/braindecode.models.MEGXL.html):
+the released `meg-xl-med-v2.ckpt` (step 113000, the checkpoint the authors' model card
+links) with its frozen BioCodec tokenizer, converted to braindecode's parameter names.
+The class is added by [braindecode#1554](https://github.com/braindecode/braindecode/pull/1554)
+(braindecode 2.1).
+
+```python
+from braindecode.models import MEGXL
+
+model = MEGXL.from_pretrained(
+    "braindecode/megxl-pretrained", chs_info=raw.info["chs"], n_outputs=2
+)
+```
+
+The files hold no classification head: it is not pretrained and its shape depends on the
+number of channels, so `from_pretrained` builds a new one (PyTorch's default init) for
+your `chs_info`. Fine-tune or linear-probe before use.
+
+## Intended input
+
+- MEG only. `chs_info` must carry sensor positions and coil orientations (`loc`);
+  magnetometers and gradiometers (planar or axial) get their own type embeddings. The
+  `chs_info` in `config.json` (the 306 MEG channels of MNE's `test_raw.fif`, Neuromag) is
+  only a default.
+- 50 Hz, preprocessed as in the authors' code: 0.1-40 Hz band-pass, baseline correction
+  and robust scaling per 3 s chunk, clipped to `[-5, 5]`.
+- Pretrained on 150 s windows (`n_times=7500`); the stored RoPE table covers those 625
+  token positions, longer windows rebuild it.
+
+## Source
+
+- Paper: D. Jayalath, O. Parker Jones. *MEG-XL: Data-Efficient Brain-to-Text via
+  Long-Context Pre-Training.* ICML 2026. [arXiv:2602.02494](https://arxiv.org/abs/2602.02494)
+- Code: [neural-processing-lab/MEG-XL @ `97da04171d22e8fe4822cfeb3047a57ecf86b318`](https://github.com/neural-processing-lab/MEG-XL/tree/97da04171d22e8fe4822cfeb3047a57ecf86b318) (MIT)
+- Weights: [pnpl/MEG-XL](https://huggingface.co/pnpl/MEG-XL) at revision
+  `79838987b0fea6bd55bb9b4cccbe1e6f59e97362`, file `meg-xl-med-v2.ckpt`,
+  sha256 `a8be0f06416bebcbea896db8bfdc8027ec25072a64b7e2ab925908dd6edfe88c`
+  (model card `license: mit`).
+- The tokenizer inside the checkpoint equals the separate BioCodec release
+  `biocodec_ckpt.pt` (sha256 `b8b5c864bd92242a2aa70c414ab1abe23231772dc2f2e39a8f509fdded728af1`,
+  [klean2050/BioCodec-release @ `575bf6ea7ae14a34754b217d9b940597a16951bd`](https://github.com/klean2050/BioCodec-release/tree/575bf6ea7ae14a34754b217d9b940597a16951bd), MIT),
+  max-abs difference 0.0.
+
+## Conversion
+
+`convert.py` (in this repository) was run with braindecode at
+`bruAristimunha/braindecode@bf95a14d2ea2847013aead661646dcacf51d9aa3` (branch `feat/meg-megxl`, PR #1554):
+
+```
+python convert.py meg-xl-med-v2.ckpt <out_dir> test_raw.fif
+```
+
+It renames the checkpoint keys to braindecode's and writes `config.json` with
+`save_pretrained`, then `model.safetensors` / `pytorch_model.bin` without the head. Tensor
+values are copied unchanged, except the RoPE table, which the authors' loader recomputes
+from `freqs` (`torch.polar`, 625 positions) and which is stored as `(cos, sin)` pairs.
+
+Dropped (pretraining only or duplicated): the BioCodec decoder, the duplicated
+`biocodec_embedding_*` codebooks, `mask_token`, the code-prediction `output_head`, the
+complex RoPE cache and the Lightning/optimizer state.
+
+| file | sha256 |
+|---|---|
+| `model.safetensors` | `0817547bc254bae86fd74c993a4c5ed941e2e2f35c16d2990046cf32916a99cf` |
+| `pytorch_model.bin` | `cb8b1cb6642bd8ee8070d859b3158b979342a9878f18aee2e7e4dde06c15974e` |
+| `config.json` | `ba072e1d8ccd8db88ed44160deedd8af1499d9ac6884c392732911512d61482c` |
+| `convert.py` | `825bfa078974adb4dc54fc3e5d93f2fbece1f2b0530b9c58cd188e448d8b675b` |
+
+## Verification
+
+`MEGXL.from_pretrained` on these files (only `final_layer.*` missing, no unexpected key;
+loaded tensors equal `convert.py` applied to the original checkpoint) against the
+authors' code with the same weights and inputs, CPU float32, eval mode, 20 s: codes
+identical, features and the gradients of all 119 parameters of one backward pass equal
+with max-abs difference 0.0, for 306 Neuromag channels (batch 1 and 2) and 157 KIT
+axial gradiometers. For that comparison the original's attention dropout is set to 0: the
+released `SelfAttention` keeps dropout 0.1 active in eval mode, braindecode turns it off.
+
+## License
+
+MIT, as the original release: [MEG-XL licence](https://github.com/neural-processing-lab/MEG-XL/blob/main/LICENSE)
+(code; the weights' model card on `pnpl/MEG-XL` says `license: mit`) and
+[BioCodec licence](https://github.com/klean2050/BioCodec-release/blob/main/LICENSE).
+Converted and re-hosted with attribution to the authors; for corrections or removal, open
+an issue at https://github.com/braindecode/braindecode/issues.
+
+## Citation
+
+```bibtex
+@inproceedings{jayalath2026megxl,
+  title     = {{MEG-XL}: Data-Efficient Brain-to-Text via Long-Context Pre-Training},
+  author    = {Jayalath, Dulhan and Parker Jones, Oiwi},
+  booktitle = {Forty-third International Conference on Machine Learning (ICML)},
+  year      = {2026},
+  note      = {arXiv:2602.02494},
+}
+
+@article{aristimunha2025braindecode,
+  title   = {Braindecode: a deep learning library for raw electrophysiological data},
+  author  = {Aristimunha, Bruno and others},
+  journal = {Zenodo},
+  year    = {2025},
+  doi     = {10.5281/zenodo.17699192},
+}
+```

convert.py

@@ -0,0 +1,78 @@
+"""Convert the released MEG-XL Lightning checkpoint to braindecode's MEGXL.
+
+usage: python convert.py meg-xl-med-v2.ckpt <out_dir> <fif with the default chs_info>
+
+The checkpoint holds the frozen BioCodec tokenizer, so ``biocodec_ckpt.pt`` is
+not needed (its tensors equal the copy inside the checkpoint, max-abs 0.0).
+The keys are renamed to braindecode's (``RENAME``). Dropped: the BioCodec
+decoder, the duplicated ``biocodec_embedding_*`` codebooks, ``mask_token``,
+``output_head`` (pretraining only) and the complex RoPE cache. The port's
+RoPE table (``rotate``, as many positions as the model's buffer: 625, the
+150 s pretraining window) is recomputed from ``freqs`` with ``torch.polar``
+on the CPU, as the authors' loader recomputes it, and stored as ``(cos, sin)``
+pairs.
+
+``<out_dir>`` gets ``config.json`` from ``save_pretrained`` (the 306 MEG
+channels of the given FIF file as the default ``chs_info``, 150 s at 50 Hz)
+and ``model.safetensors`` / ``pytorch_model.bin`` without the classification
+head: its shape depends on the number of channels and it is not pretrained,
+so ``from_pretrained`` builds a new one for the caller's ``chs_info``.
+"""
+
+import sys
+
+import torch
+
+DROP = ("tokenizer.decoder.", "biocodec_embedding_", "mask_token", "output_head.")
+RENAME = (
+    ("tokenizer.encoder.model.", "encoder."),
+    (".conv.conv.", ".conv."),
+    ("tokenizer.quantizer.vq.", "quantizer."),
+    ("criss_cross_transformer.layers.", "blocks."),
+    (".ff.layer.0.", ".ff.0."),
+    (".ff.layer.2.", ".ff.3."),
+    ("position_fourier_emb.B", "position_freqs"),
+    ("orientation_fourier_emb.B", "orientation_freqs"),
+)
+
+
+def convert(state_dict, model):
+    """Map ``state_dict`` onto ``model``'s keys."""
+    out = {}
+    for key, value in state_dict.items():
+        if key.startswith(DROP) or key.endswith("rope_embedding_layer.rotate"):
+            continue
+        for old, new in RENAME:
+            key = key.replace(old, new)
+        out[key] = value
+    for key, value in model.state_dict().items():
+        if key.endswith("rope_embedding_layer.rotate"):
+            freqs = out[key.replace("rotate", "freqs")]
+            angles = torch.outer(torch.arange(value.shape[0]).type_as(freqs), freqs)
+            out[key] = torch.view_as_real(torch.polar(torch.ones_like(angles), angles))
+    return out
+
+
+if __name__ == "__main__":
+    import mne
+    from safetensors.torch import save_file
+
+    from braindecode.models import MEGXL
+
+    ckpt_path, out_dir, fif = sys.argv[1:4]
+    raw = mne.io.read_raw_fif(fif, verbose=False)
+    raw.pick(mne.pick_types(raw.info, meg=True, ref_meg=False, exclude=[]))
+    torch.manual_seed(0)
+    model = MEGXL(chs_info=raw.info["chs"], n_outputs=2, n_times=7500, sfreq=50.0)
+    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    state = convert(ckpt["state_dict"], model)
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    assert unexpected == [] and set(missing) == {"final_layer.weight", "final_layer.bias"}
+    model.save_pretrained(out_dir)
+    weights = {
+        k: v.contiguous()
+        for k, v in model.state_dict().items()
+        if not k.startswith("final_layer.")
+    }
+    save_file(weights, f"{out_dir}/model.safetensors", metadata={"format": "pt"})
+    torch.save(weights, f"{out_dir}/pytorch_model.bin")

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0817547bc254bae86fd74c993a4c5ed941e2e2f35c16d2990046cf32916a99cf
+size 97008696

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb8b1cb6642bd8ee8070d859b3158b979342a9878f18aee2e7e4dde06c15974e
+size 97061043