Standard Attention

Standard scaled dot-product multi-head attention.

Implements the original Transformer attention mechanism with row-wise softmax normalization. Complexity is O(n^2 * d) during training and O(n) per-step inference with a KV cache. Provides full token-to-token routing with perfect expressiveness.

Reference:

Vaswani et al. (2017), “Attention Is All You Need”, arXiv:1706.03762.

class src.model.attention.standard.StandardAttention(*args: Any, **kwargs: Any)

Bases: Module

Standard multi-head scaled dot-product attention (Vaswani et al. 2017).

Projects the input into query, key, and value tensors, computes scaled dot-product attention scores, applies row-wise softmax, and aggregates values weighted by the resulting attention distribution. Supports both encoder (bidirectional) and decoder (causal) modes.

Complexity:

Training: O(n^2 * d). Inference: O(n) per step with KV cache.

Reference:

Vaswani et al. (2017), “Attention Is All You Need”, arXiv:1706.03762.

Parameters:

config – Model configuration object with attributes hidden_size, num_heads, dropout, use_bitnet, and optionally mode ("encoder" or "decoder").

hidden_size

Dimensionality of the input and output embeddings.

num_heads

Number of parallel attention heads.

head_dim

Dimensionality of each attention head (hidden_size // num_heads).

scale

Scaling factor 1 / sqrt(head_dim) applied to dot products.

q_proj

Linear (or BitLinear) projection for queries.

k_proj

Linear (or BitLinear) projection for keys.

v_proj

Linear (or BitLinear) projection for values.

out_proj

Linear (or BitLinear) output projection.

dropout

Dropout layer applied to attention weights.

mode

"encoder" for bidirectional attention, "decoder" for causal (upper-triangular) masking.

Raises:

ValueError – If hidden_size is not divisible by num_heads.

forward(x: torch.Tensor, logical_layer_idx: int | None = None) → torch.Tensor

Compute standard multi-head attention.

Parameters:

• x – Input tensor of shape (batch_size, seq_len, hidden_size).

• logical_layer_idx – Logical layer index (unused; accepted for interface compatibility with other attention modules).

Returns:

Output tensor of shape (batch_size, seq_len, hidden_size).