Self-attention combines inputs directly, with input-dependent weights

Each \(\boldsymbol{y}_\ell\) combines information from all rows of \(X\).

The learned projection matrices stay fixed; the combination weights are input-dependent.

Self-attention removes the sequential chain within a layer

Shorter paths for learning; parallel operations across positions during training.

Column vectors become transposed rows in a sequence matrix

\[X=\begin{bmatrix}\boldsymbol{x}_1^{\mathsf T}\\\vdots\\\boldsymbol{x}_L^{\mathsf T}\end{bmatrix},\qquad \boldsymbol{q}_\ell=W_Q\boldsymbol{x}_\ell,\qquad Q=XW_Q^{\mathsf T}\]

Likewise \(K=XW_K^{\mathsf T}\) and \(V=XW_V^{\mathsf T}\). Example: four tokens, three features per vector.

Each query–key inner product fills one score-matrix entry

\[S=\frac{QK^{\mathsf T}}{\sqrt{d_k}},\qquad A=\operatorname{softmax}_{\rm row}(S),\qquad a_{\ell j}=\frac{e^{s_{\ell j}}}{\sum_{r=1}^{L}e^{s_{\ell r}}}\]

Dot products give scores; row softmax gives weights.

Attention builds an input-dependent weighted combination

For each input column vector, form a value: \(\boldsymbol{v}_j=f_v(\boldsymbol{x}_j)=W_V\boldsymbol{x}_j\).

\[\boldsymbol{y}_\ell=\sum_{j=1}^{L}a_{\ell j}\boldsymbol{v}_j, \qquad a_{\ell j}\geq0,\quad\sum_{j=1}^{L}a_{\ell j}=1\]

• The score for pair \((\ell,j)\) is a function of both inputs: \(s_{\ell j}=f_s(\boldsymbol{x}_\ell,\boldsymbol{x}_j)\).
• \(a_{\ell j}=\operatorname{softmax}_j(s_{\ell1},\ldots,s_{\ell L})\): normalize against all candidate inputs.

How do queries and keys turn these input pairs into useful weights?

Attention is a soft lookup in a key–value table

Dictionary lookup: an exact key selects one stored value.

Attention: a query scores every key, then blends the paired values.

Query: features to match. Key: matching features. Value: features to combine.

Softmax turns matching scores into relative combination weights

\[s_{\ell j}=\frac{\boldsymbol{q}_\ell^{\mathsf T}\boldsymbol{k}_j}{\sqrt{d_k}}, \qquad a_{\ell j}=\frac{\exp(s_{\ell j}-m_\ell)}{\sum_{k=1}^{L}\exp(s_{\ell k}-m_\ell)}, \qquad m_\ell=\max_{k} s_{\ell k}\]

• A larger dot product gives a larger weight relative to the other scores.
• Subtracting the row maximum preserves the weights and avoids large exponentials.

These weights are not probabilities that a token is true or important.

Next: calculate one complete weighted combination.

One weighted combination uses values, not token IDs

Synthetic one-feature query \(\boldsymbol{q}=[1]\); the query is at the.

This is a context feature vector. A later vocabulary head predicts the next token.

Matrix multiplication computes all output states together

Stack token representations as rows of \(X\in\mathbb{R}^{L\times d}\).

\[S=QK^{\mathsf T}/\sqrt{d_k},\qquad A=\operatorname{softmax}_{\rm row}(S),\qquad Y=AV\]

One row of \(A\) gives the input-dependent weights for one output row of \(Y\).

Divide by \(\sqrt{d_k}\) to keep score fluctuations comparable

Assume independent query/key entries with mean \(0\) and variance \(1\).

\[z=\boldsymbol{q}^{\mathsf T}\boldsymbol{k}=\sum_{i=1}^{d_k}q_i k_i, \qquad \mathbb{E}[z]=0,\qquad \operatorname{Var}(z)=d_k\]

Positive and negative terms partly cancel. Variances add; standard deviations grow as \(\sqrt{d_k}\).

Dividing by \(\sqrt{d_k}\) keeps the score standard deviation at \(1\), for any width.

Avoid near-one-hot attention caused merely by width: saturated softmax has small gradients.

Cross-attention gets queries and context from different sequences

Query-side inputs \(X\in\mathbb{R}^{L_q\times d_x}\); context inputs \(C\in\mathbb{R}^{L_c\times d_c}\).

