Attention

The Encoder: Attention

Big picture: Attention

• The inner products \[ \mathbf{q}^T\mathbf{k}_t \] are interpreted as The degree to which token \(t\in V\) is important for predicting the current token \(q\).
• \(\mathbf{q}^T\mathbf{k}_t\) is a scalar number. Taken over all \(t\), we obtain a distribution over the context tokens that is used to combine the value vectors

Attention function

• The attention function \(A(Q,K,V)\) is typically caculated in steps. First, we calculate the attention scores as a similarity measure between the query vectors and the key vectors. \[ S_{i,j} = \dfrac{\langle \mathbf{q}_i, \mathbf{k}_j\rangle}{\sqrt{d_k}} = \dfrac{\mathbf{q}_i^T\mathbf{k}_j}{\sqrt{d_k}} \]

• $ _i$: the \(i^{th}\) query vector and \(\mathbf{k}_j\): \(j^{th}\) key vector

• \(d_k\) is the dimensionality of the key vectors

• \(\langle \mathbf{q}_i, \mathbf{k}_j\rangle\) denotes the dot product.

• The scaling factor \(\frac{1}{\sqrt{d_k}}\) is introduced to mitigate the effect of increasing dimensionality (See below)

• The attention scores \(S_{i,j}\) are then normalized using the softmax function to produce the attention weights: \[ a_{ij} = \dfrac{e^{S_{i,j}}}{\sum_{k=1}^n e^{S_{i,k}}} \]

• The \(a_{ij}\) indicate the relative importance of the \(j^{th}\) key vector to the \(i^{th}\) query vector.

Attention function

• The output of the attention mechanism is computed as a weighted sum of the value vectors: \[ \mathbf{z}_{i} = \sum_{j=1}^n a_{ij}\mathbf{v}_j \]

• \(\mathbf{v}_j\) is the \(j^{th}\) value vector

• \(\mathbf{z}_{i}\) is the output vector corresponding to the \(i^{th}\) query vector.

• We can express this in matrix form as

\[ A(\mathbf{Q},\mathbf{K},\mathbf{V})= \mathrm{softmax}\left(\dfrac{\mathbf{QK}^T}{\sqrt{d_k}}\right) \]

The attention mapping is

• linear, allowing it to combine different sets of value vectors in a controlled manner
• Invariant Under Permutations: The order of the tokens does not effect the result of the calculations (which is why the positional encodings are needed)
• Commutative: If multiple heads are applied, the order of application does not matter (see below)

\(W^K\) versus \(W^Q\)

The Functional Roles: a useful intuition is that they act as two different “filters” applied to the same input data:

• \(W^K\) (The Key Matrix): intuitively, learns to describe a token so it can be “found” by relevant queries.

• \(W^Q\) (The Query Matrix): intuitively, learns to encode what the current token needs from its context.

• This is a simplifying explanation — \(W^Q\) and \(W^K\) are trained jointly and don’t have separately interpretable roles in general.

• Empirically, it has been shown that different attention heads may specialize to recognizable syntactic/semantic relations, e.g. heads that reliably attend from verbs to their direct objects, or from determiners to nouns.¹

The scaling factor

• The scaling factor \(\frac{1}{\sqrt{d_k}}\) is needed to keep the gradients in a good range
• We will derive a simplified/intuitive explanation of why this tends to work
• Assume \(q_1, q_2, \ldots, q_{d_k}, k_1,k_2,\ldots, k_{d_k}\) are mutually independent random variables, each with mean 0 and variance 1. with dimension \(d_k=64\)
• Thus, \(\mathbf{q}_i \perp \mathbf{k}_j\) (query and key are independent of each other): this implies that \(\mathbb{E}[\mathbf{q}_i \mathbf{k}_j]=0\) and that \(\mathbb{E}[\mathbf{q}_i^2 \mathbf{k}_j^2]=\mathbb{E}[\mathbf{q}_i^2]\mathbb{E}[\mathbf{k}_j^2]\)

Proof

The dot product \(q \cdot k\) for \(d_k\) dimensions is defined as:

\[ S = \sum_{i=1}^{d_k} q_i k_i \]

