braindecode.modules.MultiHeadAttention#

class braindecode.modules.MultiHeadAttention(emb_size, num_heads, dropout=0.0, scale=None, head_dim=None, qk_norm=False)[source]#

Multi-head self-attention block.

Uses F.scaled_dot_product_attention for optimized attention kernels (flash-attention on CUDA, memory-efficient on other devices).

Parameters:

emb_size (int) – The embedding dimension.
num_heads (int) – Number of attention heads. Must evenly divide emb_size when head_dim is None.
dropout (float, optional) – Dropout probability applied to attention weights. Default: 0.0.
scale (float | None) – Multiplier applied to attention scores before softmax. When None, PyTorch’s default head_dim ** -0.5 is used.
head_dim (int | None) – Width of each head. When None, emb_size // num_heads; otherwise the query/key/value projections are num_heads * head_dim wide.
qk_norm (bool) – Apply a LayerNorm over head_dim to the queries and keys of each head (as in LaBraM and Du-IN). Default False.

Examples

>>> import torch
>>> from braindecode.modules import MultiHeadAttention
>>> module = MultiHeadAttention(emb_size=32, num_heads=4, dropout=0.1)
>>> inputs = torch.randn(2, 10, 32)
>>> outputs = module(inputs)
>>> outputs.shape
torch.Size([2, 10, 32])

Methods

forward(x, mask=None)[source]#

Forward pass.

Parameters:

x (Tensor) – Input tensor of shape (batch, seq, emb_size).
mask (Optional[Tensor]) – Attention mask following the convention of torch.nn.functional.scaled_dot_product_attention(): for boolean masks True means the position takes part in attention and False that it is ignored (the opposite of the key_padding_mask of torch.nn.MultiheadAttention); for float masks the values are added to the attention scores before softmax.

Return type:

Tensor