Transformers: self-attention, multi-head and cross attention

The course has built each architecture from one building block M1 · 1:01: the sigmoid (or other non-linear) neuron for feed-forward networks, convolution and pooling for CNNs, the recurrence \(\mathbf{h}_t = f(\mathbf{h}_{t-1}, \mathbf{x}_t)\) for RNNs, and the attention function for attention-based RNNs. The building blocks of the transformer are self-attention and cross attention. Once they are clear, the full network is easy to assemble.

Take translation, with the input "I enjoyed the movie transformers" and the Tamil target "Naan transfarmar padaththai rasiththen". The whole input sentence is available at once. The RNN encoder still processes it in order: \(\mathbf{h}_1 = f(\mathbf{h}_0, \mathbf{x}_1)\), then \(\mathbf{h}_2 = f(\mathbf{h}_1, \mathbf{x}_2)\), and so on, where \(\mathbf{x}_i\) is the embedding of word \(i\).

The order has a purpose. The embedding \(\mathbf{x}_i\) (from word2vec, GloVe or fastText as in Lecture 10, or simply a learned parameter that starts random) describes the word in general. The state \(\mathbf{h}_i\) describes the word in this sentence, because it carries information from the words before it. A bidirectional RNN also runs from right to left and concatenates the two states, so each word sees context from both sides. These contextual representations are worth keeping.

The price is speed. Each \(\mathbf{h}_t\) waits for \(\mathbf{h}_{t-1}\), so the encoder cannot process the words of a sentence in parallel. The decoder is sequential by nature, since each predicted word is the next step's input M1 · 8:11. The question is whether the encoder can do better.

A plain encoder–decoder (Lecture 13) hands only the final state \(\mathbf{h}_5\) to the decoder. The slides list its names: concept vector, thought vector, context vector or annotation. This model has no notion of alignment between source and target words. Nothing tells the decoder to focus on "I" while it produces "Naan", or on "transformers" while it produces "transfarmar". Attention adds that.

Attention keeps a copy of all the encoder states \(\mathbf{h}_1, \dots, \mathbf{h}_5\) and makes them available to the decoder. Once they are computed, the encoder block can be set aside until decoding finishes. At decoder step \(t\) the decoder receives a context vector, an attention-weighted sum of the encoder states:

\(n\) is the number of input words, \(t\) the decoder time step and \(\mathbf{s}_{t-1}\) the decoder's previous state. The matrix–vector form says that \(\mathbf{c}_t\) is a linear combination of the columns \(\mathbf{h}_i\), and the softmax makes the weights of each step sum to 1.

The slides write the weights as \(\alpha_{ti} = \text{align}(y_t, \mathbf{h}_i)\): for each output word, a distribution over input words. Plotted as a heat map, they show the alignment. In the slides' example the largest weights are \(\alpha_{11}\) (Naan, I), \(\alpha_{25}\) (transfarmar, transformers), \(\alpha_{34}\) (padaththai, film) and \(\alpha_{42}\) (rasiththen, enjoyed).

Three questions decide how parallel this model is.

• Can all \(\mathbf{h}_i\) be computed in parallel? No. The encoder is recurrent.
• At a fixed step \(t\), can all \(\alpha_{ti}\) be computed in parallel? Yes. Each score needs only \(\mathbf{s}_{t-1}\) and \(\mathbf{h}_i\), and both are available at step \(t\). All scores can be computed at once, and then the normalization.
• Can the weights of different steps be computed in parallel? No M1 · 16:27. \(\alpha_{4i}\) needs \(\mathbf{s}_3\), which needs \(\mathbf{c}_3\), which needs every \(\alpha_{3i}\). The rows of the attention matrix come out one at a time.

The slides' takeaway is that attention can be parallelized. The recurrence is what forces the waiting.

The major limitation M1 · 18:59

