The Mathematics of Transformers: Attention, Projections & Softmax Geometry

The Transformer architecture eliminated recurrence and convolutions in favor of pure self-attention mechanisms. While intuitive conceptually, its mathematical formulation reveals deep geometric properties.

---

1. Scaled Dot-Product Attention

Given an input sequence represented as a matrix $X \in \mathbb{R}^{n \times d}$, we project $X$ into three distinct vector spaces using learned weight matrices: $Q = X W_Q, \quad K = X W_K, \quad V = X W_V$ where $W_Q, W_K \in \mathbb{R}^{d \times d_k}$ and $W_V \in \mathbb{R}^{d \times d_v}$.

The attention formula is expressed as: $\text{Attention}(Q, K, V) = \text{softmax}\left( \frac{Q K^T}{\sqrt{d_k}} \right) V$

---

2. Why Scale by $\frac{1}{\sqrt{d_k}}$?

Consider two random vectors $q, k \in \mathbb{R}^{d_k}$ whose components are independent random variables with mean $0$ and variance $1$: $q_i, k_i \sim \mathcal{N}(0, 1)$

Calculating the expectation and variance of their dot product: $\mathbb{E}[S] = 0, \quad \text{Var}(S) = d_k$

As $d_k$ grows large (e.g., $d_k = 128$), the variance of the dot product scales linearly with $d_k$. Extremely large positive or negative values push the $\text{softmax}$ function into regions with near-zero gradients (vanishing gradients). Dividing by $\sqrt{d_k}$ normalizes the variance back to $1$.