\[Q=XW_Q^{\mathsf T},\qquad K=CW_K^{\mathsf T},\qquad V=CW_V^{\mathsf T}\]

\[\operatorname{CA}(X,C)= \operatorname{softmax}_{\rm row}\!\left(\frac{(XW_Q^{\mathsf T})(CW_K^{\mathsf T})^{\mathsf T}}{\sqrt{d_k}}\right)(CW_V^{\mathsf T})\]

• Example: question tokens query representations of a passage.
• Weight matrix: \(L_q\times L_c\). Output: one row per query, \(L_q\times d_v\).

Self-attention is the case \(C=X\). The same soft-lookup mechanism applies.

A causal mask removes future positions before softmax

\[A=\operatorname{softmax}_{\rm row}(S+M),\qquad Y=AV\]

Self-attention without positions is permutation equivariant

For any permutation matrix \(P\), with no mask or positional encoding:

\[\operatorname{SA}(PX)=P\operatorname{SA}(X)\]

The outputs are the same, just reordered.

Useful for an unordered set: changing the storage order should not change what each item means.

A sequence needs order beyond an unordered collection

Focus on report: its relationship to before has changed.

Without a mask or positions, self-attention gives report the same output vector in both sentences—just in a different row.

We want token representations to reflect order-dependent relationships.

Causal masks restrict access; position signals describe location

A fixed causal mask breaks arbitrary permutation equivariance: reordering tokens changes which context they may read.

The mask does not explicitly encode token indices or pairwise offsets.

Next: add position features, or rotate query/key features using their positions.

Additive positions combine token features and location features

Look up the token’s embedding, then add a vector describing its position:

\[\boldsymbol{x}_\ell=E[w_\ell,:]^{\mathsf T}+\boldsymbol{p}_\ell,\qquad \boldsymbol{x}_\ell,\boldsymbol{p}_\ell\in\mathbb{R}^{d}\]

Fixed sinusoids: \(p_{\ell,2i}=\sin(\ell\omega_i)\), \(p_{\ell,2i+1}=\cos(\ell\omega_i)\), \(\omega_i=10000^{-2i/d}\).

Pair index \(i=0,\ldots,d/2-1\); even width \(d\).

Same token “the” at positions 0 and 2: \(E[\text{the},:]^{\mathsf T}+\boldsymbol{p}_0\ne E[\text{the},:]^{\mathsf T}+\boldsymbol{p}_2\).

RoPE rotates query and key pairs by their positions

Split each query/key into 2-D feature pairs; rotate every pair by a position-dependent angle.

\[R(\alpha)=\begin{bmatrix}\cos\alpha&-\sin\alpha\\\sin\alpha&\cos\alpha\end{bmatrix},\qquad \widetilde{\boldsymbol{q}}_\ell^{(i)}=R(\ell\theta_i)\boldsymbol{q}_\ell^{(i)}\]

\[[\underbrace{q_{\ell,0},q_{\ell,1}}_{\text{pair 0: }\ell\theta_0},\quad \underbrace{q_{\ell,2},q_{\ell,3}}_{\text{pair 1: }\ell\theta_1},\quad\ldots]^{\mathsf T}, \qquad \theta_i=10000^{-2i/d_k}\]

Fast and slow pairs cover different distance scales. Keys use the same rotations.

RoPE compares positions through angle differences

Moving both tokens one step rotates both arrows by the same extra angle: their angle difference stays the same.

\[s\theta_i-\ell\theta_i=(s-\ell)\theta_i,\qquad (\widetilde{\boldsymbol{q}}_\ell^{(i)})^{\mathsf T}\widetilde{\boldsymbol{k}}_s^{(i)} =(\boldsymbol{q}_\ell^{(i)})^{\mathsf T}R((s-\ell)\theta_i)\boldsymbol{k}_s^{(i)}\]

The full dot product adds the matches from all feature pairs.

Used in Llama 2: RoPE is itself a positional encoding, an alternative to adding absolute position vectors.

Multiple heads learn different ways to combine context

\[Y=\operatorname{Concat}(Y_1,\ldots,Y_b)W_O^{\mathsf T}\]

Each head has its own query/key/value projections; every head obeys the allowed context.

Next: combine attention with residual paths, normalization and a token-wise MLP.

Implementation detail: padding masks exclude absent tokens

Two messages share a batch:

• Attention key mask: padded positions cannot supply values to real queries.
• Loss mask: padded targets contribute no training loss.

A causal mask still applies within each real prefix.