braindecode 1.3.0.dev181065563__py3-none-any.whl → 1.3.0.dev183667303__py3-none-any.whl

braindecode/models/patchedtransformer.py

@@ -0,0 +1,640 @@
+# Authors: Jose Mauricio <josemaurici3991@gmail.com>
+#          Bruno Aristimunha <b.aristimunha@gmail.com>
+# License: BSD (3-clause)
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from braindecode.models.base import EEGModuleMixin
+
+
+class PBT(EEGModuleMixin, nn.Module):
+    r"""Patched Brain Transformer (PBT) model from Klein et al. (2025) [pbt]_.
+
+    :bdg-danger:`Large Brain Models`
+
+    This implementation was based in https://github.com/timonkl/PatchedBrainTransformer/
+
+    .. figure:: https://raw.githubusercontent.com/timonkl/PatchedBrainTransformer/refs/heads/main/PBT_sketch.png
+       :align: center
+       :alt:  Patched Brain Transformer Architecture
+       :width: 680px
+
+    PBT tokenizes EEG trials into per-channel patches, linearly projects each
+    patch to a model embedding dimension, prepends a classification token and
+    adds channel-aware positional embeddings. The token sequence is processed
+    by a Transformer encoder stack and classification is performed from the
+    classification token.
+
+    .. rubric:: Macro Components
+
+    - ``PBT.tokenization`` **(patch extraction)**
+
+      *Operations.* The pre-processed EEG signal :math:`X \in \mathbb{R}^{C \times T}`
+      (with :math:`C = \text{n_chans}` and :math:`T = \text{n_times}`) is divided into
+      non-overlapping patches of size :math:`d_{\text{input}}` along the time axis.
+      This process yields :math:`N` total patches, calculated as
+      :math:`N = C \left\lfloor \frac{T}{D} \right\rfloor` (where :math:`D = d_{\text{input}}`).
+      When time shifts are applied, :math:`N` decreases to
+      :math:`N = C \left\lfloor \frac{T - T_{\text{aug}}}{D} \right\rfloor`.
+
+      *Role.* Tokenizes EEG trials into fixed-size, per-channel patches so the model
+      remains adaptive to different numbers of channels and recording lengths.
+      Process is inspired by Vision Transformers [visualtransformer]_ and
+      adapted for GPT context from [efficient-batchpacking]_.
+
+    - ``PBT.patch_projection`` **(patch embedding)**
+
+      *Operations.* The linear layer ``PBT.patch_projection`` maps the tokens from dimension
+      :math:`d_{\text{input}}` to the Transformer embedding dimension :math:`d_{\text{model}}`.
+      Patches :math:`X_P` are projected as :math:`X_E = X_P W_E^\top`, where
+      :math:`W_E \in \mathbb{R}^{d_{\text{model}} \times D}`. In this configuration
+      :math:`d_{\text{model}} = 2D` with :math:`D = d_{\text{input}}`.
+
+      *Interpretability.* Learns periodic structures similar to frequency filters in
+      the first convolutional layers of CNNs (for example :class:`~braindecode.models.EEGNet`).
+      The learned filters frequently focus on the high-frequency range (20-40 Hz),
+      which correlates with beta and gamma waves linked to higher concentration levels.
+
+    - ``PBT.cls_token`` **(classification token)**
+
+      *Operations.* A classification token :math:`[c_{\text{ls}}] \in \mathbb{R}^{1 \times d_{\text{model}}}`
+      is prepended to the projected patch sequence :math:`X_E`. The CLS token can optionally
+      be learnable (see ``learnable_cls``).
+
+      *Role.* Acts as a dedicated readout token that aggregates information through the
+      Transformer encoder stack.
+
+    - ``PBT.pos_embedding`` **(positional embedding)**
+
+      *Operations.* Positional indices are generated by ``PBT.linear_projection``, an instance
+      of :class:`~braindecode.models.patchedtransformer._ChannelEncoding`, and mapped to vectors
+      through :class:`~torch.nn.Embedding`. The embedding table
+      :math:`W_{\text{pos}} \in \mathbb{R}^{(N+1) \times d_{\text{model}}}` is added to the token
+      sequence, yielding :math:`X_{\text{pos}} = [c_{\text{ls}}, X_E] + W_{\text{pos}}`.
+
+      *Role/Interpretability.* Introduces spatial and temporal dependence to counter the
+      position invariance of the Transformer encoder. The learned positional embedding
+      exposes spatial relationships, often revealing a symmetric pattern in central regions
+      (C1-C6) associated with the motor cortex.
+
+    - ``PBT.transformer_encoder`` **(sequence processing and attention)**
+
+      *Operations.* The token sequence passes through :math:`n_{\text{blocks}}` Transformer
+      encoder layers. Each block combines a Multi-Head Self-Attention (MHSA) module with
+      ``num_heads`` attention heads and a Feed-Forward Network (FFN). Both MHSA
+      and FFN use parallel residual connections with Layer Normalization inside the blocks
+      and apply dropout (``drop_prob``) within the Transformer components.
+
+      *Role/Robustness.* Self-attention enables every token to consider all others, capturing
+      global temporal and spatial dependencies immediately and adaptively. This architecture
+      accommodates arbitrary numbers of patches and channels, supporting pre-training across
+      diverse datasets.
+
+    - ``PBT.final_layer`` **(readout)**
+
+      *Operations.* A linear layer operates on the processed CLS token only, and the model
+      predicts class probabilities as :math:`y = \operatorname{softmax}([c_{\text{ls}}] W_{\text{class}}^\top + b_{\text{class}})`.
+
+      *Role.* Performs the final classification from the information aggregated into the CLS
+      token after the Transformer encoder stack.
+
+    .. rubric:: Convolutional Details
+
+    PBT omits convolutional layers; equivalent feature extraction is carried out by the patch
+    pipeline and attention stack.
+
+    * **Temporal.** Tokenization slices the EEG into fixed windows of size :math:`D = d_{\text{input}}`
+      (for the default configuration, :math:`D=64` samples :math:`\approx 0.256\,\text{s}` at
+      :math:`250\,\text{Hz}`), while ``PBT.patch_projection`` learns periodic patterns within each
+      patch. The Transformer encoder then models long- and short-range temporal dependencies through
+      self-attention.
+
+    * **Spatial.** Patches are channel-specific, keeping the architecture adaptive to any electrode
+      montage. Channel-aware positional encodings :math:`W_{\text{pos}}` capture relationships between
+      nearby sensors; learned embeddings often form symmetric motifs across motor cortex electrodes
+      (C1–C6), and self-attention propagates information across all channels jointly.
+
+    * **Spectral.** ``PBT.patch_projection`` acts similarly to the first convolutional layer in
+      :class:`~braindecode.models.EEGNet`, learning frequency-selective filters without an explicit
+      Fourier transform. The highest-energy filters typically reside between :math:`20` and
+      :math:`40\,\text{Hz}`, aligning with beta/gamma rhythms tied to focused motor imagery.
+
+    .. rubric:: Attention / Sequential Modules
+
+    * **Attention Details.** ``PBT.transformer_encoder`` stacks :math:`n_{\text{blocks}}` Transformer
+      encoder layers with Multi-Head Self-Attention. Every token attends to all others, enabling
+      immediate global integration across time and channels and supporting heterogeneous datasets.
+      Attention rollout visualisations highlight strong activations over motor cortex electrodes
+      (C3, C4, Cz) during motor imagery decoding.
+
+
+    .. warning::
+
+        **Important:** As the other Large Brain Models in Braindecode, :class:`PBT` is
+        designed for large-scale pre-training and fine-tuning. Training from
+        scratch on small datasets may lead to suboptimal results. Cross-Dataset
+        pre-training and subsequent fine-tuning is recommended to leverage the
+        full potential of this architecture.
+
+    Parameters
+    ----------
+    d_input : int, optional
+        Size (in samples) of each patch (token) extracted along the time axis.
+    d_model : int, optional
+        Transformer embedding dimensionality.
+    n_blocks : int, optional
+        Number of Transformer encoder layers.
+    num_heads : int, optional
+        Number of attention heads.
+    drop_prob : float, optional
+        Dropout probability used in Transformer components.
+    learnable_cls : bool, optional
+        Whether the classification token is learnable.
+    bias_transformer : bool, optional
+        Whether to use bias in Transformer linear layers.
+    activation : nn.Module, optional
+        Activation function class to use in Transformer feed-forward layers.
+
+    References
+    ----------
+    .. [pbt] Klein, T., Minakowski, P., & Sager, S. (2025).
+        Flexible Patched Brain Transformer model for EEG decoding.
+        Scientific Reports, 15(1), 1-12.
+        https://www.nature.com/articles/s41598-025-86294-3
+    .. [visualtransformer]  Dosovitskiy, A., Beyer, L., Kolesnikov, A.,
+        Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M.,
+        Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J. & Houlsby,
+        N. (2021). An Image is Worth 16x16 Words: Transformers for Image
+        Recognition at Scale. International Conference on Learning
+        Representations (ICLR).
+    .. [efficient-batchpacking] Krell, M. M., Kosec, M., Perez, S. P., &
+        Fitzgibbon, A. (2021). Efficient sequence packing without
+        cross-contamination: Accelerating large language models without
+        impacting performance. arXiv preprint arXiv:2107.02027.
+    """
+
+    def __init__(
+        self,
+        # Signal related parameters
+        n_chans=None,
+        n_outputs=None,
+        n_times=None,
+        chs_info=None,
+        input_window_seconds=None,
+        sfreq=None,
+        # Model parameters
+        d_input: int = 64,
+        d_model: int = 128,
+        n_blocks: int = 4,
+        num_heads: int = 4,
+        drop_prob: float = 0.1,
+        learnable_cls=True,
+        bias_transformer=False,
+        activation: nn.Module = nn.GELU,
+    ) -> None:
+        super().__init__(
+            n_outputs=n_outputs,
+            n_chans=n_chans,
+            chs_info=chs_info,
+            n_times=n_times,
+            input_window_seconds=input_window_seconds,
+            sfreq=sfreq,
+        )
+        del n_outputs, n_chans, chs_info, n_times, input_window_seconds, sfreq
+        # Store hyperparameters
+        self.d_input = d_input
+        self.d_model = d_model
+        self.n_blocks = n_blocks
+        self.num_heads = num_heads
+        self.drop_prob = drop_prob
+
+        # number of distinct positional indices (per-channel tokens + cls)
+        self.num_embeddings = (self.n_chans * self.n_times) // self.d_input
+
+        # Classification token (learnable or fixed zero)
+        if learnable_cls:
+            self.cls_token = nn.Parameter(torch.randn(1, 1, self.d_model) * 0.002)
+        else:
+            # non-learnable zeroed tensor
+            self.cls_token = torch.full(
+                size=(1, 1, self.d_model),
+                fill_value=0,
+                requires_grad=False,
+                dtype=torch.float32,
+            )
+
+        # Patching the signal and producing channel-aware positional index generator
+        self.patch_signal = _Patcher(
+            n_chans=self.n_chans, n_times=self.n_times, d_input=self.d_input
+        )
+
+        # Linear patch projection from token raw-size -> d_model
+        self.patching_projection = nn.Linear(
+            in_features=self.d_input, out_features=self.d_model, bias=False
+        )
+
+        # actual embedding table mapping indices -> d_model
+        self.pos_embedding = nn.Embedding(
+            num_embeddings=self.num_embeddings + 1, embedding_dim=self.d_model
+        )
+
+        # Transformer encoder stack
+        self.transformer_encoder = _TransformerEncoder(
+            n_blocks=n_blocks,
+            d_model=self.d_model,
+            n_head=num_heads,
+            drop_prob=drop_prob,
+            bias=bias_transformer,
+            activation=activation,
+        )
+
+        # classification head on classify token - CLS token
+        self.final_layer = nn.Linear(
+            in_features=d_model, out_features=self.n_outputs, bias=True
+        )
+
+        # initialize weights
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def _init_weights(module: nn.Module) -> None:
+        """Weight initialization following the original implementation."""
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.002)
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        """The implementation follows the original code logic
+
+        - split input into windows of size `(num_embeddings - 1) * d_input`
+        - for each window: reshape into tokens, map positional indices to embeddings,
+          add cls token, run Transformer encoder and collect CLS outputs
+        - aggregate CLS outputs across windows (if >1) and pass through `final_layer`
+
+        Parameters
+        ----------
+        X : torch.Tensor
+            Input tensor with shape (batch_size, n_chans, n_times)
+
+        Returns
+        -------
+        torch.Tensor
+            Output logits with shape (batch_size, n_outputs).
+        """
+        X, int_pos = self.patch_signal(X)
+
+        tokens = self.patching_projection(X)
+
+        cls_token = self.cls_token.expand(X.size(0), 1, -1)
+        cls_idx = torch.zeros((X.size(0), 1), dtype=torch.long, device=X.device)
+
+        tokens = torch.cat([cls_token, tokens], dim=1)
+        pos_emb = self.pos_embedding(int_pos)
+        transformer_out = self.transformer_encoder(tokens + pos_emb)
+
+        return self.final_layer(transformer_out[:, 0])
+
+
+class _LayerNorm(nn.Module):
+    """Layer normalization with optional bias.
+
+    Simple wrapper around :func:`torch.nn.functional.layer_norm` exposing a
+    learnable scale and optional bias.
+
+    Parameters
+    ----------
+    ndim : int
+        Number of features (normalized shape).
+    bias : bool
+        Whether to include a learnable bias term.
+    """
+
+    def __init__(self, ndim: int, bias: bool) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply layer normalization.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor where the last dimension has size `ndim`.
+
+        Returns
+        -------
+        torch.Tensor
+            Normalized tensor with the same shape as input.
+        """
+        return F.layer_norm(
+            x,
+            normalized_shape=self.weight.shape,
+            weight=self.weight,
+            bias=self.bias,
+            eps=1e-5,
+        )
+
+
+class _MHSA(nn.Module):
+    """Multi-head self-attention (MHSA) block.
+
+    Implements a standard multi-head attention mechanism with optional
+    use of PyTorch's scaled_dot_product_attention (FlashAttention) when
+    available and requested.
+
+    Parameters
+    ----------
+    d_model : int
+        Dimensionality of the model / embeddings.
+    n_head : int
+        Number of attention heads.
+    bias : bool
+        Whether linear layers use bias.
+    drop_prob : float, optional
+        drop_prob probability applied to attention weights and residual projection.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_head: int,
+        bias: bool,
+        drop_prob: float = 0.0,
+    ) -> None:
+        super().__init__()
+
+        assert d_model % n_head == 0, "d_model must be divisible by n_head"
+
+        # qkv and output projection
+        self.attn = nn.Linear(d_model, 3 * d_model, bias=bias)
+        self.proj = nn.Linear(d_model, d_model, bias=bias)
+
+        # dropout modules
+        self.attn_drop_prob = nn.Dropout(drop_prob)
+        self.resid_drop_prob = nn.Dropout(drop_prob)
+
+        self.n_head = n_head
+        self.d_model = d_model
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass for MHSA.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, T, C) where C == d_model.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor of shape (B, T, C).
+        """
+
+        B, T, C = x.size()
+
+        # project to q, k, v and reshape for multi-head attention
+        q, k, v = self.attn(x).split(self.d_model, dim=2)
+
+        # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+
+        y = torch.nn.functional.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=None,
+            dropout_p=self.drop_prob if self.training else 0.0,
+            is_causal=False,
+        )
+
+        # re-assemble heads -> (B, T, C)
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+
+        # final linear projection + residual drop_prob
+        y = self.resid_drop_prob(self.proj(y))
+        return y
+
+
+class _FeedForward(nn.Module):
+    """Position-wise feed-forward network from Transformer.
+
+    Implements the two-layer MLP with GELU activation and dropout used in
+    Transformer architectures.
+
+    Parameters
+    ----------
+    d_model : int
+        Input and output dimensionality.
+    dim_feedforward : int, optional
+        Hidden dimensionality of the feed-forward layer. If None, must be provided by caller.
+    drop_prob : float, optional
+        Dropout probability.
+    bias : bool, optional
+        Whether linear layers use bias.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        dim_feedforward: Optional[int] = None,
+        drop_prob: float = 0.0,
+        bias: bool = False,
+        activation: nn.Module = nn.GELU,
+    ) -> None:
+        super().__init__()
+
+        if dim_feedforward is None:
+            raise ValueError("dim_feedforward must be provided")
+
+        self.proj_in = nn.Linear(d_model, dim_feedforward, bias=bias)
+        self.activation = activation()
+        self.proj = nn.Linear(dim_feedforward, d_model, bias=bias)
+        self.drop_prob = nn.Dropout(drop_prob)
+        self.drop_prob1 = nn.Dropout(drop_prob)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the feed-forward block."""
+        x = self.proj_in(x)
+        x = self.activation(x)
+        x = self.drop_prob1(x)
+        x = self.proj(x)
+        x = self.drop_prob(x)
+        return x
+
+
+class _TransformerEncoderLayer(nn.Module):
+    """Single Transformer encoder layer (pre-norm) combining MHSA and feed-forward.
+
+    The block follows the pattern:
+    x <- x + MHSA(_LayerNorm(x))
+    x <- x + FF(_LayerNorm(x))
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_head: int,
+        drop_prob: float = 0.0,
+        dim_feedforward: Optional[int] = None,
+        bias: bool = False,
+        activation: nn.Module = nn.GELU,
+    ) -> None:
+        super().__init__()
+
+        if dim_feedforward is None:
+            dim_feedforward = 4 * d_model
+            # note: preserve the original behaviour (print) from the provided code
+            print(
+                "dim_feedforward is set to 4*d_model, the default in Vaswani et al. (Attention is all you need)"
+            )
+
+        self.layer_norm_att = _LayerNorm(d_model, bias=bias)
+        self.mhsa = _MHSA(d_model, n_head, bias, drop_prob=drop_prob)
+        self.layer_norm_ff = _LayerNorm(d_model, bias=bias)
+        self.feed_forward = _FeedForward(
+            d_model=d_model,
+            dim_feedforward=dim_feedforward,
+            drop_prob=drop_prob,
+            bias=bias,
+            activation=activation,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Execute one encoder layer.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input of shape (B, T, d_model).
+
+        Returns
+        -------
+        torch.Tensor
+            Output of the same shape as input.
+        """
+        x = x + self.mhsa(self.layer_norm_att(x))
+        x = x + self.feed_forward(self.layer_norm_ff(x))
+        return x
+
+
+class _TransformerEncoder(nn.Module):
+    """Stack of Transformer encoder layers.
+
+    Parameters
+    ----------
+    n_blocks : int
+        Number of encoder layers to stack.
+    d_model : int
+        Dimensionality of embeddings.
+    n_head : int
+        Number of attention heads per layer.
+    drop_prob : float
+        Dropout probability.
+    bias : bool
+        Whether linear layers use bias.
+    """
+
+    def __init__(
+        self,
+        n_blocks: int,
+        d_model: int,
+        n_head: int,
+        drop_prob: float,
+        bias: bool,
+        activation: nn.Module = nn.GELU,
+    ) -> None:
+        super().__init__()
+
+        self.encoder_block = nn.ModuleList(
+            [
+                _TransformerEncoderLayer(
+                    d_model=d_model,
+                    n_head=n_head,
+                    drop_prob=drop_prob,
+                    dim_feedforward=None,
+                    bias=bias,
+                    activation=activation,
+                )
+                for _ in range(n_blocks)
+            ]
+        )
+
+        # GPT2-like initialization for linear layers
+        self.apply(self._init_weights)
+        for pn, p in self.named_parameters():
+            if pn.endswith("proj.weight"):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * n_blocks))
+
+    @staticmethod
+    def _init_weights(module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward through all encoder blocks sequentially."""
+        for block in self.encoder_block:
+            x = block(x)
+        return x
+
+
+class _Patcher(nn.Module):
+    """Patching encoding helper.
+
+    This module "patchifies" the original X entry in a ViT manner.
+
+    Parameters
+    ----------
+    n_chans : int
+        Number of EEG channels.
+    n_times : int
+        Number of time samples per trial.
+    d_input : int, optional
+        Number of positional embedding slots allocated per channel.
+    """
+
+    def __init__(self, n_chans=2, n_times=1000, d_input=64) -> None:
+        super().__init__()
+        self.d_input = d_input
+        self.num_tokens = (n_chans * n_times) // d_input
+
+    def forward(self, X: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Returns patched input together with positional indices
+        expanded to batch dimension.
+
+        Parameters
+        ----------
+        X : torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        x_patched: torch.LongTensor
+            Tensor of shape (B, (C*T)//d_input, d_input) containing positional token indices.
+
+        position: torch.LongTensor
+            Tensor of shape ((C*T)//d_input + 1) containing positional token indices.
+        """
+
+        X_flat = X.view(X.size(0), -1)
+        x_patched = X_flat[:, : self.num_tokens * self.d_input].view(
+            X.size(0), self.num_tokens, self.d_input
+        )
+
+        position = torch.arange(x_patched.size(1) + 1, device=X.device)
+
+        return x_patched, position