To find the variance of \(S\), we first look at the variance of a single term \(q_i k_i\):\[Var(q_i k_i) = E[(q_i k_i)^2] - (E[q_i k_i])^2\]Mean of the product: Since \(q_i\) and \(k_i\) are independent, \(E[q_i k_i] = E[q_i]E[k_i] = 0 \cdot 0 = 0\).

• Expected value of the square: \(E[q_i^2 k_i^2] = E[q_i^2]E[k_i^2]\). Since \(Var(X) = E[X^2] - (E[X])^2\), and \(E[X]=0\), then \(E[X^2] = Var(X) = 1\).
• Thus, \(E[q_i^2 k_i^2] = 1 \cdot 1 = 1\).
• Variance of one term: \(Var(q_i k_i) = 1 - 0 = 1\).

Proof (2)

Why variance adds under independence. For any two random variables \(X, Y\): \[ \text{Var}(X+Y) = \text{Var}(X) + \text{Var}(Y) + 2\,\text{Cov}(X,Y) \] The covariance term vanishes when \(X\) and \(Y\) are independent, since \(\text{Cov}(X,Y) = \mathbb{E}[XY]-\mathbb{E}[X]\mathbb{E}[Y] = 0\) whenever \(X \perp Y\). This extends to a sum of \(d_k\) mutually independent terms: all cross-covariance terms vanish, leaving only the individual variances.

• Because \(q_1,\ldots,q_{d_k}\) and \(k_1,\ldots,k_{d_k}\) are mutually independent (by assumption ), the products \(q_m k_m\) for different \(m\) are independent, and their covariances are all 0.
• So the variance of the sum collapses to the sum of the variances:

\[ \text{Var}(S) = \text{Var}\left(\sum_{m=1}^{d_k} q_m k_m\right) = \sum_{m=1}^{d_k} \text{Var}(q_m k_m) = \sum_{m=1}^{d_k} 1 = d_k \]

Proof (3): Dividing by \(\sqrt{d_k}\)

Recall the scaling rule for variance: for a constant \(c\), \(\text{Var}(cX) = c^2\,\text{Var}(X)\).

• Applying this with \(c = \dfrac{1}{\sqrt{d_k}}\) and \(X = S\):

\[ \text{Var}\left(\frac{S}{\sqrt{d_k}}\right) = \left(\frac{1}{\sqrt{d_k}}\right)^2 \text{Var}(S) = \frac{1}{d_k}\cdot d_k = 1 \]

• So dividing the raw dot product by \(\sqrt{d_k}\) (not \(d_k\) itself) exactly restores unit variance, regardless of \(d_k\).
• This keeps the inputs to softmax in a stable range: without scaling, large-magnitude scores push softmax into extremely peaked (near one-hot) distributions, where gradients with respect to the inputs vanish almost everywhere — the “exploding/vanishing gradients” problem.

Softmax and attention

Recall the definition of softmax:

\[ \begin{bmatrix} 1.3 \\ 5.1\\ 2.2 \\ 0.7 \\ 1.1 \end{bmatrix} \rightarrow \dfrac{e^{z_i}}{\sum_{j=1}^K e^{z_j}} \rightarrow \begin{bmatrix} 0.02 \\ 0.90\\ 0.05 \\ 0.01 \\ 0.02 \end{bmatrix} \]

• Each row in the \(QK^\top\) matrix represents a single word’s “view” of the entire sentence.
• Row 1: How much Word 1 cares about [Word 1, Word 2, …, Word \(n\)].
• Row 2: How much Word 2 cares about [Word 1, Word 2, …, Word \(n\)].
• etc.
• Since we want each word to “distribute” its attention across the sentence, the scores in each row must sum to 1.0.

\[ \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right) \]

Softmax and attention (2)

• In detail, the operation looks like this:

