The core mechanism of the Transformer is Scaled Dot-Product Attention. It operates on three matrices: Query (Q), Key (K), and Value (V). These matrices are linear projections of the input embeddings. For a given token, its Query vector is matched against the Key vectors of all other tokens in the sequence to produce attention scores. These scores are then used to create a weighted sum of the Value vectors.

The Q, K, and V matrices are derived by multiplying the input embedding matrix $X$ with learned weight matrices $W_Q$, $W_K$, and $W_V$. This allows the model to project the inputs into different representation subspaces, optimizing them for the task of calculating attention.

The scaling factor, $\frac{1}{\sqrt{d_k}}$, is critical. For large values of $d_k$, the dot products $QK^T$ can grow very large in magnitude. This pushes the softmax function into regions where its gradients are extremely small, effectively vanishing and stalling the learning process. Scaling mitigates this by keeping the variance of the dot products at 1, ensuring stable gradients.

Instead of performing a single attention function, Transformers employ Multi-Head Attention (MHA). This mechanism runs the Scaled Dot-Product Attention process multiple times in parallel. Each parallel run, or "head," uses different, learned linear projections for Q, K, and V. This allows the model to jointly attend to information from different representation subspaces at different positions.

The outputs of the individual heads are concatenated and then passed through another learned linear projection to produce the final output. This allows the model to integrate the information captured by the different heads.

The self-attention mechanism is permutation-invariant; it does not inherently understand the order of tokens. To address this, positional information is injected. While the original Transformer used sinusoidal functions, modern architectures often use more advanced techniques.

Rotary Positional Embeddings (RoPE) are a popular choice. Instead of adding positional information to the embeddings, RoPE modifies the Query and Key vectors by applying a rotation matrix whose angle depends on the token's absolute position. When the dot product is taken between a rotated Query and a rotated Key, the result depends only on their relative positions, elegantly encoding relative distance directly into the self-attention mechanism.

Layer Normalization is another key component for stabilizing training. Its placement within the Transformer block has evolved. The original architecture used Post-LayerNorm, applying normalization after the residual connection. However, this can lead to unstable training in very deep models as the output magnitudes can explode.

Let's check your understanding of these core mechanics.

Understanding these mathematical underpinnings is essential for anyone working with or fine-tuning large language models.