braindecode.models.ZUNA#

class braindecode.models.ZUNA(n_outputs=None, n_chans=None, chs_info=None, n_times=None, input_window_seconds=None, sfreq=None, *, dim=1024, n_layers=16, n_heads=8, head_dim=64, fine_time_pts=32, latent_dim=32, max_seqlen=256, rope_theta=10000.0, pos_bins=50, pos_half_range=0.12, norm_eps=1e-05, multiple_of=256, ffn_dim_multiplier=None, sandwich_norm=True, qk_norm=True, activation=<class 'torch.nn.modules.activation.SiLU'>)[source]#

ZUNA from Warner et al. (2026) [Warner2026ZUNA].

Foundation Model Channel Attention/Transformer

Added in version 1.7.

ZUNA pretraining encoder-decoder architecture

ZUNA was introduced as a diffusion autoencoder for masked EEG channel reconstruction and super-resolution [Warner2026ZUNA]. ZUNA1.1 retains that objective and adds query-key normalization, sandwich normalization, pretraining windows from 0.5 to 30 seconds, and eight channel and time dropout schemes [Warner2026ZUNA11].

This Braindecode class contains the ZUNA1.1 encoder followed by a classification head. It does not contain the diffusion decoder, channel masking, or reconstruction sampler shown in the figure. forward() returns logits of shape (batch, n_outputs). With return_features=True, it returns the per-channel encoder features used by the classification head.

Signal size and sampling frequency follow the standard Braindecode model arguments. Supply either n_times or both input_window_seconds and sfreq; ZUNA does not set these values implicitly. ZUNA1.1 was trained at 256 Hz with 32-sample tokens. This implementation does not resample, filter, or normalize the input. Channel coordinates are read from chs_info when the model is constructed and stored in fixed rotary position buffers.

Architecture Overview