\[ \text{Softmax} \left( \frac{1}{\sqrt{d_k}} \times \underbrace{ \begin{bmatrix} q_1 \cdot k_1 & q_1 \cdot k_2 & \dots & q_1 \cdot k_n \\ q_2 \cdot k_1 & q_2 \cdot k_2 & \dots & q_2 \cdot k_n \\ \vdots & \vdots & \ddots & \vdots \\ q_n \cdot k_1 & q_n \cdot k_2 & \dots & q_n \cdot k_n \end{bmatrix} }_{\text{Raw Scores } (n \times n)} \right) = \underbrace{ \begin{bmatrix} \alpha_{1,1} & \alpha_{1,2} & \dots & \alpha_{1,n} \\ \alpha_{2,1} & \alpha_{2,2} & \dots & \alpha_{2,n} \\ \vdots & \vdots & \ddots & \vdots \\ \alpha_{n,1} & \alpha_{n,2} & \dots & \alpha_{n,n} \end{bmatrix} }_{\text{Attention Weights } (n \times n)} \]

In this resulting matrix:

• \(\sum_{j=1}^{n} \alpha_{i,j} = 1\) (Each row sums to 1).
• \(\alpha_{1,2}\) is the percentage of attention Word 1 pays to Word 2.

Value Vectors

• Value vectors contain detailed information about each token
• They are transformed embeddings that are weighted and summed based on the attention scores
• Intuitively, the value vector contains the potential meanings of a word, and the attention scores determine which meanings are most relevant within a given context
• We get the value vectors analogously from a learned \(W^V\)

\[ V = XW^V = \begin{bmatrix} v_{11} & v_{12} & \dots & v_{1d_v} \\ v_{21} & v_{22} & \dots & v_{2d_v} \\ \vdots & \vdots & \ddots & \vdots \\ v_{n1} & v_{n2} & \dots & v_{nd_v} \end{bmatrix}_{n \times d_v} \]

Calculating Values

Now that we have a matrix of Attention Weights (\(\alpha\)), we use them to take a weighted average of the information stored in the Value matrix (\(V\)).

• The Value Matrix (\(V\)): Just like \(Q\) and \(K\), the Value matrix is created by projecting the input \(X\) through a learnable weight matrix \(W^V\): \[\underbrace{V}_{n \times d_v} = \underbrace{X}_{n \times d} \times \underbrace{W^V}_{d \times d_v}\]

• \(d\) is the dimension of the input embeddings (typically, \(d>d_k\))

• The output dimension of attention is \(d_v\), i.e., attention’s output width is set entirely by the Value projection

• The Value represents the actual “content” of the word.

The Weighted Sum

We take our \(n \times n\) attention weights and multiply them by our \(n \times d_v\) Values.

\[\text{Output} = \underbrace{\text{Softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)}_{n \times n} \times \underbrace{V}_{n \times d_v}\]

To see what happens to Word 1, look at the first row of the multiplication:

\[\text{Output}_1 = \alpha_{1,1}v_1 + \alpha_{1,2}v_2 + \dots + \alpha_{1,n}v_n\]

• The Final Output is a \(n \times d_v\) matrix representing the context-aware version of each token.

Attention: Example

• Many of the attention heads attend to a distant dependency of the verb ‘making’, completing the phrase ‘making…more difficult’. Attentions here shown only for the word ‘making’.

Multi-Head Assembly

• We will cover multihead attention in the section lecture on transformers

Layer Normalization

• For an input to a layer \(\mathbf{x} = \left[x_1, x_2, \ldots, x_d \right]\) where \(d\) is the dimensionality
• The mean is \(\mu_x = \frac{1}{d}\sum_{i=1}^d x_i\)
• The variance is \(\sigma^2 = \frac{1}{d}\sum_{i=1}^d (x_i-\mu_x)^2\)
• \(\hat{x}_i = \frac{x_i-\mu_x}{\sqrt{\sigma^2+\epsilon}}\) (\(\epsilon\) is a small constant for numerical stability)
• \(\mathrm{LayerNorm}(\mathbf{x}) = \gamma \odot \hat{x}_i + \beta\)
• \(\gamma\) and \(\beta\) are learnable parameters that scale and shift the normalized output

Layer Normalization (LN)

• LN greatly contributes to stabilizing training and optimizing performance in deep learning models, especially in transformers. It works by normalizing the activations at each layer, ensuring a more consistent gradient flow during training.

