ZUNA (braindecode re-host)

Re-host of the ZUNA1.1 EEG foundation-model encoder weights for use with braindecode, converted to the parameter names of braindecode's ZUNA class.

• Original model: Zyphra/ZUNA1.1 at commit 10b9d6c8 (model-00001-of-00001.safetensors, sha256 295fcddb42ab4623dc1346fba12176798d97ab95389c666907e32cbce23f7f33)
• Original code: https://github.com/Zyphra/zuna (checked against commit c7b0195)
• Papers:
  • Warner, C., Mago, J., Huml, J.R. and Millidge, B. (2026). ZUNA1.1: A More Flexible EEG Foundation Model for Denoising and Super-Resolution. arXiv:2607.27308
  • Warner, C., Mago, J., Huml, J.R., Osman, M. and Millidge, B. (2026). ZUNA: Flexible EEG Superresolution with Position-Aware Diffusion Autoencoders. arXiv:2602.18478
• Original authors (Zyphra): Chris Warner, Jonas Mago, Jon Huml, et al.
• License: Apache-2.0 (inherited from the upstream release)

Why this re-host

The braindecode ZUNA port loads these weights through ZUNA.from_pretrained("braindecode/ZUNA", ...). Re-hosting under the braindecode org gives a stable, permanent location that the library can point to, so the integration does not depend on the upstream repository staying unchanged.

The upstream file holds the full ZUNA1.1 pretraining model (encoder and rectified-flow decoder, 639 tensors). model.safetensors here holds its 213 encoder tensors, unchanged in value (float32, bit-identical), renamed to the braindecode class's parameter names; the 426 decoder tensors are not included. convert_zuna11.py in this repository produces the file from the upstream checkpoint: it downloads the pinned revision, checks its sha256, renames the keys, loads them into ZUNA with every non-head key required (the only missing keys are the classification head) and saves the encoder with save_pretrained.

config.json records ZUNA1.1's settings: the class defaults, sfreq 256, and pos_bins 100. ZUNA1.1 buckets each electrode coordinate into 100 bins over ±0.12 m (arXiv:2607.27308, Table II; ZUNA1 used 50). The input geometry (channels, window length, number of outputs) is left to you.

Parity with the original code

On identical random inputs (19 standard 10-05 channels and the 32-channel BioSemi montage, 256 Hz, 5 s and 10 s windows, several seeds), braindecode's encoder was compared with the encoder of the original Zyphra code (EncoderTransformer, eval mode, same weights):

• the patch tokens and the discretized electrode positions are identical;
• float64 against float64: per-token latents differ by at most 9.1e-13, the per-channel features by at most 2.8e-14 (the original code's float32 casts redirected to float64 for this check);
• float32 against float32: per-token latents (values up to ~10) differ by at most 3.0e-4, the per-channel features used by the classification head by at most 8.3e-6. Reordering the tokens in the original code alone (its own time-major against channel-major layout) changes its float32 output by up to 2.9e-4, so the float32 gap is rounding.

ZUNA.from_pretrained on this repository loads all 213 encoder tensors; its features are bit-identical to the converted model's.

What is (and is not) pretrained

These are the pretrained encoder weights of ZUNA1.1 (a masked diffusion autoencoder trained to reconstruct, denoise and upsample EEG). The braindecode wrapper adds a classification head that is randomly initialised and must be fine-tuned on your downstream task — loading these weights alone does not give a trained classifier. The diffusion decoder and the reconstruction sampler are not part of the braindecode class; use the upstream zuna package for reconstruction.

Training data

datasets: id	Dataset	Stage	Modality
TUEG	Temple University Hospital EEG Corpus (TUH EEG)	pretraining	EEG

• ZUNA1.1 was pretrained on approximately 3.5M channel-hours of public EEG, the ZUNA1 corpus (about 2M channel-hours from the TUH EEG Corpus and OpenNeuro, 208 datasets) grown through quality-aware loading of partially noisy channels and pre-epoched datasets.
• TUEG (the TUH EEG Corpus) contains the recordings of its sub-corpora TUAB, TUEV, TUSZ, TUAR, TUSL and TUEP; an evaluation on one of those is not held out from these weights unless the authors removed it.
• Incomplete list. The corpus also includes OpenNeuro datasets that the papers, the upstream model cards and the upstream code do not list, so they are not in datasets:. A dataset missing here is therefore not proof that these weights never saw it.

Usage

from braindecode.models import ZUNA

# Channel locations come from chs_info (metres, MNE head frame); a channel
# without a location takes the standard 10-05 position of its name.
model = ZUNA.from_pretrained(
    "braindecode/ZUNA",
    chs_info=raw.info["chs"],  # your montage
    n_times=1280,              # e.g. 5 s at 256 Hz; ZUNA1.1 saw 0.5 s to 30 s
    n_outputs=4,               # your task; this head is new and untrained
)

Inputs are EEG sampled at 256 Hz. braindecode does not resample, filter or normalise; the upstream inference pipeline high-pass filters at 0.5 Hz, average-references each segment and z-scores each channel before the encoder.

Revisions

• Earlier revision 10cd7745 (2026-08-01) held the raw ZUNA1 file of Zyphra/ZUNA (sha256 31e68d14…, upstream key names model.*). Its keys did not match the braindecode class, so ZUNA.from_pretrained on it loaded no tensor and left the encoder randomly initialised (without an error). The current revision holds the ZUNA1.1 encoder in the class's own parameter names.

Citation

@article{warner2026zuna11,
  title   = {{ZUNA1.1}: A More Flexible {EEG} Foundation Model for Denoising and Super-Resolution},
  author  = {Warner, Christopher and Mago, Jonas and Huml, JR and Millidge, Beren},
  journal = {arXiv preprint arXiv:2607.27308},
  year    = {2026}
}

@article{warner2026zuna,
  title   = {{ZUNA}: Flexible {EEG} Superresolution with Position-Aware Diffusion Autoencoders},
  author  = {Warner, Christopher and Mago, Jonas and Huml, JR and Osman, Mohamed and Millidge, Beren},
  journal = {arXiv preprint arXiv:2602.18478},
  year    = {2026}
}

Cite Braindecode

If you use this braindecode port in scientific work, cite the software on Zenodo, doi:10.5281/zenodo.17699192, and the original model and datasets cited above.

@software{braindecode,
  author = {Aristimunha, Bruno and
            Guetschel, Pierre and
            Wimpff, Martin and
            Gemein, Lukas and
            Rommel, Cedric and
            Banville, Hubert and
            Sliwowski, Maciej and
            Wilson, Daniel and
            Brandt, Simon and
            Gnassounou, Théo and
            Paillard, Joseph and
            {Junqueira Lopes}, Bruna and
            Sedlar, Sara and
            Moreau, Thomas and
            Chevallier, Sylvain and
            Gramfort, Alexandre and
            Schirrmeister, Robin Tibor},
  title = {Braindecode: toolbox for decoding raw electrophysiological brain data
           with deep learning models},
  url = {https://github.com/braindecode/braindecode},
  doi = {10.5281/zenodo.17699192},
  publisher = {Zenodo},
  license = {BSD-3-Clause},
}