braindecode.models.MVPFormer#

class braindecode.models.MVPFormer(n_outputs=None, n_chans=None, chs_info=None, n_times=None, input_window_seconds=None, sfreq=None, *, segment_len=2560, d_model=256, n_layers=4, n_heads=8, n_head_kv=4, d_inner=512, local_window=10, global_att=True, max_segments=110, max_channels=128, drop_prob=0.1, activation=<class 'torch.nn.modules.activation.SiLU'>, pooling='mean', channel_strategy='native', channel_strategy_kwargs=None)[source]#

MVPFormer from Carzaniga et al. (2026) [Carzaniga2026].

Foundation Model Attention/Transformer

MVPFormer model overview

MVPFormer model overview from [Carzaniga2026].#

Multi-variate parallel attention mechanism

Multi-variate parallel attention (MVPA) from [Carzaniga2026].#

Added in version 1.6.1.

A foundation model with multi-variate parallel attention for heterogeneous multi-variate iEEG. The raw signal is tokenised channel-wise into continuous wavelet embeddings and processed by a decoder-only (Llama2-style) transformer whose self-attention is decomposed into content, time-relative and channel-relative terms over a (segment, channel) token grid [Carzaniga2026].

Architecture Overview

MVPFormer combines three ideas to model heterogeneous multi-variate iEEG [Carzaniga2026]:

  1. Continuous wavelet tokenization. A db4 wavelet encoder maps each fixed-length signal segment (per channel) to a continuous embedding, rather than a discrete vocabulary, preserving fine signal structure.

  2. Multi-variate parallel attention (MVPA). Self-attention over the 2D (segment, channel) token grid is decomposed into three additive terms – content, time-relative and channel-relative – so the temporal and spatial axes are modelled jointly at the attention level.

  3. Llama2-style decoder. Parallel attention and MLP blocks, grouped-query attention and causal masking in time, trained to predict the next-in-time embedding (generative pre-training).

Macro Components

MVPFormer.patch_tokenizer + MVPFormer.patch_embed (Wavelet encoder)

Operations. PatchTokenizer splits the raw signal \(\mathbf{x} \in \mathbb{R}^{C \times T}\) into non-overlapping segments of segment_len samples. Each segment passes independently through a fixed db4 discrete wavelet decomposition, is RMS-normalised, and is linearly projected to a continuous embedding \(\mathbf{e}_{c,t} \in \mathbb{R}^{d}\) (_WaveletPatchEmbed). The wavelet transform has no learnable parameters; only the norm and projection are trained. Role. Turns the heterogeneous raw signal into a (segment, channel) grid of continuous tokens whose frequency content is exposed by the wavelet sub-bands.

MVPFormer.blocks (MVPA decoder stack)

Operations. A stack of _MVPABlock modules in the parallel attention/MLP configuration \(\mathbf{o} = \mathbf{o} + \text{attn}(\text{ln}(\mathbf{o})) + \text{mlp}(\text{ln}(\mathbf{o}))\) (Megatron-LM style) with RMSNorm and a SwiGLU MLP. Each block computes multi-variate parallel attention (_MVPAttention) between tokens at grid positions \((c,t)\) and \((c',t')\) as a sum of three terms:

\[\begin{split}\mathbf{a}^{\text{MVPA}}_{c,t,c',t'} &= \mathbf{x}_{c,t}^\top W_q^\top W_{k_e}\,\mathbf{x}_{c',t'} + \mathbf{u}^\top W_{k_e}\,\mathbf{x}_{c',t'} && \text{(content)} \\ &\;+ \mathbf{x}_{c,t}^\top W_q^\top W_{k_t}\,\mathcal{T}_{t-t'} + \mathbf{v}^\top W_{k_t}\,\mathcal{T}_{t-t'} && \text{(time-relative)} \\ &\;+ \mathbf{x}_{c,t}^\top W_q^\top W_{k_c}\,\mathcal{C}_{c-c'} + \mathbf{w}^\top W_{k_c}\,\mathcal{C}_{c-c'} && \text{(channel-relative)}\end{split}\]

where \(\mathcal{T}\) and \(\mathcal{C}\) are learnable time and channel positional codebooks and \(\mathbf{u}, \mathbf{v}, \mathbf{w}\) are learnable bias terms. Attention is causal in time and uses grouped-query attention. Role. Models content, temporal and spatial dependencies jointly, producing one output embedding per input segment.

MVPFormer.final_layer (Classification head)

Operations. The last time-segment of the decoder output is pooled over channels (pooling="mean" or "concat") and passed through a linear layer. Role. Maps the pooled representation to n_outputs class logits.

