Attention & Transformer

Full Transformer model with encoder and decoder stacks.

Post-norm layers with ReLU feed-forward networks by default (see the normFirst and activation options). Unlike PyTorch's nn.Transformer, the encoder and decoder stacks have no final LayerNorm unless you ask for it with options.finalNorm. Parameters follow the dtype option (default: the global default dtype, float32 unless changed).

Multi-Head Attention mechanism.

Allows the model to jointly attend to information from different representation subspaces at different positions. This is the core building block of Transformers.

Mathematical Formulation:

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * VMultiHead(Q, K, V) = Concat(head_1, ..., head_h) * W_Owhere head_i = Attention(Q * W_Q^i, K * W_K^i, V * W_V^i)

Parameters are float32 unless the dtype option (or the global default dtype) says otherwise. The layer computes in the parameter dtype and casts its inputs to it. A GradTensor input gives a GradTensor; plain tensors give a GradTensor that tracks the weights while they require grad and gradient tracking is on, and a plain Tensor otherwise (inside noGrad() or with frozen weights).

Sinusoidal Positional Encoding.

Adds sinusoidal positional information to input embeddings using the formula from "Attention Is All You Need":

PE(pos, 2i)   = sin(pos / 10000^(2i/d_model))PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

For an odd dModel the last column pairs with the same frequency as the one before it. The table is stored as the pe buffer and cast to the input's float dtype on each call; integer inputs are cast to float32.

Transformer Decoder, a stack of N decoder layers.

The first layer is the decoderLayer you pass in; the other N - 1 layers are independent copies of it (see TransformerDecoderLayer.clone).

Transformer Decoder Layer.

A single layer of the Transformer decoder, consisting of: 1. Masked multi-head self-attention 2. Add & Norm 3. Multi-head cross-attention (attending to encoder output) 4. Add & Norm 5. Feed-forward network 6. Add & Norm

By default the self-attention is causal: position i only attends to positions up to i. Pass tgtMask to forward to use a different mask. The layout is post-norm with a ReLU feed-forward network unless normFirst and activation say otherwise.

Parameters are float32 unless the dtype option (or the global default dtype) says otherwise; inputs of other dtypes are cast to the parameter dtype. Gradient tracking follows the rule of MultiheadAttention.

Transformer Encoder, a stack of N encoder layers.

The first layer is the encoderLayer you pass in; the other N - 1 layers are independent copies of it (see TransformerEncoderLayer.clone), so every layer starts from the same weights and then trains separately.

Transformer Encoder Layer.

A single layer of the Transformer encoder, consisting of: 1. Multi-head self-attention 2. Add & Norm (residual connection + layer normalization) 3. Feed-forward network (FFN) 4. Add & Norm

Uses the post-norm layout and a ReLU feed-forward network by default, like PyTorch. Pass activation: 'gelu' and normFirst: true for a GELU network and the pre-norm layout. Parameters are float32 unless the dtype option (or the global default dtype) says otherwise; inputs of other dtypes are cast to the parameter dtype. Gradient tracking follows the rule of MultiheadAttention.