braindecode 1.2.0.dev182094932__py3-none-any.whl → 1.3.0.dev173691341__py3-none-any.whl

braindecode/models/attentionbasenet.py

@@ -26,25 +26,150 @@ from braindecode.modules.attention import (
 class AttentionBaseNet(EEGModuleMixin, nn.Module):
     """AttentionBaseNet from Wimpff M et al. (2023) [Martin2023]_.
 
+    :bdg-success:`Convolution` :bdg-info:`Small Attention`
+
     .. figure:: https://content.cld.iop.org/journals/1741-2552/21/3/036020/revision2/jnead48b9f2_hr.jpg
-       :align: center
-       :alt: Attention Base Net
+        :align: center
+        :alt: AttentionBaseNet Architecture
+        :width: 640px
+
+
+    .. rubric:: Architectural Overview
+
+    AttentionBaseNet is a *convolution-first* network with a *channel-attention* stage.
+    The end-to-end flow is:
+
+    - (i) :class:`_FeatureExtractor` learns a temporal filter bank and per-filter spatial
+      projections (depthwise across electrodes), then condenses time by pooling;
+    - (ii) **Channel Expansion** uses a ``1x1`` convolution to set the feature width;
+    - (iii) :class:`_ChannelAttentionBlock` refines features via depthwise–pointwise temporal
+      convs and an optional channel-attention module (SE/CBAM/ECA/…);
+    - (iv) **Classifier** flattens the sequence and applies a linear readout.
+
+    This design mirrors shallow CNN pipelines (EEGNet-style stem) but inserts a pluggable
+    attention unit that *re-weights channels* (and optionally temporal positions) before
+    classification.
+
+
+    .. rubric:: Macro Components
+
+    - :class:`_FeatureExtractor` **(Shallow conv stem → condensed feature map)**
+
+        - *Operations.*
+        - **Temporal conv** (:class:`torch.nn.Conv2d`) with kernel ``(1, L_t)`` creates a learned
+          FIR-like filter bank with ``n_temporal_filters`` maps.
+        - **Depthwise spatial conv** (:class:`torch.nn.Conv2d`, ``groups=n_temporal_filters``)
+          with kernel ``(n_chans, 1)`` learns per-filter spatial projections over the full montage.
+        - **BatchNorm → ELU → AvgPool → Dropout** stabilize and downsample time.
+        - Output shape: ``(B, F2, 1, T₁)`` with ``F2 = n_temporal_filters x spatial_expansion``.
+
+    *Interpretability/robustness.* Temporal kernels behave as analyzable FIR filters; the
+    depthwise spatial step yields rhythm-specific topographies. Pooling acts as a local
+    integrator that reduces variance on short EEG windows.
+
+    - **Channel Expansion**
+
+        - *Operations.*
+        - A ``1x1`` conv → BN → activation maps ``F2 → ch_dim`` without changing
+          the temporal length ``T₁`` (shape: ``(B, ch_dim, 1, T₁)``).
+          This sets the embedding width for the attention block.
+
+    - :class:`_ChannelAttentionBlock` **(temporal refinement + channel attention)**
+
+        - *Operations.*
+        - **Depthwise temporal conv** ``(1, L_a)`` (groups=``ch_dim``) + **pointwise ``1x1``**,
+          BN and activation → preserves shape ``(B, ch_dim, 1, T₁)`` while refining timing.
+        - **Optional attention module** (see *Additional Mechanisms*) applies channel reweighting
+          (some variants also apply temporal gating).
+        - **AvgPool (1, P₂)** with stride ``(1, S₂)`` and **Dropout** → outputs
+          ``(B, ch_dim, 1, T₂)``.
+
+    *Role.* Emphasizes informative channels (and, in certain modes, salient time steps)
+    before the classifier; complements the convolutional priors with adaptive re-weighting.
+
+    - **Classifier (aggregation + readout)**
+
+    *Operations.* :class:`torch.nn.Flatten` → :class:`torch.nn.Linear` from
+    ``(B, ch_dim·T₂)`` to classes.
+
+
+    .. rubric:: Convolutional Details
+
+    - **Temporal (where time-domain patterns are learned).**
+        Wide kernels in the stem (``(1, L_t)``) act as a learned filter bank for oscillatory
+        bands/transients; the attention block’s depthwise temporal conv (``(1, L_a)``) sharpens
+        short-term dynamics after downsampling. Pool sizes/strides (``P₁,S₁`` then ``P₂,S₂``)
+        set the token rate and effective temporal resolution.
+
+    - **Spatial (how electrodes are processed).**
+        A depthwise spatial conv with kernel ``(n_chans, 1)`` spans the full montage to
+        learn *per-temporal-filter* spatial projections (no cross-filter mixing at this step),
+        mirroring the interpretable spatial stage in shallow CNNs.
 
-    Neural Network from the paper: EEG motor imagery decoding:
-    A framework for comparative analysis with channel attention
-    mechanisms
+    - **Spectral (how frequency content is captured).**
+        No explicit Fourier/wavelet transform is used in the stem—spectral selectivity
+        emerges from learned temporal kernels. When ``attention_mode="fca"``, a frequency
+        channel attention (DCT-based) summarizes frequencies to drive channel weights.
 
-    The paper and original code with more details about the methodological
-    choices are available at the [Martin2023]_ and [MartinCode]_.
 
-    The AttentionBaseNet architecture is composed of four modules:
-    - Input Block that performs a temporal convolution and a spatial
-    convolution.
-    - Channel Expansion that modifies the number of channels.
-    - An attention block that performs channel attention with several
-    options
-    - ClassificationHead
+    .. rubric:: Attention / Sequential Modules
 
+    - **Type.** Channel attention chosen by ``attention_mode`` (SE, ECA, CBAM, CAT, GSoP,
+        EncNet, GE, GCT, SRM, CATLite). Most operate purely on channels; CBAM/CAT additionally
+        include temporal attention.
+
+    - **Shapes.** Input/Output around attention: ``(B, ch_dim, 1, T₁)``. Re-arrangements
+        (if any) are internal to the module; the block returns the same shape before pooling.
+
+    - **Role.** Re-weights channels (and optionally time) to highlight informative sources
+        and suppress distractors, improving SNR ahead of the linear head.
+
+
+    .. rubric:: Additional Mechanisms
+
+        - **Attention variants at a glance.**
+        - ``"se"``: Squeeze-and-Excitation (global pooling → bottleneck → gates).
+        - ``"gsop"``: Global second-order pooling (covariance-aware channel weights).
+        - ``"fca"``: Frequency Channel Attention (DCT summary; uses ``seq_len`` and ``freq_idx``).
+        - ``"encnet"``: EncNet with learned codewords (uses ``n_codewords``).
+        - ``"eca"``: Efficient Channel Attention (local 1-D conv over channel descriptor; uses ``kernel_size``).
+        - ``"ge"``: Gather–Excite (context pooling with optional MLP; can use ``extra_params``).
+        - ``"gct"``: Gated Channel Transformation (global context normalization + gating).
+        - ``"srm"``: Style-based recalibration (mean–std descriptors; optional MLP).
+        - ``"cbam"``: Channel then temporal attention (uses ``kernel_size``).
+        - ``"cat"`` / ``"catlite"``: Collaborative (channel ± temporal) attention; *lite* omits temporal.
+        - **Auto-compatibility on short inputs.**
+
+    If the input duration is too short for the configured kernels/pools, the implementation
+    **automatically rescales** temporal lengths/strides downward (with a warning) to keep
+    shapes valid and preserve the pipeline semantics.
+
+
+    .. rubric:: Usage and Configuration
+
+    - ``n_temporal_filters``, ``temporal_filter_length`` and ``spatial_expansion``:
+        control the capacity and the number of spatial projections in the stem.
+    - ``pool_length_inp``, ``pool_stride_inp`` then ``pool_length``, ``pool_stride``:
+        trade temporal resolution for compute; they determine the final sequence length ``T₂``.
+    - ``ch_dim``: width after the ``1x1`` expansion and the effective embedding size for attention.
+    - ``attention_mode`` + its specific hyperparameters (``reduction_rate``,
+        ``kernel_size``, ``seq_len``, ``freq_idx``, ``n_codewords``, ``use_mlp``):
+        select and tune the reweighting mechanism.
+    - ``drop_prob_inp`` and ``drop_prob_attn``: regularize stem and attention stages.
+    - **Training tips.**
+
+    Start with moderate pooling (e.g., ``P₁=75,S₁=15``) and ELU activations; enable attention
+    only after the stem learns stable filters. For small datasets, prefer simpler modes
+    (``"se"``, ``"eca"``) before heavier ones (``"gsop"``, ``"encnet"``).
+
+    Notes
+    -----
+    - Sequence length after each stage is computed internally; the final classifier expects
+      a flattened ``ch_dim x T₂`` vector.
+    - Attention operates on *channel* dimension by design; temporal gating exists only in
+      specific variants (CBAM/CAT).
+    - The paper and original code with more details about the methodological
+      choices are available at the [Martin2023]_ and [MartinCode]_.
     .. versionadded:: 0.9
 
     Parameters
@@ -73,18 +198,18 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         the depth of the network after the initial layer. Default is 16.
     attention_mode : str, optional
         The type of attention mechanism to apply. If `None`, no attention is applied.
-        - "se" for Squeeze-and-excitation network
-        - "gsop" for Global Second-Order Pooling
-        - "fca" for Frequency Channel Attention Network
-        - "encnet" for context encoding module
-        - "eca" for Efficient channel attention for deep convolutional neural networks
-        - "ge" for Gather-Excite
-        - "gct" for Gated Channel Transformation
-        - "srm" for Style-based Recalibration Module
-        - "cbam" for Convolutional Block Attention Module
-        - "cat" for Learning to collaborate channel and temporal attention
-        from multi-information fusion
-        - "catlite" for Learning to collaborate channel attention
+            - "se" for Squeeze-and-excitation network
+            - "gsop" for Global Second-Order Pooling
+            - "fca" for Frequency Channel Attention Network
+            - "encnet" for context encoding module
+            - "eca" for Efficient channel attention for deep convolutional neural networks
+            - "ge" for Gather-Excite
+            - "gct" for Gated Channel Transformation
+            - "srm" for Style-based Recalibration Module
+            - "cbam" for Convolutional Block Attention Module
+            - "cat" for Learning to collaborate channel and temporal attention
+            from multi-information fusion
+            - "catlite" for Learning to collaborate channel attention
         from multi-information fusion (lite version, cat w/o temporal attention)
     pool_length : int, default=8
         The length of the window for the average pooling operation.

braindecode/models/{sleep_stager_eldele_2021.py → attn_sleep.py}

@@ -8,18 +8,19 @@ from copy import deepcopy
 
 import torch
 import torch.nn.functional as F
+from mne.utils import deprecated
 from torch import nn
 
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import CausalConv1d
 
 
-class SleepStagerEldele2021(EEGModuleMixin, nn.Module):
+class AttnSleep(EEGModuleMixin, nn.Module):
     """Sleep Staging Architecture from Eldele et al. (2021) [Eldele2021]_.
 
     .. figure:: https://raw.githubusercontent.com/emadeldeen24/AttnSleep/refs/heads/main/imgs/AttnSleep.png
         :align: center
-        :alt: SleepStagerEldele2021 Architecture
+        :alt: AttnSleep Architecture
 
     Attention based Neural Net for sleep staging as described in [Eldele2021]_.
     The code for the paper and this model is also available at [1]_.
@@ -533,3 +534,12 @@ class _PositionwiseFeedForward(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Implements FFN equation."""
         return self.w_2(self.dropout(self.activate(self.w_1(x))))
+
+
+@deprecated(
+    "`SleepStagerEldele2021` was renamed to `AttnSleep` in v1.12 to follow original author's name; this alias will be removed in v1.14."
+)
+class SleepStagerEldele2021(AttnSleep):
+    """Deprecated alias for SleepStagerEldele2021."""
+
+    pass

braindecode/models/ctnet.py

@@ -39,7 +39,7 @@ class CTNet(EEGModuleMixin, nn.Module):
     The architecture consists of three main components:
 
     1. **Convolutional Module**:
-        - Apply EEGNetV4 to perform some feature extraction, denoted here as
+        - Apply :class:`EEGNet` to perform some feature extraction, denoted here as
         _PatchEmbeddingEEGNet module.
 
     2. **Transformer Encoder Module**:

braindecode/models/deep4.py

@@ -19,9 +19,13 @@ from braindecode.modules import (
 class Deep4Net(EEGModuleMixin, nn.Sequential):
     """Deep ConvNet model from Schirrmeister et al (2017) [Schirrmeister2017]_.
 
-     .. figure:: https://onlinelibrary.wiley.com/cms/asset/fc200ccc-d8c4-45b4-8577-56ce4d15999a/hbm23730-fig-0001-m.jpg
+    :bdg-success:`Convolution`
+
+    .. figure:: https://onlinelibrary.wiley.com/cms/asset/fc200ccc-d8c4-45b4-8577-56ce4d15999a/hbm23730-fig-0001-m.jpg
         :align: center
-        :alt: CTNet Architecture
+        :alt: Deep4Net Architecture
+        :width: 600px
+
 
     Model described in [Schirrmeister2017]_.
 

braindecode/models/deepsleepnet.py

@@ -8,14 +8,128 @@ from braindecode.models.base import EEGModuleMixin
 
 
 class DeepSleepNet(EEGModuleMixin, nn.Module):
-    """Sleep staging architecture from Supratak et al. (2017) [Supratak2017]_.
+    """DeepSleepNet from Supratak et al. (2017) [Supratak2017]_.
 
-     .. figure:: https://raw.githubusercontent.com/akaraspt/deepsleepnet/refs/heads/master/img/deepsleepnet.png
+    :bdg-success:`Convolution` :bdg-info:`Recurrent`
+
+    .. figure:: https://raw.githubusercontent.com/akaraspt/deepsleepnet/master/img/deepsleepnet.png
         :align: center
         :alt: DeepSleepNet Architecture
+        :width: 700px
+
+    .. rubric:: Architectural Overview
+
+    DeepSleepNet couples **dual-path convolution neural network representation learning** with
+    **sequence residual learning** via bidirectional LSTMs.
+
+    The network have:
+
+    - (i) learns complementary, time-frequency features from each
+      30-s epoch using **two parallel CNNs** (small vs. large first-layer filters), then
+    - (ii) models **temporal dependencies across epochs** using **two-layer BiLSTMs**
+      with a **residual shortcut** from the CNN features, and finally
+    - (iii) outputs per-epoch sleep stages. This design encodes both
+      epoch-local patterns and longer-range transition rules used by human scorers.
+
+    In term of implementation:
+
+    - (i) :class:`_RepresentationLearning` two CNNs extract epoch-wise features
+      (small-filter path for temporal precision; large-filter path for frequency precision);
+    - (ii) :class:`_SequenceResidualLearning` stacked BiLSTMs with peepholes + residual shortcut
+      inject temporal context while preserving CNN evidence;
+    - (iii) :class:`_Classifier` linear readout (softmax) for the five sleep stages.
+
+    .. rubric:: Macro Components
+
+    - :class:`_RepresentationLearning` **(dual-path CNN → epoch feature)**
+
+       - *Operations.*
+       - **Small-filter CNN** 4 times:
+       - :class:`~torch.nn.Conv1d`
+       - :class:`~torch.nn.BatchNorm1d`
+       - :class:`~torch.nn.ReLU`
+       - :class:`~torch.nn.MaxPool1d` after.
+       First conv uses **filter length ≈ Fs/2** and **stride ≈ Fs/16** to emphasize *timing* of graphoelements.
+    - **Large-filter CNN**:
+        - Same stack but first conv uses **filter length ≈ 4·Fs** and
+        - **stride ≈ Fs/2** to emphasize *frequency* content.
+    - Outputs from both paths are **concatenated** into the epoch embedding ``a_t``.
+
+    - *Rationale.*
+      Two first-layer scales provide a **learned, dual-scale filter bank** that trades
+      temporal vs. frequency precision without hand-crafted features.
+
+    - :class:`_SequenceResidualLearning` (:class:`~torch.nn.BiLSTM` **context + residual fusion)**
+
+        - *Operations.*
+        - **Two-layer BiLSTM** with **peephole connections** processes the sequence of epoch embeddings
+          ``{a_t}`` forward and backward; hidden states from both directions are **concatenated**.
+        - A **shortcut MLP** (fully connected + :class:`~torch.nn.BatchNorm1d` + :class:`~torch.nn.ReLU`) projects ``a_t`` to the BiLSTM output
+          dimension and is **added** (residual) to the :class:`~torch.nn.BiLSTM` output at each time step.
+        - *Role.* Encodes **stage-transition rules** and smooths predictions over time while preserving
+          salient CNN features via the residual path.
+
+    - :class:`_Classifier` **(epoch-wise prediction)**
+
+        - *Operations.*
+        - :class:`~torch.nn.Linear` to produce per-epoch class probabilities.
 
-    Convolutional neural network and bidirectional-Long Short-Term
-    for single channels sleep staging described in [Supratak2017]_.
+    Original training uses two-step optimization: CNN pretraining on class-balanced data,
+    then end-to-end fine-tuning with sequential batches.
+
+    .. rubric:: Convolutional Details
+
+    - **Temporal (where time-domain patterns are learned).**
+
+    Both CNN paths use **1-D temporal convolutions**. The *small-filter* path (first kernel ≈ Fs/2,
+    stride ≈ Fs/16) captures *when* characteristic transients occur; the *large-filter* path
+    (first kernel ≈ 4·Fs, stride ≈ Fs/2) captures *which* frequency components dominate over the
+    epoch. Deeper layers use **small kernels** to refine features with fewer parameters, interleaved
+    with **max pooling** for downsampling.
+
+    - **Spatial (how channels are processed).**
+    The original model operates on **single-channel** raw EEG; convolutions therefore mix only
+    along time (no spatial convolution across electrodes).
+
+    - **Spectral (how frequency information emerges).**
+    No explicit Fourier/wavelet transform is used. The **large-filter path** serves as a
+    *frequency-sensitive* analyzer, while the **small-filter path** remains *time-sensitive*,
+    together functioning as a **two-band learned filter bank** at the first layer.
+
+    .. rubric:: Attention / Sequential Modules
+
+    - **Type.** **Bidirectional LSTM** (two layers) with **peephole connections**; forward and
+      backward streams are independent and concatenated.
+    - **Shapes.** For a sequence of ``N`` epochs, the CNN produces ``{a_t} ∈ R^{D}``;
+      BiLSTM outputs ``h_t ∈ R^{2H}``; the shortcut MLP maps ``a_t → R^{2H}`` to enable
+      **element-wise residual addition**.
+    - **Role.** Models **long-range temporal dependencies** (e.g., persisting N2 without visible
+      K-complex/spindles), stabilizing per-epoch predictions.
+
+
+    .. rubric:: Additional Mechanisms
+
+    - **Residual shortcut over sequence encoder.** Adds projected CNN features to BiLSTM outputs,
+      improving gradient flow and retaining discriminative content from representation learning.
+    - **Two-step training.**
+        - (i) **Pretrain** the CNN paths with class-balanced sampling;
+        - (ii) **fine-tune** the full network with sequential batches, using **lower LR** for CNNs and **higher LR** for the
+        sequence encoder.
+    - **State handling.** BiLSTM states are **reinitialized per subject** so that temporal context
+      does not leak across recordings.
+
+
+    .. rubric:: Usage and Configuration
+
+    - **Epoch pipeline.** Use **two parallel CNNs** with the first conv sized to **Fs/2** (small path)
+      and **4·Fs** (large path), with strides **Fs/16** and **Fs/2**, respectively; stack three more
+      conv blocks with small kernels, plus **max pooling** in each path. Concatenate path outputs
+      to form epoch embeddings.
+    - **Sequence encoder.** Apply **two-layer BiLSTM (peepholes)** over the sequence of embeddings;
+      add a **projection MLP** on the CNN features and **sum** with BiLSTM outputs (residual).
+      Finish with :class:`~torch.nn.Linear` per epoch.
+    - **Reference implementation.** See the official repository for a faithful implementation and
+      training scripts.
 
     Parameters
     ----------
@@ -32,7 +146,6 @@ class DeepSleepNet(EEGModuleMixin, nn.Module):
     drop_prob : float, default=0.5
         The dropout rate for regularization. Values should be between 0 and 1.
 
-
     References
     ----------
     .. [Supratak2017] Supratak, A., Dong, H., Wu, C., & Guo, Y. (2017).

braindecode/models/eegconformer.py

@@ -12,33 +12,129 @@ from braindecode.modules import FeedForwardBlock, MultiHeadAttention
 
 
 class EEGConformer(EEGModuleMixin, nn.Module):
-    """EEG Conformer from Song et al. (2022) from [song2022]_.
+    """EEG Conformer from Song et al. (2022) [song2022]_.
 
-     .. figure:: https://raw.githubusercontent.com/eeyhsong/EEG-Conformer/refs/heads/main/visualization/Fig1.png
+    :bdg-success:`Convolution` :bdg-info:`Small Attention`
+
+    .. figure:: https://raw.githubusercontent.com/eeyhsong/EEG-Conformer/refs/heads/main/visualization/Fig1.png
         :align: center
         :alt: EEGConformer Architecture
+        :width: 600px
+
+
+    .. rubric:: Architectural Overview
+
+    EEG-Conformer is a *convolution-first* model augmented with a *lightweight transformer
+    encoder*. The end-to-end flow is:
+
+    - (i) :class:`_PatchEmbedding` converts the continuous EEG into a compact sequence of tokens via a
+      :class:`ShallowFBCSPNet` temporal–spatial conv stem and temporal pooling;
+    - (ii) :class:`_TransformerEncoder` applies small multi-head self-attention to integrate
+      longer-range temporal context across tokens;
+    - (iii) :class:`_ClassificationHead` aggregates the sequence and performs a linear readout.
+      This preserves the strong inductive biases of shallow CNN filter banks while adding
+      just enough attention to capture dependencies beyond the pooling horizon [song2022]_.
+
+    .. rubric:: Macro Components
+
+    - :class:`_PatchEmbedding` **(Shallow conv stem → tokens)**
+
+        - *Operations.*
+        - A temporal convolution (`:class:torch.nn.Conv2d`) ``(1 x L_t)`` forms a data-driven "filter bank";
+        - A spatial convolution (`:class:torch.nn.Conv2d`) (n_chans x 1)`` projects across electrodes,
+          collapsing the channel axis into a virtual channel.
+        - **Normalization function** :class:`torch.nn.BatchNorm`
+        - **Activation function** :class:`torch.nn.ELU`
+        - **Average Pooling** :class:`torch.nn.AvgPool` along time (kernel ``(1, P)`` with stride ``(1, S)``)
+        -  final ``1x1`` :class:`torch.nn.Linear` projection.
+
+    The result is rearranged to a token sequence ``(B, S_tokens, D)``, where ``D = n_filters_time``.
+
+    *Interpretability/robustness.* Temporal kernels can be inspected as FIR filters;
+    the spatial conv yields channel projections analogous to :class:`ShallowFBCSPNet`’s learned
+    spatial filters. Temporal pooling stabilizes statistics and reduces sequence length.
+
+    - :class:`_TransformerEncoder` **(context over temporal tokens)**
+
+        - *Operations.*
+        - A stack of ``att_depth`` encoder blocks. :class:`_TransformerEncoderBlock`
+        - Each block applies LayerNorm :class:`torch.nn.LayerNorm`
+        - Multi-Head Self-Attention (``att_heads``) with dropout + residual :class:`MultiHeadAttention` (:class:`torch.nn.Dropout`)
+        - LayerNorm :class:`torch.nn.LayerNorm`
+        - 2-layer feed-forward (≈4x expansion, :class:`torch.nn.GELU`) with dropout + residual.
+
+    Shapes remain ``(B, S_tokens, D)`` throughout.
 
-    Convolutional Transformer for EEG decoding.
+    *Role.* Small attention focuses on interactions among *temporal patches* (not channels),
+    extending effective receptive fields at modest cost.
 
-    The paper and original code with more details about the methodological
-    choices are available at the [song2022]_ and [ConformerCode]_.
+    - :class:`ClassificationHead` **(aggregation + readout)**
 
-    This neural network architecture receives a traditional braindecode input.
-    The input shape should be three-dimensional matrix representing the EEG
-    signals.
+        - *Operations*.
+        - Flatten, :class:`torch.nn.Flatten` the sequence ``(B, S_tokens·D)`` -
+        - MLP (:class:`torch.nn.Linear` → activation (default: :class:`torch.nn.ELU`) → :class:`torch.nn.Dropout` → :class:`torch.nn.Linear`)
+        - final Linear to classes.
 
-         `(batch_size, n_channels, n_timesteps)`.
+    With ``return_features=True``, features before the last Linear can be exported for
+    linear probing or downstream tasks.
 
-    The EEG Conformer architecture is composed of three modules:
-        - PatchEmbedding
-        - TransformerEncoder
-        - ClassificationHead
+    .. rubric:: Convolutional Details
+
+    - **Temporal (where time-domain patterns are learned).**
+        The initial ``(1 x L_t)`` conv per channel acts as a *learned filter bank* for oscillatory
+        bands and transients. Subsequent **AvgPool** along time performs local integration,
+        converting activations into “patches” (tokens). Pool length/stride control the
+        token rate and set the lower bound on temporal context within each token.
+
+    - **Spatial (how electrodes are processed).**
+        A single conv with kernel ``(n_chans x 1)`` spans the full montage to learn spatial
+        projections for each temporal feature map, collapsing the channel axis into a
+        virtual channel before tokenization. This mirrors the shallow spatial step in
+        :class:`ShallowFBCSPNet` (temporal filters → spatial projection → temporal condensation).
+
+    - **Spectral (how frequency content is captured).**
+        No explicit Fourier/wavelet stage is used. Spectral selectivity emerges implicitly
+        from the learned temporal kernels; pooling further smooths high-frequency noise.
+        The effective spectral resolution is thus governed by ``L_t`` and the pooling
+        configuration.
+
+    .. rubric:: Attention / Sequential Modules
+
+    - **Type.** Standard multi-head self-attention (MHA) with ``att_heads`` heads over the token sequence.
+    - **Shapes.** Input/Output: ``(B, S_tokens, D)``; attention operates along the ``S_tokens`` axis.
+    - **Role.** Re-weights and integrates evidence across pooled windows, capturing dependencies
+      longer than any single token while leaving channel relationships to the convolutional stem.
+      The design is intentionally *small*—attention refines rather than replaces convolutional feature extraction.
+
+    .. rubric:: Additional Mechanisms
+
+    - **Parallel with ShallowFBCSPNet.** Both begin with a learned temporal filter bank,
+        spatial projection across electrodes, and early temporal condensation.
+        :class:`ShallowFBCSPNet` then computes band-power (via squaring/log-variance), whereas
+        EEG-Conformer applies BN/ELU and **continues with attention** over tokens to
+        refine temporal context before classification.
+
+    - **Tokenization knob.** ``pool_time_length`` and especially ``pool_time_stride`` set
+        the number of tokens ``S_tokens``. Smaller strides → more tokens and higher attention
+        capacity (but higher compute); larger strides → fewer tokens and stronger inductive bias.
+
+    - **Embedding dimension = filters.** ``n_filters_time`` serves double duty as both the
+        number of temporal filters in the stem and the transformer’s embedding size ``D``,
+        simplifying dimensional alignment.
+
+    .. rubric:: Usage and Configuration
+
+    - **Instantiation.** Choose ``n_filters_time`` (embedding size ``D``) and
+        ``filter_time_length`` to match the rhythms of interest. Tune
+        ``pool_time_length/stride`` to trade temporal resolution for sequence length.
+        Keep ``att_depth`` modest (e.g., 4–6) and set ``att_heads`` to divide ``D``.
+        ``final_fc_length="auto"`` infers the flattened size from PatchEmbedding.
 
     Notes
     -----
     The authors recommend using data augmentation before using Conformer,
     e.g. segmentation and recombination,
-    Please refer to the original paper and code for more details.
+    Please refer to the original paper and code for more details [ConformerCode]_.
 
     The model was initially tuned on 4 seconds of 250 Hz data.
     Please adjust the scale of the temporal convolutional layer,
@@ -47,7 +143,10 @@ class EEGConformer(EEGModuleMixin, nn.Module):
     .. versionadded:: 0.8
 
     We aggregate the parameters based on the parts of the models, or
-    when the parameters were used first, e.g. n_filters_time.
+    when the parameters were used first, e.g. ``n_filters_time``.
+
+    .. versionadded:: 1.1
+
 
     Parameters
     ----------

braindecode/models/eeginception_erp.py

@@ -15,12 +15,84 @@ from braindecode.modules import DepthwiseConv2d, Ensure4d, InceptionBlock
 class EEGInceptionERP(EEGModuleMixin, nn.Sequential):
     """EEG Inception for ERP-based from Santamaria-Vazquez et al (2020) [santamaria2020]_.
 
+    :bdg-success:`Convolution`
+
     .. figure:: https://braindecode.org/dev/_static/model/eeginceptionerp.jpg
         :align: center
         :alt: EEGInceptionERP Architecture
 
-    The code for the paper and this model is also available at [santamaria2020]_
-    and an adaptation for PyTorch [2]_.
+        Figure: Overview of EEG-Inception architecture. 2D convolution blocks and depthwise 2D convolution blocks include batch normalization, activation and dropout regularization. The kernel size is displayed for convolutional and average pooling layers.
+
+    .. rubric:: Architectural Overview
+
+    A two-stage, multi-scale CNN tailored to ERP detection from short (0-1000 ms) single-trial epochs. Signals are mapped through
+    * (i) :class:`_InceptionModule1` multi-scale temporal feature extraction plus per-branch spatial mixing;
+    * (ii) :class:`_InceptionModule2` deeper multi-scale refinement at a reduced temporal resolution; and
+    * (iii) :class:`_OutputModule` compact aggregation and linear readout.
+
+    .. rubric:: Macro Components
+
+    - :class:`_InceptionModule1` **(multi-scale temporal + spatial mixing)**
+
+        - *Operations.*
+        - `EEGInceptionERP.c1`: :class:`torch.nn.Conv2d` ``k=(64,1)``, stride ``(1,1)``, *same* pad on input reshaped to ``(B,1,128,8)`` → BN → activation → dropout.
+        - `EEGInceptionERP.d1`: :class:`torch.nn.Conv2d` (depthwise) ``k=(1,8)``, *valid* pad over channels → BN → activation → dropout.
+        - `EEGInceptionERP.c2`: :class:`torch.nn.Conv2d` ``k=(32,1)`` → BN → activation → dropout; then `EEGInceptionERP.d2` depthwise ``k=(1,8)`` → BN → activation → dropout.
+        - `EEGInceptionERP.c3`: :class:`torch.nn.Conv2d` ``k=(16,1)`` → BN → activation → dropout; then `EEGInceptionERP.d3` depthwise ``k=(1,8)`` → BN → activation → dropout.
+        - `EEGInceptionERP.n1`: :class:`torch.nn.Concat` over branch features.
+        - `EEGInceptionERP.a1`: :class:`torch.nn.AvgPool2d` ``pool=(4,1)``, stride ``(4,1)`` for temporal downsampling.
+
+    *Interpretability/robustness.* Depthwise `1 x n_chans` layers act as learnable montage-wide spatial filters per temporal scale; pooling stabilizes against jitter.
+
+    - :class:`_InceptionModule2` **(refinement at coarser timebase)**
+
+        - *Operations.*
+        - `EEGInceptionERP.c4`: :class:`torch.nn.Conv2d` ``k=(16,1)`` → BN → activation → dropout.
+        - `EEGInceptionERP.c5`: :class:`torch.nn.Conv2d` ``k=(8,1)`` → BN → activation → dropout.
+        - `EEGInceptionERP.c6`: :class:`torch.nn.Conv2d` ``k=(4,1)`` → BN → activation → dropout.
+        - `EEGInceptionERP.n2`: :class:`torch.nn.Concat` (merge C4-C6 outputs).
+        - `EEGInceptionERP.a2`: :class:`torch.nn.AvgPool2d` ``pool=(2,1)``, stride ``(2,1)``.
+        - `EEGInceptionERP.c7`: :class:`torch.nn.Conv2d` ``k=(8,1)`` → BN → activation → dropout; then `EEGInceptionERP.a3`: :class:`torch.nn.AvgPool2d` ``pool=(2,1)``.
+        - `EEGInceptionERP.c8`: :class:`torch.nn.Conv2d` ``k=(4,1)`` → BN → activation → dropout; then `EEGInceptionERP.a4`: :class:`torch.nn.AvgPool2d` ``pool=(2,1)``.
+
+    *Role.* Adds higher-level, shorter-window evidence while progressively compressing temporal dimension.
+
+    - :class:`_OutputModule` **(aggregation + readout)**
+
+        - *Operations.*
+        - :class:`torch.nn.Flatten`
+        - :class:`torch.nn.Linear` ``(features → 2)``
+
+    .. rubric:: Convolutional Details
+
+    - **Temporal (where time-domain patterns are learned).**
+    First module uses 1D temporal kernels along the 128-sample axis: ``64``, ``32``, ``16``
+    (≈500, 250, 125 ms at 128 Hz). After ``pool=(4,1)``, the second module applies ``16``,
+    ``8``, ``4`` (≈125, 62.5, 31.25 ms at the pooled rate). All strides are ``1`` in convs;
+    temporal resolution changes only via average pooling.
+
+    - **Spatial (how electrodes are processed).**
+    Depthwise convs with ``k=(1,8)`` span all channels and are applied **per temporal branch**,
+    yielding scale-specific channel projections (no cross-branch mixing until concatenation).
+    There is no full 2D mixing kernel; spatial mixing is factorized and lightweight.
+
+    - **Spectral (how frequency information is captured).**
+    No explicit transform; multiple temporal kernels form a *learned filter bank* over
+    ERP-relevant bands. Successive pooling acts as low-pass integration to emphasize sustained
+    post-stimulus components.
+
+    .. rubric:: Additional Mechanisms
+
+    - Every conv/depthwise block includes **BatchNorm**, nonlinearity (paper used grid-searched activation), and **dropout**.
+    - Two Inception stages followed by short convs and pooling keep parameters small (≈15k reported) while preserving multi-scale evidence.
+    - Expected input: epochs of shape ``(B,1,128,8)`` (time x channels as a 2D map) or reshaped from ``(B,8,128)`` with an added singleton feature dimension.
+
+    .. rubric:: Usage and Configuration
+
+    - **Key knobs.** Number of filters per branch; kernel lengths in both Inception modules; depthwise kernel over channels (typically ``n_chans``); pooling lengths/strides; dropout rate; choice of activation.
+    - **Training tips.** Use 0-1000 ms windows at 128 Hz with CAR; tune activation and dropout (they strongly affect performance); early-stop on validation loss when overfitting emerges.
+
+    .. rubric:: Implementation Details
 
     The model is strongly based on the original InceptionNet for an image. The main goal is
     to extract features in parallel with different scales. The authors extracted three scales
@@ -33,12 +105,9 @@ class EEGInceptionERP(EEGModuleMixin, nn.Sequential):
     The winners of BEETL Competition/NeurIps 2021 used parts of the
     model [beetl]_.
 
-    The model is fully described in [santamaria2020]_.
+    The code for the paper and this model is also available at [santamaria2020]_
+    and an adaptation for PyTorch [2]_.
 
-    Notes
-    -----
-    This implementation is not guaranteed to be correct, has not been checked
-    by original authors, only reimplemented from the paper based on [2]_.
 
     Parameters
     ----------