Math Foundations — Lecture 6

Attention score functions, self-attention, QKV projections, multi-head, positional encodings, residuals, layer norm, causal mask.

Attention score functions

All attention mechanisms compute a relevance score \(e\) between a query and a key, then softmax over keys.

Dot-product

\[e = \mathbf{s}^\top\mathbf{h}\]

Bilinear (general)

\[e = \mathbf{s}^\top W\,\mathbf{h}\]

Learnable \(W\) allows the model to weight dimensions differently.

Additive (Bahdanau)

\[e = \mathbf{b}^\top\tanh(W_1\mathbf{h}+W_2\mathbf{s})\]

Projects query and key into a shared space, then scores with a learned vector \(\mathbf{b}\).

Scaled dot-product (Vaswani et al.)

\[e = \frac{\mathbf{q}^\top\mathbf{k}}{\sqrt{d_k}}\]

For random unit-variance vectors in \(\mathbb{R}^{d_k}\), \(\mathbb{E}[\mathbf{q}^\top\mathbf{k}]=0\) but \(\mathrm{Var}[\mathbf{q}^\top\mathbf{k}]=d_k\). Dividing by \(\sqrt{d_k}\) restores unit variance → softmax stays in a non-saturated regime → healthier gradients.

Quick Check: If \(d_k=64\), what is the scaling denominator?

\(\sqrt{64}=8\)

\(64\)

\(\log 64\)

Self-attention

Every position attends to every other position in the same sequence:

\[e_{ij}=\mathbf{x}_i^\top\mathbf{x}_j,\quad \alpha_{ij}=\frac{e^{e_{ij}}}{\sum_k e^{e_{ik}}},\quad \mathbf{o}_i=\sum_j\alpha_{ij}\,\mathbf{x}_j\]

Output \(\mathbf{o}_i\) is a weighted average of all input vectors — weights determined by pairwise similarity.

Query, key, value projections

Instead of raw \(\mathbf{x}\), project into three roles:

\[\mathbf{q}_i=W^Q\mathbf{x}_i,\quad \mathbf{k}_j=W^K\mathbf{x}_j,\quad \mathbf{v}_j=W^V\mathbf{x}_j\]

\[\alpha_{ij}=\mathrm{softmax}_j\!\left(\frac{\mathbf{q}_i^\top\mathbf{k}_j}{\sqrt{d_k}}\right),\qquad \mathbf{t}_i=\sum_j\alpha_{ij}\,\mathbf{v}_j\]

In matrix form for all positions simultaneously:

\[\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V\]

where \(Q,K,V\in\mathbb{R}^{n\times d_k}\) (or \(d_v\) for values).

Multi-head attention

\[\mathrm{head}_h=\mathrm{Attention}(XW_h^Q,\;XW_h^K,\;XW_h^V)\] \[\mathrm{MultiHead}(X)=\mathrm{Concat}(\mathrm{head}_1,\dots,\mathrm{head}_H)\,W^O\]

\(H\) heads, each with \(d_k=d_\text{model}/H\). Heads can specialise (e.g. one attends to syntax, another to coreference). Output projection \(W^O\in\mathbb{R}^{d_\text{model}\times d_\text{model}}\) mixes them.

Quick Check: With \(d_\text{model}=512\) and \(H=8\) heads, each head operates in dimension:

\(512/8=64\)

\(512\)

Positional encodings

Self-attention is permutation-invariant. Add position information to inputs:

Sinusoidal (original Transformer)

\[\mathrm{PE}(t,2i)=\sin\!\left(\frac{t}{10000^{2i/d}}\right),\quad \mathrm{PE}(t,2i+1)=\cos\!\left(\frac{t}{10000^{2i/d}}\right)\]

Different frequencies per dimension; the model can learn to attend to relative positions.

Learned positional embeddings

A trainable vector per position index — simple and effective; limited to max seen length.

RoPE (rotary)

Rotates query/key pairs by position-dependent angles so that \(\mathbf{q}_m^\top\mathbf{k}_n\) depends only on \(m-n\). Preserves relative distance in cosine similarity.

Causal masking

For autoregressive decoding, position \(i\) must not see positions \(j>i\):

\[e_{ij}=\begin{cases}\mathbf{q}_i^\top\mathbf{k}_j/\sqrt{d_k}&j\le i\\-\infty&j>i\end{cases}\]

After softmax, \(\alpha_{ij}=0\) for future positions. Enables parallel training of all positions while preserving the autoregressive property.

Residual connections

\[\mathbf{x}' = \mathbf{x} + \mathrm{SubLayer}(\mathbf{x})\]

Gradient flows directly through the addition, avoiding degradation in deep stacks. Also stabilises training by keeping the input "highway" intact.

Layer normalisation

\[\mathrm{LayerNorm}(\mathbf{x})=\gamma\odot\frac{\mathbf{x}-\mu}{\sigma+\epsilon}+\beta\]

\(\mu,\sigma\)Mean and std computed over the feature dimension of a single example.

\(\gamma,\beta\)Learned scale and shift (element-wise).

Normalises activations per token; independent of batch size (unlike batch norm).

Quick Check: Layer norm computes statistics over:

The feature dimension of one example

The batch dimension

Transformer encoder block (summary)

\[\mathbf{a}=\mathrm{LayerNorm}\bigl(\mathbf{x}+\mathrm{MultiHeadAttn}(\mathbf{x})\bigr)\] \[\mathbf{h}=\mathrm{LayerNorm}\bigl(\mathbf{a}+\mathrm{FFN}(\mathbf{a})\bigr)\]

where \(\mathrm{FFN}(\mathbf{a})=\mathrm{ReLU}(\mathbf{a}W_1+\mathbf{b}_1)W_2+\mathbf{b}_2\). Stack \(N\) such blocks (e.g. 6 in the original paper).

Cross-attention (decoder)

Decoder block has three sub-layers: masked self-attention, cross-attention, FFN.

\[Q=\text{decoder hidden},\quad K=V=\text{encoder output}\]

Queries from the decoder, keys and values from the encoder — this replaces the single bottleneck vector of vanilla seq2seq.

Quick Check: In cross-attention, where do Q, K, V come from?

Q from decoder; K,V from encoder

All three from encoder

All three from decoder