Two flavors:

• Post-LN transformer applies normalization after the addition of the layer’s output with the residual connection’s output.¹
• Pre-LN transformers apply normalization before self-attention and feed-forward layers. This configuration is used in modern architectures such as GPT, Llama, and Vision Transformers.²

LN step 1: Add

The layer performs an Add and Norm operation:

• Add is a Residual Connection (also called a Skip Connection).

• Output: \[\text{Output} = X + Sublayer(X)\]

• strictly element-wise addition, e.g., for token \(i\) and feature \(j\):

\[\text{Output}_{i,j} = x_{i,j} + \text{Sublayer}(X)_{i,j}\]

• Intuition: This operation can be regarded as \(X + f(X)\). During backpropagation, we calculate the first derivative as

\[\frac{\partial(X + f(X))}{\partial X} = 1 + f'(X)\]

• Even if the sublayer’s gradient \(f'(X)\) is near zero (the “vanishing gradient” problem), the total gradient is still at least \(1\).¹

LN step 2: Norm

• Mean (\(\mu\)): Average of the \(d\) values for a specific token.
• Variance (\(\sigma^2\)): Corresponding variance.
• Standardization: mean of 0 and a variance of 1. Before this step, values can be highly divergent which would lead to softmax saturation problems
• \(\gamma\) and \(\beta\) are learnable parameters. They are vectors of size \(d\) (same dimension as the input embedding):

\[\text{Output}_j = (\gamma_j \cdot \hat{x}_j) + \beta_j\]

• Before training, all \(\gamma\) values are initialized to 1 and all \(\beta\) values are initialized to 0.

LN: Effects

• In a transformer model, layer normalization is applied after the multi-head atten- tion mechanism and after the feedforward network.
• This ensures that the outputs of these components have a stable distribution, which is crucial for the model’s ability to learn effectively.

Feed-Forward Network (FFN)

• the FFN transforms that information for each single word, adds necessary nonlinearity, stores factual knowledge, and holds most of the model’s parameters
• The FFN enhances performance of the transformer by introducing non-linear transformations at each token position.
• It consists of two linear transformations and a non-linear activation function (typically ReLU or GELU).

FFN layer

\[ FFN(x) = \text{GELU}(xW_1 + b_1)W_2 + b_2 \]

Rectified Linear Unit (ReLU)

• The Rectified Linear Unit (ReLU) outputs the input if positive; otherwise, it outputs zero.

Gaussian Error Linear Unit (GELU)

Layer Normalization (2)

• The First Add & Norm stabilized the Attention layer.
• This Second Add & Norm stabilizes the FFN layer.
• (This ordering reflects the original Post-LN design;¹ Pre-LN architectures norm before each sublayer instead.)

\[\text{Layer Output} = \text{LayerNorm} \left( \underbrace{X_{\text{post-attn}}}_{\text{Base}} + \underbrace{\text{FFN}(X_{\text{post-attn}})}_{\text{FFN output}} \right) \]

Similar to the previous add & norm layer, this one has two purposes

• The “Add” (Residual): ensures that even if the FFN didn’t learn anything useful for a particular token (e.g. a common word like “the” or “a”), the original context from the attention layer is preserved.
• The “Norm” (LayerNorm): rescales the \(d\) features so their magnitude doesn’t grow progressively larger from layer to layer (e.g. \(d=768\) across \(L=12\) layers for BERT-Base).
• The output of this final Norm has exactly the same shape as the input at the very start of the layer: \(n \times d\).

Conclusion

Sources for this lecture include

• Vaswani A et al. (2017) Attention Is All You Need. arXiv 1706.03762
• Wang et al. (2018), On the Dimensionality of Word Embeddings, (NeurIPS 2018)
• Singh P and Raman B (2025) The Geometry of Intelligence: Foundations of Transformer Networks in Deep Learning. January 2025. ISBN: 978-981-96-4705-7.
• Ba JL, et al. (2016) Layer Normalization. arXiv 1607.06450
• Allamar J (blog) The Illustrated GPT-2 (Visualizing Transformer Language Models)