RNN sequence-to-sequence models with attention perform well on translation, summarization and image or video captioning. Their weakness is training: for a given training example, the computations cannot be parallelized across time steps. The lecture has no complaint about context itself. Recurrence is how an RNN makes its representations contextual, and that part is welcome. The cost is the sequential computation that comes with it.

Wish list: an architecture that keeps the attention mechanism, allows parallel computation, and gets rid of the vanishing and exploding gradients of long recurrences (Lecture 11).

The transformer was introduced in 2017 in the paper Attention Is All You Need. It is still an encoder–decoder, but its blocks are new. The encoder passes word embeddings through a self-attention layer and a feed-forward network. The decoder has a self-attention layer, an encoder–decoder attention layer and a feed-forward network. The lecture studies these components, and a few more, one at a time, and then assembles the full architecture.

Several things carry over from RNN encoder–decoders.

• The input sentence is given in full, and the output is produced one word at a time.
• The inputs are word embeddings. In practice they are randomly initialized and learned with the rest of the network.
• The encoder maps the inputs \(\mathbf{x}_1, \dots, \mathbf{x}_5\) to outputs \(\mathbf{z}_1, \dots, \mathbf{z}_5\), one per word, just as an RNN maps them to \(\mathbf{h}_1, \dots, \mathbf{h}_5\). Each \(\mathbf{z}_i\) is a contextual representation: it knows about word \(i\) and about the other words of the sentence.

What changes is how the context gets in. The transformer does it without recurrent connections, so all \(T\) outputs of a layer can be computed in parallel.

Notation from here on: \(\mathbf{h}_j\) is the input embedding of word \(j\) (in the RNN lectures it was the encoder output), and \(\mathbf{z}_j\) is the output of a block. More intermediate outputs need more letters, so the names shift.

Attention needs a pair of vectors. In RNN attention the pair was a decoder state and an encoder state, \((\mathbf{s}_{t-1}, \mathbf{h}_i)\). Inside the encoder there is no decoder, so the pair is two words of the same sentence. To build a new representation of word \(i\) (say "movie", word 4), take an attention-weighted sum over all the words of the sentence:

This is self-attention: attention among the words of the input itself. It is different from cross attention, between the decoder and the encoder, which comes later. All the \(\mathbf{h}_j\) are known as soon as the sentence is looked up, and no \(\alpha_{ij}\) depends on an earlier output. So the weights for "enjoyed" can be computed at the same time as those for "I" and for "movie".

Why a word needs its context

Compare two sentences.

• "The animal didn't cross the street because it was too tired." Here "it" refers to the animal.
• "The animal didn't cross the street because it was congested." Now "it" refers to the street.

The word "it" and its static embedding are the same in both. Its representation should still lean on "animal" in the first sentence and on "street" in the second, and learned attention weights can do exactly that. The goal is a \(T \times T\) table: for each word, a score against every word of the sentence (itself included) that is higher when the two are related in context. The slides illustrate it with made-up numbers.

Illustrative numbers from the slides, which leave out the middle rows. The row for "it" puts its largest weight, 0.6, on "animal".

For convenience, call the row headers \(\mathbf{s}_i\) and the column headers \(\mathbf{h}_j\), as in RNN attention M2 · 16:04. Here both are the same embeddings. The difference from the seq2seq model is that every \(\mathbf{s}_i\) and every \(\mathbf{h}_j\) is available all the time, where the decoder states \(\mathbf{s}_i\) used to arrive one per step. So all rows of the table can be computed at once.

Which attention function should self-attention use? Recall the score from RNN attention (Lecture 13):

It involves three vectors (\(\mathbf{s}\), \(\mathbf{h}\) and \(\mathbf{v}_a\)), two linear transformations (\(U_{att}\) and \(W_{att}\)), one non-linearity (\(\tanh\)) and, at the end, a dot product with \(\mathbf{v}_a\).