PatchTokenizer splits every channel into non-overlapping patches of fine_time_pts samples. The resulting n_chans * (n_times // fine_time_pts) patches are serialized in channel-major order. A learned register is placed before each patch, and both are projected to dim before entering the transformer blocks.

Self-attention is bidirectional. Four-dimensional rotary positions encode each token’s discretized scalp coordinates (x, y, z) and coarse-time index. The encoder reads the register positions, projects them to latent_dim, restores the channel and patch axes, and averages over the patch axis. A linear layer maps the concatenated channel features to n_outputs.

Macro Components

  • ZUNA.patch_embedding (torch.nn.Sequential)

    Operations: split (batch, channel, time) into patches with PatchTokenizer, then rearrange (channel, temporal_patch) into one token axis.

    Role: produce one continuous-valued token for every channel and time patch.

  • ZUNA.encoder (torch.nn.Module)

    Operations: tok_embeddings (linear patch embedding) → interleave registers → n_layers × _TransformerBlock (RMS-normed multi-head attention with 4D RoPE and QK-norm, SwiGLU feed-forward, sandwich norm) → norm → output linear projection to latent_dim per token.

    Role: encode channel-time patches while retaining their spatial and temporal coordinates.

  • ZUNA.final_layer (torch.nn.Sequential)

    Operations: rearrange the (n_chans, latent_dim) channel features into one axis → linear map to n_outputs.

    Role: produce task logits from the pooled encoder features.

Temporal, Spatial, and Spectral Encoding

  • Temporal: fixed-size patches provide local samples, and the fourth rotary axis identifies each patch’s coarse-time index. The patch axis is averaged after encoding.

  • Spatial: the first three rotary axes contain bucketed 3D coordinates from chs_info. A model instance uses the montage supplied at construction.

  • Spectral: there is no Fourier or filter-bank stage. The encoder receives raw time-domain patches.

Additional Mechanisms

  • Register tokens: one learned register is paired with every patch. Only register outputs enter the latent projection.

  • ZUNA1.1 normalization: query-key RMS normalization and optional post-attention and post-feed-forward RMS normalization match the changes introduced in ZUNA1.1.

  • Fixed positional grid: channel and coarse-time rotary values are computed during construction. Inputs passed to forward() must have the configured channel count and window length.

Parameters:

  • n_outputs (Optional[int]) – Number of outputs of the model. This is the number of classes in the case of classification.

  • n_chans (Optional[int]) – Number of EEG channels.

  • chs_info (Optional[list[dict]]) – Information about each individual EEG channel. This should be filled with info["chs"]. Refer to mne.Info for more details.

  • n_times (Optional[int]) – Number of time samples of the input window.

  • input_window_seconds (Optional[float]) – Length of the input window in seconds.

  • sfreq (Optional[float]) – Sampling frequency of the EEG recordings.

  • dim (int) – Transformer embedding dimension. The default is 1024.

  • n_layers (int) – Number of transformer blocks. The default is 16.

  • n_heads (int) – Number of attention heads per block. The default is 8.

  • head_dim (int) – Dimension of each attention head. It must be divisible by eight. The default is 64.

  • fine_time_pts (int) – Number of fine time points per token (the encoder input dimension). n_times must be divisible by this value. The default is 32, or 0.125 seconds for data sampled at 256 Hz.

  • latent_dim (int) – Per-token output dimension of the encoder. The default is 32.

  • max_seqlen (int) – Length of the rotary frequency table. It must be at least max(pos_bins, n_times // fine_time_pts). The default is 256.

  • rope_theta (float) – Base period of the rotary positional embedding. The default is 10000.0.

  • pos_bins (int) – Number of buckets per spatial coordinate. The default is 50.

  • pos_half_range (float) – Half-range (in metres) used to normalise channel coordinates before bucketing. Coordinates at or beyond this range are clipped to the first or last bucket. The default is 0.12.

  • norm_eps (float) – Epsilon of the RMS normalization layers. The default is 1e-5.

  • multiple_of (int) – Feed-forward hidden dimension is rounded up to a multiple of this value. The default is 256.

  • ffn_dim_multiplier (Optional[float]) – Multiplier applied before rounding the feed-forward hidden dimension. The default is None.

  • sandwich_norm (bool) – Apply RMS normalization after attention and the feed-forward layer. The default is True.

  • qk_norm (bool) – Apply RMS normalization to queries and keys. The default is True.

  • activation (type[Module]) – Feed-forward activation class. The default is torch.nn.SiLU.

Raises:

ValueError – If some input signal-related parameters are not specified: and can not be inferred.

Notes

The full ZUNA1.1 pretraining model has an encoder and a rectified-flow decoder. Its encoder latent is regularized with a maximum mean discrepancy loss. The decoder and both pretraining losses are outside this class.

In the paper’s downstream experiments, token latents were averaged and the encoder was fine-tuned with the classifier. A frozen encoder with a linear head performed worse. The authors also found that checkpoints with lower reconstruction error were not necessarily better for classification [Warner2026ZUNA11].

References

[Warner2026ZUNA] (1,2)

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. https://arxiv.org/abs/2602.18478

[Warner2026ZUNA11] (1,2)

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. https://arxiv.org/abs/2607.27308

Hugging Face Hub integration

When the optional huggingface_hub package is installed, all models automatically gain the ability to be pushed to and loaded from the Hugging Face Hub. Install with:

pip install braindecode[hub]

Pushing a model to the Hub:

from braindecode.models import ZUNA

# Train your model
model = ZUNA(n_chans=22, n_outputs=4, n_times=1000)
# ... training code ...

# Push to the Hub
model.push_to_hub(
    repo_id="username/my-zuna-model",
    commit_message="Initial model upload",
)

Loading a model from the Hub:

from braindecode.models import ZUNA

# Load pretrained model
model = ZUNA.from_pretrained("username/my-zuna-model")

# Load with a different number of outputs (head is rebuilt automatically)
model = ZUNA.from_pretrained("username/my-zuna-model", n_outputs=4)

Extracting features and replacing the head:

import torch

x = torch.randn(1, model.n_chans, model.n_times)
# Extract encoder features (consistent dict across all models)
out = model(x, return_features=True)
features = out["features"]

# Replace the classification head
model.reset_head(n_outputs=10)

Saving and restoring full configuration:

import json

config = model.get_config()            # all __init__ params
with open("config.json", "w") as f:
    json.dump(config, f)

model2 = ZUNA.from_config(config)    # reconstruct (no weights)

All model parameters (both EEG-specific and model-specific such as dropout rates, activation functions, number of filters) are automatically saved to the Hub and restored when loading.

See Loading and Adapting Pretrained Foundation Models for a complete tutorial.

Methods

forward(input_tensor, return_features=False)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

  • input_tensor (Tensor) – The description is missing.

  • return_features (bool) – The description is missing.

reset_head(n_outputs)[source]#

Replace the classification head for a new number of outputs.

Parameters:

n_outputs (int) – New number of output classes.

Return type:

None

Examples

>>> from braindecode.models import BENDR
>>> model = BENDR(n_chans=22, n_times=1000, n_outputs=4)
>>> model.reset_head(10)
>>> model.n_outputs
10

Added in version 1.4.