peal.architectures.module_blocks.TransformerBlock

class peal.architectures.module_blocks.TransformerBlock

Bases: Sequential

Pre-activation transformer block: attention and position-wise MLP, each residual.

Sub-block 1 is SkipConnection(SelfAttentionLayer -> LayerNorm); sub-block 2 is SkipConnection(1x1 Conv1d over the embedding axis -> activation -> LayerNorm). Inputs and outputs have shape (B, T, embedding_dim).

Parameters:

• embedding_dim (int) – Token embedding width.

• num_heads (int) – Number of self-attention heads.

• activation (type of nn.Module) – Activation class for the feed-forward sub-block.

• use_masking (bool, optional) – Whether the attention layer applies a causal mask. Default False.

Methods

__init__(embedding_dim, num_heads, activation, use_masking=False)

The __init__ method initializes the TransformerBlock class.

Parameters:

• embedding_dim (int) – The number of input channels.

• num_heads (int) – The number of self-attention heads.

• activation (int) – The activation function.

• use_masking (bool, optional) – Whether to use masking. Defaults to False.