Self-attention starts from the embeddings \(\mathbf{h}_j\) alone, so three vectors have to be made from each embedding. The transformer does it with three matrices, one per role, whose entries are learned:

\(\mathbf{q}_j\) is the query vector of word \(j\), \(\mathbf{k}_j\) its key and \(\mathbf{v}_j\) its value. \(W_Q\), \(W_K\) and \(W_V\) are the corresponding linear transformation (parameter) matrices. For \(T\) words this makes \(3T\) vectors, each computed independently of the others.

The output for one word M3 · 2:03

To compute \(\mathbf{z}_1\), the output for "I", hold its query \(\mathbf{q}_1\) fixed and score it against every key, its own included. The score function is a plain dot product M3 · 6:11. A softmax turns the scores into weights, and the weights combine the values M3 · 7:14:

The roles match the old score function. The query plays the part of \(\mathbf{s}_{t-1}\): it stays fixed while the keys vary. Queries and keys together produce the weights. The output is a weighted sum of the values, much as \(\mathbf{v}_a\) sat outside the rest of the old expression. For \(\mathbf{z}_2\) the same steps run with \(\mathbf{q}_2\) as the query, and so on for every word.

The slides' running example embeds "I enjoyed the movie transformers" in \(\mathbb{R}^3\) as \([1,0,0]\), \([0,1,0]\), \([1,1,0]\), \([1,0,1]\) and \([1,1,1]\), with two-dimensional queries, keys and values, so each \(W\) is \(2 \times 3\). The interactive figure in the next section uses the same embeddings.

Computing \(\mathbf{z}_1\), then \(\mathbf{z}_2\), then \(\mathbf{z}_3\) one after another would bring back the RNN's problem. The computation vectorizes. Stack the embeddings as the columns of \(H = [\mathbf{h}_1, \dots, \mathbf{h}_T]\). Since \(\mathbf{q}_j = W_Q\mathbf{h}_j\) for every \(j\), one matrix product gives every query, and likewise for the keys and values:

With 64-dimensional embeddings and \(64 \times 64\) matrices (the lecture's example), \(Q, K, V \in \mathbb{R}^{64 \times T}\) M3 · 11:18. The matrices could just as well project to a smaller or a larger space. What matters is that \(T\) inputs give \(T\) queries, \(T\) keys and \(T\) values, all at once. Then the whole output is one formula.

To check that this is the same computation as before, follow row 2 M3 · 16:28. Entry \((i, j)\) of \(Q^\top K\) is \(\mathbf{q}_i^\top\mathbf{k}_j\), so row 2 is \([\mathbf{q}_2^\top\mathbf{k}_1, \dots, \mathbf{q}_2^\top\mathbf{k}_T]\), and after scaling and the softmax it is \([\alpha_{21}, \dots, \alpha_{2T}]\). The rows of \(V^\top\) are \(\mathbf{v}_1^\top, \dots, \mathbf{v}_T^\top\). Row 2 of a product is a combination of the rows of the right factor, weighted by row 2 of the left factor:

Every step is a matrix product or an element-wise operation, so no output waits for another. The block is a chain of four operations M3 · 19:37: MatMul (\(Q^\top K\)), Scale (\(1/\sqrt{d_k}\)), Softmax, and MatMul with \(V^\top\). Together with the three projections it is called one head of attention.

The lecture states the scaling without justifying it. Not in the lecture: the original paper's reason is that if the entries of \(\mathbf{q}\) and \(\mathbf{k}\) are independent with mean 0 and variance 1, then \(\mathbf{q}\cdot\mathbf{k}\) has variance \(d_k\). Large scores push the softmax toward a one-hot vector, where its gradients are tiny, and dividing by \(\sqrt{d_k}\) brings the variance back to 1. The scale slider below shows the saturation.

One head takes \(H\), forms \(Q, K, V\) with its own \(W_Q, W_K, W_V\), and returns \(\mathbf{z}_1, \dots, \mathbf{z}_T\). Nothing stops us from running a second head with a second set of matrices, which returns a second set of outputs. Two questions follow: why do it, and what to do with two outputs per word?

Why more than one head

A CNN layer has many filters (Lecture 8). Each runs the same operation with its own weights and, we hope, picks up a different characteristic of the image, such as edges or blur, and more filters give more abstract representations. Heads play the same role. One head learns one way of gathering context, and a word may need more than one.

Take "it" in "The animal didn't cross the street because it was too tired" M4 · 5:41. One useful link is to "animal", the noun it stands for. Another is to "was", the verb it is the subject of. With two heads, one can put a large weight on "animal" and the other a large weight on "was". The slides show this in a trained transformer (a Tensor2Tensor visualization). With a single head, "it" is most strongly linked to "was". With two heads, the first head links it to "was" and the second to "animal".

Nobody tells a head what to look at M4 · 7:13. The heads only have more parameters and so more freedom. If focusing on "animal" in one head lowers the loss, training will make it do so. Each head computes its attention independently, so the heads also run in parallel, as the filters of a CNN layer do.

Combining the heads M4 · 12:49

For each word, concatenate the outputs of the heads and pass the result through a linear transformation \(W_O\). The sizes need care. With 512-dimensional embeddings, eight heads that each returned 512 dimensions would concatenate to 4096, and the next transformation would need many more parameters. So each head projects down: \(W_Q^i, W_K^i, W_V^i \in \mathbb{R}^{64 \times 512}\), every head returns 64-dimensional vectors, and the eight heads concatenate back to \(8 \times 64 = 512\). (The video's two-head example uses 256 dimensions per head for the same reason.)

The input is projected into \(h = 8\) different representation subspaces. In the paper's words, this lets the model "jointly attend to information from different representation subspaces at different positions".

The concatenation happens per word M4 · 14:20: word 1's eight head outputs are joined, then word 2's, and so on. Every layer keeps \(T\) vectors in and \(T\) vectors out.

Back to the basic block. The encoder's self-attention becomes multi-head attention, and its output goes to a feed-forward network. In the slides' terminology these two sublayers form one encoder layer, and the encoder is a stack of \(N = 6\) such layers.

The feed-forward network is small. Each attention output \(\mathbf{z}_i \in \mathbb{R}^{512}\) goes through one hidden layer of \(l = 2048\) ReLU units and back to 512 dimensions:

The slide writes \(\text{FFN}(\mathbf{z}) = \max(0, W_1\mathbf{z} + \mathbf{b}_1)W_2 + \mathbf{b}_2\), half in the row-vector style of the original paper; the computation is the same.

The same network, with the same \(W_1, \mathbf{b}_1, W_2, \mathbf{b}_2\), is applied to every position separately M4 · 17:54. The slides draw it as FFN 1 for \(\mathbf{z}_1\), FFN 2 for \(\mathbf{z}_2\), up to FFN 5, all identical. The words do not interact here: the attention sublayer mixes information across positions, and the feed-forward sublayer transforms each position on its own. The video adds that any non-linearity can be used, typically one of the ReLU family such as GELU (Lecture 7).

The slides count the parameters of one encoder layer with \(h = 8\) heads, \(d_k = d_q = d_v = 64\), \(d_{model} = 512\) and a sentence of \(T = 40\) words. Biases of the attention projections are not counted.

The slide prints 32,178 for \(64 \times 512\); the product is 32,768. No entry depends on \(T\), so the same weights serve sentences of any length.

Each layer takes a \(T \times 512\) input and returns a \(T \times 512\) output, and the output of one layer is the input of the next. The parallelism is horizontal. Within a layer, for a given input, all the queries, weights and outputs of the \(T\) positions are computed at once, where an RNN had to produce \(\mathbf{z}_1\), then \(\mathbf{z}_2\), and so on. It is not vertical: layer 2 still waits for the output of layer 1.

The output of the top encoder layer is written \(\mathbf{e}_1, \dots, \mathbf{e}_T\), or \(E\) as a matrix M4 · 20:29. Each \(\mathbf{e}_j\) is contextual and has passed through six layers of processing. It is what the decoder receives.

The decoder works one output at a time: it produces a word, and that word becomes an input for producing the next. It has two kinds of input. One is the encoder output \(\mathbf{e}_1, \dots, \mathbf{e}_T\), which tells it what is being translated. The other is its own words decoded so far.

At every position the output is a probability distribution over the target vocabulary, of size \(1 \times |V|\), with \(|V| = 37{,}000\) in the slides. The predicted word is the argmax. The first input is the start symbol <Go>; ideally position 1 peaks at "Naan", position 2 at "transfarmar", and so on until a stop symbol is produced.

Like the encoder, the decoder is a stack of \(N = 6\) identical layers (common models use 6, 8 or 12, depending on the data and the task). Each decoder layer has three sublayers:

1. masked multi-head (self) attention over the decoder's own inputs;
2. multi-head (cross) attention between the decoder and the encoder output;
3. a feed-forward network, the same kind as in the encoder.

Teacher forcing Slides

During training, the target sentence is fed to the decoder as one of its inputs. The usual alternative feeds the decoder's own previous prediction back in. If an early prediction is wrong, the later ones are likely to go wrong too, because each prediction is conditioned on those before it, and the errors accumulate. In the slides' English example, one wrong word sends the rest of the sentence away from the ground truth. Training would correct this eventually, but slowly. Teacher forcing feeds the ground-truth previous words instead, whatever the model predicted.

The first decoder sublayer computes queries, keys and values exactly as in the encoder, \(Q = W_Q H\) and so on, where \(\mathbf{h}_1, \dots, \mathbf{h}_T\) are now the embeddings of the target sentence. There is one important difference, a mask.

With teacher forcing the whole target sentence is fed in at once, yet the output at position \(t\) has to be predicted from the words before it. So position \(t\) must not attend to the positions after it. The video explains the same rule from the inference side. The decoder assumes a maximum output length \(T_1\). When only <Go> has been produced, the other inputs are junk, so their attention weights are set to 0 and the first representation depends on <Go> alone. In general, positions that have not been decoded yet never take part in the computation, and the remaining weights are rescaled to sum to 1 M5 · 10:26.

The slides implement this by inserting \(-\infty\) at the masked positions. A \(T \times T\) mask \(M\), 0 on and below the diagonal and \(-\infty\) above it, is added to the scaled scores before the softmax:

Since \(e^{-\infty} = 0\), each masked entry gets a weight of exactly 0, and the softmax normalizes the surviving weights of each row to sum to 1. Row \(t\) uses keys \(1, \dots, t\) only. The rescaling the video mentions happens automatically.

• At inference there is no target sentence and so no teacher forcing. The decoder acts as an auto-regressor, feeding its own predictions back one step at a time.
• In practice the encoder also uses masking in its attention, to hide the padding tokens of sequences shorter than \(T\).

Tick "Causal mask" in the self-attention figure to see the triangular pattern.

The second sublayer brings the two inputs of the decoder together. From the masked self-attention sublayer come the vectors \(\mathbf{s}_1, \dots, \mathbf{s}_T\) (the matrix \(S\)). From the top layer of the encoder come \(\mathbf{e}_1, \dots, \mathbf{e}_T\) (the matrix \(E\)), shared by all layers of the decoder stack. Queries are built from the decoder and keys and values from the encoder, with a new set of matrices \(W_{Q_2}, W_{K_2}, W_{V_2}\):

Because it attends from one half of the model to the other, this is called encoder–decoder attention or cross attention. It is computed with several heads whose outputs are concatenated, as before, and the result goes through the feed-forward network to give the layer's output vectors \(O\).

The slide's formula leaves out the \(1/\sqrt{d_k}\). The heads here are the same scaled dot-product heads as elsewhere, as in the original paper.

For decoder position \(i\), the query \(\mathbf{q}_i\) is scored against the keys of all the source words, so \(\alpha_{ij}\) is the attention that target position \(i\) pays to source word \(j\). The new representation is a weighted sum of the encoder's values, a refinement of \(\mathbf{s}_i\) in the light of the input sentence. This is the RNN attention of Lecture 13 (query from the decoder, keys and values from the encoder), computed for all decoder positions in one matrix product. Source and target may have different lengths, say \(T\) and \(T_1\); then \(Q_2^\top K_2\) is \(T_1 \times T\).

1. One decoder layer, in the slides' notation
2. Masked self-attention: \(\{Q_1, K_1, V_1\}\) \(= f(W_{Q_1}, W_{K_1}, W_{V_1}, \mathbf{h}_{1:T})\), then \(\mathbf{s}_{1:T} = f(Q_1, K_1, V_1; \text{Mask}; W_{O_1})\).
3. Cross attention: \(Q_2 = W_{Q_2}S\) from the decoder, \(K_2 = W_{K_2}E\) and \(V_2 = W_{V_2}E\) from the encoder output \(\mathbf{e}_{1:T}\), then \(f(Q_2, K_2, V_2; W_{O_2})\).
4. The feed-forward network on each position, giving the layer's output.

Parameters of one decoder layer (slides): the FFN has \(2 \times (512 \times 2048) + 2048 + 512 \approx 2\) million, the masked multi-head attention about 1 million and the cross attention about 1 million, so about 4 million per decoder layer.

The top decoder layer outputs one 512-dimensional vector per position. A linear transformation \(W_D\) of size \(512 \times |V|\) maps it to \(|V|\) scores, and a softmax turns them into a probability distribution over the vocabulary. At inference the argmax word is fed back as the next input.

This output layer alone is large: \(512 \times 37{,}000 \approx 19\) million parameters, out of about 65 million for the whole architecture.

To summarize the model so far: an encoder layer has two sublayers (multi-head self-attention and a feed-forward network), and a decoder layer has three (masked multi-head self-attention, multi-head cross attention and a feed-forward network). The key difference in the decoder is the mask: inputs that are not available yet must not take part in the computation, so their weights are set to 0.

The M5 video, the last recording of the course, ends here. The rest of the lecture, from positional encoding to the warm-up schedule, is in the slides only.

An RNN encodes the position of each word in its hidden states, because it reads the words in order. Nothing in a transformer layer does this. The slides call the output of self-attention permutation-invariant: shuffle the input words and each word gets the same output vector as before, only in a different place in the list, since a weighted sum does not care about order. Position has to be put in explicitly.

The fix is to add a positional encoding vector \(\mathbf{p}_j\) to the embedding \(\mathbf{h}'_j\) of the word at position \(j\), before the first layer:

How should \(\mathbf{p}_j\) be filled? The slides consider three simple options and set them aside as unsuitable when the sentence length varies:

• a constant vector whose entries all equal the position \(j\);
• a one-hot vector for the position \(j = 0, 1, \dots, T\);
• a learned embedding for every possible position.

The method chosen rests on a hypothesis. If each position gets a unique pattern of features, the model will learn to attend by relative position.

The transformer generates the features with sines and cosines of the position, at frequencies that fall with the dimension index. For position \(j\) and \(d_{model} = 512\):

The slide writes the cosine's exponent as \((2i+1)/d_{model}\) and lets \(i\) run to 256. The form above is the original paper's, which matches the slide's description: for a fixed position, sine at even and cosine at odd dimensions.

At position 0 every sine is 0 and every cosine is 1, so \(\mathbf{p}_0 = [0, 1, 0, 1, \dots, 0, 1]^\top\). This alternating vector is added to the first word of every sentence. The slides evaluate \(PE_{(j,i)}\) for positions \(j = 0, \dots, 8\) and dimensions \(i = 0, \dots, 63\) and draw it as a heat map. Each even column is a sine and each odd column a cosine, with the frequency decreasing (the wavelength increasing) as \(i\) grows.

Does it respect distance?

Words can be compared by how far apart they are. For "I enjoyed the film transformer" the distance matrix \(|j - k|\) is:

In every row the distance grows to the left and right of the 0, and the matrix is symmetric about the diagonal. A good positional encoding should behave the same way: \(\lVert\mathbf{p}_j - \mathbf{p}_k\rVert\) should grow with \(|j - k|\).

One-hot encoding fails this test. Two different one-hot vectors differ in exactly two entries, so their Euclidean distance is always \(\sqrt{2}\), whatever the positions. The sinusoidal encoding passes it. The slides check this graphically, and the figure below does the same.

Not in the lecture, a short reason. Expanding the squares with \(\sin^2 + \cos^2 = 1\) and \(\cos a\cos b + \sin a \sin b = \cos(a - b)\) gives the expression below. It depends only on \(j - k\), and is symmetric in it.

For a rough comparison with a plain feed-forward network, count attention layers and hidden layers. Every encoder layer has one attention layer and two hidden layers (\(W_1\) and \(W_2\) of the FFN). Every decoder layer has two attention layers and two hidden layers. With 6 encoder and 6 decoder layers the network is \(6 \times 3 + 6 \times 4 = 42\) layers deep, followed by the linear and softmax output.

That raises the two usual questions for a deep network, and the slides give the usual answers.

• How do we ensure the gradient flows across the network? Residual connections, as in ResNet (Lecture 8).
• How do we speed up training? Normalization.

Batch normalization (Lecture 7) works per neuron, across a mini-batch. Let \(x_i^j\) be the activation of neuron \(i\) of layer \(l\) for training example \(j\), and let an accumulator at layer \(l\) store these activations for the \(m\) examples of the batch. The three indices \(l\), \(i\), \(j\) can be pictured as the axes of a cube. For each neuron:

Batch norm can be applied to transformers, but it has limitations.

• The mean and variance are estimates whose accuracy depends on \(m\). A small batch gives a large error.
• A batch size of 1 is ruled out: the statistics degenerate to \(\mu_i = x_i\) and \(\sigma_i = 0\).
• Empirically, naive batch norm degrades performance on NLP tasks. A systematic study confirmed this and proposed a modified batch norm called PowerNorm.

Layer normalization is a simple alternative that works well. It averages over the \(H\) hidden units of the layer, one example at a time:

The statistics no longer depend on the number of examples in the batch, so layer norm works even with a batch size of 1, as an RNN may need.

The complete layer adds a residual connection and a layer norm, an Add & Layer Norm block, after every sublayer: after each multi-head attention, each feed-forward network and each cross-attention block.

Standard form, from the original paper (the slides draw the block without a formula): each sublayer outputs \(\text{LayerNorm}\big(\mathbf{x} + \text{Sublayer}(\mathbf{x})\big)\), where the "Add" is the residual sum.

The full transformer puts the pieces together.

• The output of the top encoder layer is fed to every decoder layer, where cross attention uses it for keys and values.
• The input embeddings are learned while the model trains. No pretrained embedding model such as word2vec is used, so the embeddings are another set of weights.
• The positional encoding is added to the input embeddings. The sinusoidal function is fixed, but it could also be parameterized.

1. A training forward pass, end to end
2. Look up the embeddings of the \(T\) source words and add \(\mathbf{p}_j\) to word \(j\), giving a \(T \times 512\) input.
3. Pass it through 6 encoder layers, each: multi-head self-attention, Add & Layer Norm, feed-forward network, Add & Layer Norm. The top layer gives \(E = \mathbf{e}_{1:T}\).
4. Look up the embeddings of <Go> and the target words (teacher forcing) and add their positional encodings.
5. Pass them through 6 decoder layers, each: masked multi-head self-attention, Add & Layer Norm, multi-head cross attention with keys and values from \(E\), Add & Layer Norm, feed-forward network, Add & Layer Norm.
6. Map each position's output with \(W_D\) and a softmax to a distribution over the vocabulary, and compare it with the next target word.

Module 16.5 builds the transformer's learning-rate schedule in three attempts, for training with mini-batches.

1. A decaying rate, \(\eta = 1/\sqrt{\text{steps}} = \text{steps}^{-0.5}\). Suppose the model takes too long to converge: the error decays very slowly, especially over the first 4000 steps. So try a larger learning rate.
2. A growing rate, increased linearly: \(\eta = \text{steps} \times 4000^{-1.5}\). Suppose the error now decays well for the first 4000 steps and increases after that.
3. Both: increase the rate for the first 4000 steps and decrease it afterwards.

This is counterintuitive. Training usually starts with a high learning rate and keeps decreasing it. Here the rate warms up during the initial steps and decreases only after the warm-up steps. The number of warm-up steps (4000 here) is another hyperparameter.

The two branches combine into one expression. The decreasing curve \(\text{step}^{-0.5}\) and the increasing curve \(\text{step} \cdot \text{warmup}^{-1.5}\) are equal at step = warmup, where both are \(\text{warmup}^{-0.5}\). Before that point the increasing curve is the smaller one; after it, the increasing curve is always above the decreasing one. So the piecewise rule is a minimum, multiplied by the scaling factor \(d_{model}^{-0.5}\):

Warm-up and other learning-rate schemes for general training are also discussed in Lecture 5.

Contextual representation

A vector for a word that depends on the other words of its sentence. RNN states and transformer outputs are both contextual.

Self-attention

Attention among the words of one sequence: \(\mathbf{z}_i = \sum_j \alpha_{ij}\mathbf{v}_j\), with weights from \(\mathbf{q}_i\) and every key \(\mathbf{k}_j\).

Query, key, value

\(\mathbf{q}_j = W_Q\mathbf{h}_j\), \(\mathbf{k}_j = W_K\mathbf{h}_j\), \(\mathbf{v}_j = W_V\mathbf{h}_j\): the vector a word scores others with, the one it is scored by, and the one it contributes.

Scaled dot-product attention

\(\operatorname{softmax}(Q^\top K/\sqrt{d_k})\,V^\top\): dot-product scores, scaled by \(1/\sqrt{d_k}\), a softmax per row, then a weighted sum of values.

Head

One set of \(W_Q, W_K, W_V\) together with its scaled dot-product attention.

Multi-head attention

8 heads of size 64 in parallel, concatenated per word to 512 dimensions and mixed by \(W_O\).

Position-wise FFN

The same 512 → 2048 → 512 ReLU network applied to each position separately.

Encoder layer

Multi-head self-attention and an FFN, each followed by Add & Layer Norm; about 3 million parameters, stacked 6 times.

Teacher forcing

Training the decoder on the ground-truth previous words instead of its own predictions.

Masked self-attention

Self-attention with \(-\infty\) added above the diagonal of the scores, so position \(t\) attends only to positions \(1, \dots, t\).

Cross attention

Encoder–decoder attention: queries from the decoder, keys and values from the top encoder output \(E\).

Positional encoding

A vector \(\mathbf{p}_j\) added to the embedding at position \(j\); sinusoidal in the transformer, with \(\mathbf{p}_0 = [0, 1, 0, 1, \dots]\).

Layer normalization

Normalizing over the hidden units of a layer for each example, so it does not depend on the batch size.

Warm-up

Increasing the learning rate linearly for the first steps (4000) before decaying it as \(\text{step}^{-0.5}\).