braindecode.models.MEGXL#
- class braindecode.models.MEGXL(n_outputs=None, n_chans=None, chs_info=None, n_times=None, input_window_seconds=None, sfreq=None, *, latent_dim=512, num_layers=8, num_heads=8, drop_prob=0.1, fourier_pos_dim=250, position_scale=1.8, orientation_scale=1.0, tokenizer_dim=128, n_filters=32, ratios=(3, 2, 2), n_q=6, codebook_size=256, codebook_dim=16)[source]#
MEG-XL from Jayalath and Parker Jones (2026) [megxl].
Foundation Model Attention/Transformer
Architecture Overview
A frozen BioCodec tokenizer turns every channel into discrete codes; a criss-cross transformer contextualizes them over time and sensors:
(batch, n_chans, n_times) -> BioCodec (SEANet + RVQ codes) -> codebook embeddings -> projection + sensor embeddings -> criss-cross blocks -> mean over time -> (batch, n_chans * latent_dim) -> (batch, n_outputs)
Macro Components
MEGXL.encoderandMEGXL.quantizerOperations. Each channel is divided by its RMS, a causal SEANet encoder downsamples it by
prod(ratios), andn_qresidual codebooks pick their entry by cosine similarity. Role. Frozen BioCodec tokenizer (no gradients, no codebook updates).MEGXL.rvq_projectorOperations. The
n_qcodebook vectors of a token are concatenated and projected tolatent_dim.MEGXL.position_projector,MEGXL.orientation_projector,MEGXL.sensor_type_layerOperations. Gaussian random Fourier features of the sensor position and orientation, each projected to
latent_dim, plus a gradiometer/magnetometer embedding, added to every token of the channel. Role. Montage-agnostic channel identity.MEGXL.blocksOperations. Half of the features attend over time (RoPE), the other half over the channels, then a feed-forward layer (BrainOmni’s block).
MEGXL.final_layerOperations. Linear layer on the time-averaged, flattened
n_chans * latent_dimfeatures. Role. Classification head (not pretrained).
Temporal, Spatial, and Spectral Encoding
Temporal:
ceil(n_times / 12)tokens of 240 ms at 50 Hz (the SEANet pads the end, as BioCodec does); RoPE attention spans the whole window (150 s in pretraining).Spatial: attention over the channels, keyed by sensor geometry.
Spectral: learned implicitly by the tokenizer’s convolutions.
Additional Mechanisms
Sensor geometry comes from
chs_infoas in the authors’ preprocessing:loc[:3]centred and divided bysqrt(3 * mean(|r|^2)), the coil x-axis of planar gradiometers or the coil normal of other coils as orientation. MEG coils whose name containsMAGare magnetometers, the other MEG coils gradiometers. EEG channels also get the gradiometer embedding, with zero orientation and positions normalised apart from the MEG ones; other channel kinds raise aValueError. MEG-XL is pretrained on MEG only.The RoPE table holds the 625 positions of the 150 s pretraining window (
torch.polarvalues in the converted weights); a longer window rebuilds it fromcos/sin, one ulp away fromtorch.polar.The released attention keeps its dropout (0.1) active in eval mode; here dropout is off in eval mode, so the eval forward is deterministic.
encode()andencode_codes()accept a temporal code-corruptionmaskand caller-ownedmask_valuebefore metadata and context. They return actual temporal hidden states for native masked RVQ-code prediction. The mask sampler, learned mask vector and prediction head belong to the caller, not the default encoder checkpoint. Channel padding (ch_mask) is separate and remains unsupported.
Important
Weights converted from the released
meg-xl-med-v2.ckpt(pnpl/MEG-XL, step 113000, with its frozen BioCodec tokenizer) are on the Hugging Face Hub atbraindecode/megxl-pretrained(revision5b5dcbcc); the head is not pretrained:from braindecode.models import MEGXL model = MEGXL.from_pretrained("braindecode/megxl-pretrained", chs_info=raw.info["chs"], n_outputs=2)
Input is expected at 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]).Added in version 2.1.
- 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 tomne.Infofor 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.
latent_dim (
int) – Transformer dimension.num_layers (
int) – Number of criss-cross blocks.num_heads (
int) – Attention heads (even: half temporal, half spatial).drop_prob (
float) – Dropout of the attention and feed-forward layers.fourier_pos_dim (
int) – Size of the Fourier features of the sensor position and orientation.position_scale (
float) – Standard deviation of the position Fourier frequencies.orientation_scale (
float) – Standard deviation of the orientation Fourier frequencies.tokenizer_dim (
int) – SEANet output dimension.n_filters (
int) – Base number of SEANet filters.n_q (
int) – Number of residual codebooks.codebook_size (
int) – Entries per codebook.codebook_dim (
int) – Codebook dimension.
- Raises:
ValueError – If some input signal-related parameters are not specified: and can not be inferred.
Notes
If some input signal-related parameters are not specified, there will be an attempt to infer them from the other parameters.
References
[megxl]Jayalath, D., Parker Jones, O. (2026). MEG-XL: Data-Efficient Brain-to-Text via Long-Context Pre-Training. ICML 2026. https://arxiv.org/abs/2602.02494
Hugging Face Hub integration
When the optional
huggingface_hubpackage 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 MEGXL # Train your model model = MEGXL(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-megxl-model", commit_message="Initial model upload", model_card_kwargs={"datasets": ["BNCI2014_001"]}, )
Loading a model from the Hub:
from braindecode.models import MEGXL # Load pretrained model model = MEGXL.from_pretrained("username/my-megxl-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 = MEGXL.from_pretrained("username/my-megxl-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 = MEGXL.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
- encode(x, mask=None, mask_value=None)[source]#
Encode clean signals, optionally hiding code content before context.
xhas shape(batch, n_chans, n_times). Returns the actual contextual tokens(batch, n_chans, n_tokens, latent_dim). Seeencode_codes()for the corruption mask and replacement vector. All samples must have the same valid length and the fullchs_infomontage: do not right-pad shorter raw recordings. Global RMS scaling, convolutions and the tokenizer LSTM would change their valid codes.With hop
H = prod(ratios),n_tokens = ceil(n_times / H)includes the native boundary-extension token of a partial final hop. An encoded mask is not a waveform mask: neighboring codes can contain information from overlapping receptive fields and the whole-window RMS.
- encode_codes(codes, mask=None, mask_value=None)[source]#
Contextualize clean BioCodec codes, tokenizing only once for SSL.
- Parameters:
codes (
Tensor) – Clean integer targets fromtokenize(), shape(batch, n_chans, n_tokens, n_q). Keep these targets detached.mask (
Optional[Tensor]) – Boolean(batch, n_chans, n_tokens), True to hide the entire projected RVQ content at that token. Replacement is BEFORE adding sensor metadata and BEFORE attention. Corrupted valid tokens still participate in attention; this is not a padding mask. The native sampler selects temporal spans shared across all channels.mask_value (
Optional[Tensor]) – Replacement vector of shape(latent_dim,), required withmask. A caller can own a learnable vector and a prediction headLinear(latent_dim, n_q * codebook_size). The native objective is selected-code cross-entropy per RVQ level, averaged across levels; it is not causal next-token prediction or raw-signal MSE.
- Returns:
Actual contextual tokens
(batch, n_chans, n_tokens, latent_dim). No new parameters are allocated; default checkpoint keys and the ordinary forward/return_features dictionary are unchanged.- Return type:
Notes
All tokens must belong to equal-valid-length, full-montage windows. There is no temporal key-padding mask. For complete-hop-only loss, exclude token indices >=
n_times // prod(ratios); this excludes the partial boundary token, not its context. Do not batch padded code sequences and merely hide their loss. Bucket recordings by length.
- forward(x, return_features=False, ch_mask=None)[source]#
Classify
xor return the pooled features.- Parameters:
x (
Tensor) – Input of shape(batch, n_chans, n_times).return_features (
bool) – If True, return{"features": (batch, n_chans * latent_dim), "cls_token": None}instead of the logits.ch_mask (
Optional[Tensor]) – Not supported: any value butNoneraises aValueError. The head flattens then_chans * latent_dimfeatures and the sensor geometry is the fixedchs_infomontage, so every sample needs alln_chanschannels.
- Return type: