braindecode.models.MEGGPT#

class braindecode.models.MEGGPT(n_outputs=None, n_chans=None, chs_info=None, n_times=None, input_window_seconds=None, sfreq=None, *, n_tokens=66, n_codes=128, rnn_hidden_size=128, embed_dim=256, num_heads=4, feedforward_dim=1024, n_patches=(100, 25, 20, 10, 8), patch_lens=(1, 4, 4, 8, 8), n_last_steps=16, patch_band=(8, None, None, None, None), activation=<class 'torch.nn.modules.activation.GELU'>, drop_prob=0.2, on_non_divisible='pad', return_features=False)[source]#

MEG-GPT from Huang et al. (2026) [huang2026meggpt].

Foundation Model Attention/Transformer Recurrent

Architecture Overview

  1. The EphysTokenizer encoder [ephystokenizer] turns each channel into a stream of discrete tokens, one per sample: a GRU runs over the channel, a linear layer and a layer norm give one logit per codebook entry, and the argmax is remapped to a token id through the released vocabulary.

  2. The token stream is cut into segments of n_patches[0] * patch_lens[0] tokens (100, i.e. 0.4 s at 250 Hz). Token, time and channel embeddings are summed.

  3. A stack of causal time-attention layers runs on every channel. In layer i the keys and values are pooled into n_patches[i] patches of patch_lens[i] steps and the output keeps the last steps only, so the sequence shrinks to n_last_steps outputs per segment.

  4. The outputs are averaged over time and segments into (batch, n_chans, embed_dim) features (Section 2.4 of the paper), flattened and classified by a linear layer.

Macro Components

MEGGPT.tokenizer

Operations: per-channel GRU, linear layer, layer norm, argmax and vocabulary lookup.

Role: discretise each sample of each channel. The argmax has no gradient, so the tokenizer is frozen (its parameters do not require gradients).

MEGGPT.layers

Operations: pre-norm causal attention over time with patched keys and values (a linear layer mixes the heads and steps of a patch), a banded mask in layers with a patch_band, and a feed-forward block.

Role: model the temporal dependencies of each brain region.

MEGGPT.final_layer

Operations: flatten and linear projection.

Role: map the time-averaged features to class logits.

Differences from the released code

  • The released model also adds a learned embedding of the Cam-CAN session (621 sessions). It is specific to the pre-training data, and the paper discards it for downstream use (Section 2.5.4), so this class has none: the transplanted checkpoint drops those weights, and the forward equals the released model called without session labels.

  • The next-token prediction head is replaced by final_layer.

  • Only the parts of the released configuration are ported: time attention with patching. The reference code’s optional channel and cross attention, unpatched keys, rotary or sinusoidal positions and batch/group norms are off in the released configuration and not implemented.

  • The tokenizer encodes the whole input window in one GRU pass; the reference tokenizes recordings in 200-sample windows, which gives the same tokens for windows of up to 200 samples.

Input

Source-space MEG: 52 parcels of the Glasser 52 parcellation at 250 Hz, each parcel z-scored over the recording (tokenize_session(standardize=True) in the reference). The model does not normalise its input. The channel embedding is indexed by position, so the parcels must come in the order of the released parcellation, sign-flip aligned as in the paper (Section 2.5.3). sfreq is not checked: 250 Hz is assumed. A window that is not a multiple of the segment length is padded or cropped as set by on_non_divisible; with "pad", the zero-padded last segment counts as much as the others in the feature mean.

Channel mask

With return_features=True, forward takes ch_mask, the (batch, n_chans) mask of real channels that pad_channels_collate() returns (real channels first). The tokenizer and the attention read each channel alone, so a recording with n_i real channels gets the features of its first n_i parcels run alone, and the features and tokens of padded channels are zero. The flattened final_layer is tied to n_chans, so ch_mask without return_features=True raises a ValueError. The argmax of the tokenizer passes no gradient to the input.

Token-level self-supervised use

token_mask in forward hides selected clean token embeddings (True=hide) before the causal decoder, without changing the frozen tokenizer or attention masks. encode_tokens() accepts clean cached IDs and returns their tail-output/next-target alignment, channel/time validity and selected-target masks for a downstream vocabulary head and shifted next-token cross entropy. No vocabulary head is installed in this encoder; see that method’s training example. The optional zero-embedding corruption is not the released pretraining method or a bidirectional masked-sample reconstruction objective.

Pre-trained weights

The released Cam-CAN MEG-GPT and EphysTokenizer weights, converted to this class’s keys, are on the Hugging Face Hub at braindecode/meggpt-pretrained (revision 5c49e27c):

from braindecode.models import MEGGPT
model = MEGGPT.from_pretrained("braindecode/meggpt-pretrained", n_outputs=4)

The conversion omits the session embedding, next-token head and tokenizer signal decoder; the classification head is randomly initialized.

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 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.

  • n_tokens (int) – Size of the GPT vocabulary (65 used codebook entries plus id 0 for the unused ones in the released vocabulary).

  • n_codes (int) – Codebook size of the tokenizer.

  • rnn_hidden_size (int) – Hidden size of the tokenizer GRU.

  • embed_dim (int) – Embedding and model dimension.

  • num_heads (int) – Number of attention heads.

  • feedforward_dim (int) – Hidden width of the feed-forward blocks.

  • n_patches (Sequence[int]) – Number of key/value patches in each layer.

  • patch_lens (Sequence[int]) – Steps per key/value patch in each layer. Layer i reads n_patches[i] * patch_lens[i] steps and outputs as many steps as the next layer reads; n_patches[0] * patch_lens[0] is the segment length.

  • n_last_steps (int) – Number of steps the last layer outputs per segment.

  • patch_band (Sequence[Optional[int]]) – Half-width, in patches, of the banded attention of each layer; None for no band.

  • activation (type[Module]) – Activation of the feed-forward blocks.

  • drop_prob (float) – Dropout after the embeddings, the attention and the feed-forward blocks.

  • on_non_divisible (str) – How a window that is not a multiple of the segment length is handled: right zero-padded, cropped, or rejected.

  • return_features (bool) – Whether forward returns the feature dictionary by default.

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.

Methods

encode_tokens(token_ids, token_mask=None, ch_mask=None, lengths=None)[source]#

Encode clean IDs with optional corruption and shifted-target metadata.

Parameters:

  • token_ids (Tensor) – Clean int64 IDs, (batch, n_chans, n_times). Obtain them with model.tokenizer(x) or the reference-compatible cached IDs. The tokenizer is frozen: this is token-level training, not raw masked-sample reconstruction.

  • token_mask (Optional[Tensor]) – Boolean input shape, True=hide (separate from True=real in ch_mask). Zero selected embeddings before positions/context. This optional corruption is not part of the released method.

  • ch_mask (Optional[Tensor]) – Boolean (batch, n_chans) real-channel mask, real parcels first.

  • lengths (Optional[Tensor]) – Int64 (batch,) original valid time counts, default n_times. Invalid IDs (even -1) are sanitized before embedding and targets.

Returns:

features and tokens as in forward; token_ids are sanitized clean prepared IDs. input_indices indexes prepared time; output_indices indexes the last n_last_steps input positions per segment, corresponding to the returned tokens. next_token_targets predicts clean IDs at output_indices+1. target_valid excludes padded channels/times and unavailable next tokens. prediction_mask is target_valid if no corruption mask, else additionally selects hidden TARGET positions (at t+1, not t). These three tensors are (batch, n_chans, n_segments * n_last_steps). Invalid targets are zero, not labels.

Return type:

Dict[str, Tensor]

Notes

A separate vocabulary head can train on the selected valid targets:

encoded = model.encode_tokens(ids, token_mask=hide, lengths=lengths)
selected = encoded["prediction_mask"]
if not selected.any():
    raise ValueError("no selected valid next-token targets")
logits = vocabulary_head(encoded["tokens"])  # (B, C, O, n_tokens)
loss = torch.nn.functional.cross_entropy(
    logits[selected], encoded["next_token_targets"][selected]
)

This is shifted next-token supervision, not same-position masked reconstruction. Later GRU tokens may encode earlier raw samples, so causal attention is retained. No vocabulary head/loss is installed in the encoder, and no reference pretraining quality is claimed.

forward(x, return_features=None, ch_mask=None, token_mask=None)[source]#

ch_mask: optional (batch, n_chans) mask of real channels, see check_ch_mask(). Needs return_features=True; the features and tokens of padded channels are zero. token_mask is a boolean (batch, n_chans, n_times) mask (True=hide): selected clean token embeddings are replaced by zeros before positions and causal context, not before tokenization. See encode_tokens() for next-token targets and loss selection.

Parameters:

  • x (Tensor) – The description is missing.

  • return_features (Optional[bool]) – The description is missing.

  • ch_mask (Optional[Tensor]) – The description is missing.

  • token_mask (Optional[Tensor]) – The description is missing.

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.

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.