Temporal, Spatial, and Spectral Encoding

  • Temporal: Time is the first grid axis. The time-relative attention term encodes the signed distance \(t-t'\) between segments (shared across channels) in Transformer-XL fashion, and attention is causally masked so a segment only attends to the past. Content attention is restricted to a local window of local_window segments while the relative term still spans the full context.

  • Spatial: Channels are the second grid axis. The channel-relative attention term encodes the relative distance \(c-c'\) between channels (shared across time); being relative, it is montage-agnostic and can recover a connectivity map from random initialisation, independent of the absolute electrode layout.

  • Spectral: Frequency content is captured at tokenization: the db4 discrete wavelet decomposition of each segment yields multi-resolution sub-band coefficients before projection, so spectral structure is embedded in every token rather than learned by explicit band-pass filters.

Additional Mechanisms

  • Grouped-query attention. n_head_kv key/value heads are shared across n_heads query heads to reduce memory and compute.

  • Relative terms and content window. The time- and channel-relative terms are quadratic in one axis and constant in the other (computed once and broadcast), and a relative-shift trick yields all relative offsets in one pass. The local content window bounds each query’s effective context, but this reference implementation still materialises the full (segment*channel)-squared content-attention matrix before masking; the Triton FlashMVPA kernel that avoids that cost is not ported, so memory scales quadratically with the total number of tokens.

  • Variable channel counts. With ch_mask and pooling="mean", one batch can mix channel counts; each sample gives its unpadded output.

  • Generative pre-training (reference only). Upstream, MVPFormer is pre-trained to predict the next-in-time embedding with a contrastive loss and fine-tuned with LoRA. This braindecode implementation provides the architecture; the Triton FlashMVPA kernel and the contrastive pre-training loop are not included.

Parameters:

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

  • n_chans (int) – Number of EEG channels.

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

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

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

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

  • segment_len (int) – Length of each temporal segment in samples.

  • d_model (int) – Token embedding dimension.

  • n_layers (int) – Number of decoder blocks.

  • n_heads (int) – Number of attention (query) heads.

  • n_head_kv (int) – Number of key/value heads (grouped-query attention).

  • d_inner (int) – Hidden dimension of the SwiGLU MLP.

  • local_window (int) – Size, in segments, of the causal local content-attention lookback window.

  • global_att (bool) – Whether to use the global content-attention term.

  • max_segments (int) – Size of the segment (time) embedding table; must be >= number of segments produced from the input.

  • max_channels (int) – Size of the channel embedding table; must be >= n_chans.

  • drop_prob (float) – Dropout rate (embedding, attention and residual).

  • activation (type[Module]) – Activation class used in the SwiGLU MLP (default torch.nn.SiLU).

  • pooling (str) – How to pool the last segment over channels before the head. "mean" (default, as used for seizure detection) is montage-agnostic; "concat" flattens channels and ties the head to n_chans.

  • channel_strategy (str) – How any montage reaches the backbone (pretrained models only; see Channel strategies: any montage in). "native" keeps the model as it is. "exact", "zero", "nearest", "idw", "spline", "field", "source", "region", "wiener" (call model.channel_layer.fit first) or "latent" map the montage of chs_info (or of the chs_info given to forward) onto the backbone’s channels with a ChannelLayer. Saved in the config. model(x) and model.forward(x) both apply the layer.

  • channel_strategy_kwargs (dict | None) – Options of the strategy (e.g. {"reg": 1e-2} for "spline").

Raises:

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

Notes

The defaults are a small, CI-friendly configuration. The published model sizes are MVPFormer-S (75M, d_model=768, n_layers=12, n_heads=12, n_head_kv=4, d_inner=1728) and MVPFormer-M (1.2B, d_model=2048, n_layers=24, n_heads=16, n_head_kv=8, d_inner=5632); both use 5-second segments at 512 Hz, i.e. segment_len=2560.

References

[Carzaniga2026] (1,2,3,4,5)

Carzaniga, F. S., Hersche, M., Sebastian, A., Schindler, K., & Rahimi, A. (2026). A foundation model with multi-variate parallel attention to generate neuronal activity. In The Fourteenth International Conference on Learning Representations. https://openreview.net/forum?id=5M1YOW3bRq

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 MVPFormer

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

# Push to the Hub; ``datasets`` and ``base_model`` go to the card
model.push_to_hub(
    repo_id="username/my-mvpformer-model",
    commit_message="Initial model upload",
    model_card_kwargs={"datasets": ["BNCI2014_001"]},
)

Loading a model from the Hub:

from braindecode.models import MVPFormer

# Load pretrained model
model = MVPFormer.from_pretrained("username/my-mvpformer-model")
# Data the weights were trained on, from the Hub model card
print(model.card_data.datasets if model.card_data else None)

# Load with a different number of outputs (head is rebuilt automatically)
model = MVPFormer.from_pretrained("username/my-mvpformer-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 = MVPFormer.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(x, return_features=False, ch_mask=None)[source]#

Decode a batch of signals.

Parameters:

  • x (torch.Tensor) – Input of shape (batch, n_chans, n_times). Samples with fewer channels are padded along the channel axis and flagged by ch_mask.

  • return_features (bool) – If True, return the pooled last-segment embedding instead of the class logits, as {"features": pooled, "cls_token": None} (braindecode foundation-model convention). MVPFormer has no class token, hence cls_token is None.

  • ch_mask (Optional[Tensor]) – (batch, n_chans) mask of real channels, see check_ch_mask(). Each sample is processed as if unpadded. Needs pooling="mean".

Returns:

Class logits of shape (batch, n_outputs), or the feature dict {"features", "cls_token"} when return_features is set.

Return type:

Union[Tensor, Dict[str, Optional[Tensor]]]

reset_head(n_outputs)[source]#

Replace the classification head for a new number of outputs.

This is called automatically by from_pretrained() when the user passes an n_outputs that differs from the saved config. Override in subclasses that need a model-specific head structure. Implementations keep changed constructor arguments in sync with self._update_init_kwargs, so that a saved model can be loaded back. Implementations requiring positive outputs can also use self._set_n_outputs to validate and record the new value.

Parameters:

n_outputs (int) – New number of output classes.

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.