Speech and Language Processing. Daniel Jurafsky & James H. Martin. Copyright © 2025.

All
rights reserved. Draft of August 24, 2025.

CHAPTER

8 Transformers

“The true art of memory is the art of attention ”

Samuel Johnson, Idler #74, September 1759

In this chapter we introduce the transformer, the standard architecture for build-
ing large language models. As we discussed in the prior chapter, transformer-based
large language models have completely changed the field of speech and language
processing. Indeed, every subsequent chapter in this textbook will make use of them.
As with the previous chapter, we’ll focus for this chapter on the use of transformers
to model left-to-right (sometimes called causal or autoregressive) language model-
ing, in which we are given a sequence of input tokens and predict output tokens one
by one by conditioning on the prior context.
The transformer is a neural network with a specific structure that includes a
mechanism called self-attention or multi-head attention.1 Attention can be thought
of as a way to build contextual representations of a token’s meaning by attending to
and integrating information from surrounding tokens, helping the model learn how
tokens relate to each other over large spans.

Next token long and thanks for all

Language
Modeling
logits logits logits logits logits …
Head U U U U U

Stacked
… … … … …
Transformer …
Blocks

x1 x2 x3 x4 x5 …
+ 1 + 2 + 3 + 4 + 5
Input
Encoding E E E E E
…

Input tokens So long and thanks for

Figure 8.1 The architecture of a (left-to-right) transformer, showing how each input token
get encoded, passed through a set of stacked transformer blocks, and then a language model
head that predicts the next token.

Fig. 8.1 sketches the transformer architecture. A transformer has three major
components. At the center are columns of transformer blocks. Each block is a
multilayer network (a multi-head attention layer, feedforward networks and layer
1 Although multi-head attention developed historically from the RNN attention mechanism (Chap-
ter 13), we’ll define attention from scratch here.
2 C HAPTER 8 • T RANSFORMERS

normalization steps) that maps an input vector xi in column i (corresponding to input

token i) to an output vector hi . The set of n blocks maps an entire context window
of input vectors (x1 , ..., xn ) to a window of output vectors (h1 , ..., hn ) of the same
length. A column might contain from 12 to 96 or more stacked blocks.
The column of blocks is preceded by the input encoding component, which pro-
cesses an input token (like the word thanks) into a contextual vector representation,
using an embedding matrix E and a mechanism for encoding token position. Each
column is followed by a language modeling head, which takes the embedding out-
put by the final transformer block, passes it through an unembedding matrix U and
a softmax over the vocabulary to generate a single token for that column.
Transformer-based language models are complex, and so the details will un-
fold over the next few chapters. Chapter 7 already discussed how language models
are pretrained, and how tokens are generated via sampling. In the next sections
we’ll introduce multi-head attention, the rest of the transformer block, and the input
encoding and language modeling head components of the transformer. Chapter 10
introduces masked language modeling and the BERT family of bidirectional trans-
former encoder models. Chapter 9 shows how to instruction-tune language models
to perform NLP tasks, and how to align the model with human preferences. Chap-
ter 12 will introduce machine translation with the encoder-decoder architecture.
We’ll see further use of the encoder-decoder architecture in Chapter 15.

8.1 Attention
Recall from Chapter 5 that for word2vec and other static embeddings, the repre-
sentation of a word’s meaning is always the same vector irrespective of the context:
the word chicken, for example, is always represented by the same fixed vector. So
a static vector for the word it might somehow encode that this is a pronoun used
for animals and inanimate entities. But in context it has a much richer meaning.
Consider it in one of these two sentences:
(8.1) The chicken didn’t cross the road because it was too tired.
(8.2) The chicken didn’t cross the road because it was too wide.
In (8.1) it is the chicken (i.e., the reader knows that the chicken was tired), while
in (8.2) it is the road (and the reader knows that the road was wide).2 That is, if
we are to compute the meaning of this sentence, we’ll need the meaning of it to be
associated with the chicken in the first sentence and associated with the road in
the second one, sensitive to the context.
Furthermore, consider reading left to right like a causal language model, pro-
cessing the sentence up to the word it:
(8.3) The chicken didn’t cross the road because it
At this point we don’t yet know which thing it is going to end up referring to! So a
representation of it at this point might have aspects of both chicken and road as
the reader is trying to guess what happens next.
This fact that words have rich linguistic relationships with other words that may
be far away pervades language. Consider two more examples:
(8.4) The keys to the cabinet are on the table.
2 We say that in the first example it corefers with the chicken, and in the second it corefers with the
road; we’ll return to this in Chapter 23.
 8.1 • ATTENTION 3

(8.5) I walked along the pond, and noticed one of the trees along the bank.
In (8.4), the phrase The keys is the subject of the sentence, and in English and many
languages, must agree in grammatical number with the verb are; in this case both are
plural. In English we can’t use a singular verb like is with a plural subject like keys
(we’ll discuss agreement more in Chapter 18). In (8.5), we know that bank refers
to the side of a pond or river and not a financial institution because of the context,
including words like pond. (We’ll discuss word senses more in Chapter 10.)
The point of all these examples is that these contextual words that help us com-
pute the meaning of words in context can be quite far away in the sentence or para-
graph. Transformers can build contextual representations of word meaning, contex-
contextual
embeddings tual embeddings, by integrating the meaning of these helpful contextual words. In a
transformer, layer by layer, we build up richer and richer contextualized representa-
tions of the meanings of input tokens. At each layer, we compute the representation
of a token i by combining information about i from the previous layer with infor-
mation about the neighboring tokens to produce a contextualized representation for
each word at each position.
Attention is the mechanism in the transformer that weighs and combines the
representations from appropriate other tokens in the context from layer k to build
the representation for tokens in layer k + 1.

columns corresponding to input tokens

chicken

because
didn’t
cross

tired
Layer k+1
road
The

the

was
too
it

self-attention distribution
chicken

because
didn’t
cross

tired

Layer k
road
The

the

was
too
it

Figure 8.2 The self-attention weight distribution α that is part of the computation of the
representation for the word it at layer k + 1. In computing the representation for it, we attend
differently to the various words at layer k, with darker shades indicating higher self-attention
values. Note that the transformer is attending highly to the columns corresponding to the
tokens chicken and road , a sensible result, since at the point where it occurs, it could plausibly
corefer with the chicken or the road, and hence we’d like the representation for it to draw on
the representation for these earlier words. Figure adapted from Uszkoreit (2017).

Fig. 8.2 shows a schematic example simplified from a transformer (Uszkoreit,

2017). The figure describes the situation when the current token is it and we need
to compute a contextual representation for this token at layer k +1 of the transformer,
drawing on the representations (from layer k) of every prior token. The figure uses
color to represent the attention distribution over the contextual words: the tokens
chicken and road both have a high attention weight, meaning that as we are com-
puting the representation for it, we will draw most heavily on the representation for
chicken and road. This will be useful in building the final representation for it,
since it will end up coreferring with either chicken or road.
Let’s now turn to how this attention distribution is represented and computed.
4 C HAPTER 8 • T RANSFORMERS

8.1.1 Attention more formally

As we’ve said, the attention computation is a way to compute a vector representation
for a token at a particular layer of a transformer, by selectively attending to and
integrating information from prior tokens at the previous layer. Attention takes an
input representation xi corresponding to the input token at position i, and a context
window of prior inputs x1 ..xi−1 , and produces an output ai .
In causal, left-to-right language models, the context is any of the prior words.
That is, when processing xi , the model has access to xi as well as the representations
of all the prior tokens in the context window (context windows consist of thousands
of tokens) but no tokens after i. (By contrast, in Chapter 10 we’ll generalize attention
so it can also look ahead to future words.)
Fig. 8.3 illustrates this flow of information in an entire causal self-attention layer,
in which this same attention computation happens in parallel at each token position
i. Thus a self-attention layer maps input sequences (x1 , ..., xn ) to output sequences
of the same length (a1 , ..., an ).

a1 a2 a3 a4 a5

Self-Attention attention attention attention attention attention

Layer

x1 x2 x3 x4 x5

Figure 8.3 Information flow in causal self-attention. When processing each input xi , the
model attends to all the inputs up to, and including xi .

Simplified version of attention At its heart, attention is really just a weighted

sum of context vectors, with a lot of complications added to how the weights are
computed and what gets summed. For pedagogical purposes let’s first describe a
simplified intuition of attention, in which the attention output ai at token position i
is simply the weighted sum of all the representations x j , for all j ≤ i; we’ll use αi j
to mean how much x j should contribute to ai :
X
Simplified version: ai = αi j x j (8.6)
j≤i

Each αi j is a scalar used for weighing the value of input x j when summing up
the inputs to compute ai . How shall we compute this α weighting? In attention we
weight each prior embedding proportionally to how similar it is to the current token
i. So the output of attention is a sum of the embeddings of prior tokens weighted
by their similarity with the current token embedding. We compute similarity scores
via dot product, which maps two vectors into a scalar value ranging from −∞ to
∞. The larger the score, the more similar the vectors that are being compared. We’ll
normalize these scores with a softmax to create the vector of weights αi j , j ≤ i.

Simplified Version: score(xi , x j ) = xi · x j (8.7)

αi j = softmax(score(xi , x j )) ∀ j ≤ i (8.8)

Thus in Fig. 8.3 we compute a3 by computing three scores: x3 · x1 , x3 · x2 and x3 · x3 ,

normalizing them by a softmax, and using the resulting probabilities as weights
indicating each of their proportional relevance to the current position i. Of course,
 8.1 • ATTENTION 5

the softmax weight will likely be highest for xi , since xi is very similar to itself,
resulting in a high dot product. But other context words may also be similar to i, and
the softmax will also assign some weight to those words. Then we use these weights
as the α values in Eq. 8.6 to compute the weighted sum that is our a3 .
The simplified attention in equations 8.6 – 8.8 demonstrates the attention-based
approach to computing ai : compare the xi to prior vectors, normalize those scores
into a probability distribution used to weight the sum of the prior vector. But now
we’re ready to remove the simplifications.
A single attention head using query, key, and value matrices Now that we’ve
attention head seen a simple intuition of attention, let’s introduce the actual attention head, the
head version of attention that’s used in transformers. (The word head is often used in
transformers to refer to specific structured layers). The attention head allows us to
distinctly represent three different roles that each input embedding plays during the
course of the attention process:
• As the current element being compared to the preceding inputs. We’ll refer to
query this role as a query.
• In its role as a preceding input that is being compared to the current element
key to determine a similarity weight. We’ll refer to this role as a key.
value • And finally, as a value of a preceding element that gets weighted and summed
up to compute the output for the current element.
To capture these three different roles, transformers introduce weight matrices
WQ , WK , and WV . These weights will project each input vector xi into a represen-
tation of its role as a query, key, or value:

qi = xi WQ ; ki = xi WK ; vi = xi WV (8.9)

Given these projections, when we are computing the similarity of the current ele-
ment xi with some prior element x j , we’ll use the dot product between the current
element’s query vector qi and the preceding element’s key vector k j . Furthermore,
the result of a dot product can be an arbitrarily large (positive or negative) value, and
exponentiating large values can lead to numerical issues and loss of gradients during
training. To avoid this, we scale the dot product by a factor related to the size of the
embeddings, via dividing by the square root of the dimensionality of the query and
key vectors (dk ). We thus replace the simplified Eq. 8.7 with Eq. 8.11. The ensuing
softmax calculation resulting in αi j remains the same, but the output calculation for
headi is now based on a weighted sum over the value vectors v (Eq. 8.13).
Here’s a final set of equations for computing self-attention for a single self-
attention output vector ai from a single input vector xi . This version of attention
computes ai by summing the values of the prior elements, each weighted by the
similarity of its key to the query from the current element:

qi = xi WQ ; k j = x j WK ; v j = x j WV (8.10)
qi · k j
score(xi , x j ) = √ (8.11)
dk
αi j = softmax(score(xi , x j )) ∀ j ≤ i (8.12)
X
headi = αi j v j (8.13)
j≤i

ai = headi WO (8.14)
6 C HAPTER 8 • T RANSFORMERS

8. Output of self-attention a3 [1 × d]

7. Reshape to [1 x d] WO [dv × d]

[1 × dv]
6. Sum the weighted
value vectors

[1 × dv] [1 × dv] [1 × dv]

𝛼3,1 𝛼3,2 𝛼3,3

5. Weigh each value vector

×
×
4. Turn into 𝛼i,j weights via softmax

3. Divide scalar score by √dk √d ÷ √dk

÷
√dk
÷
k

2. Compare x3’s query with

the keys for x1, x2, and x3
[1 × dv] [1 × dv] [1 x dv]

1. Generate k q v k q v k q v
key, query, value WK WQ WV WK WQ WV WK WQ WV
vectors

x1 x2 x3
[1 × d] [1 × d] [1 × d]

Figure 8.4 Calculating the value of a3 , the third element of a sequence using causal (left-
to-right) self-attention.

We illustrate this in Fig. 8.4 for the case of calculating the value of the third output
a3 in a sequence.
Note that we’ve also introduced one more matrix, WO , which is right-multiplied
by the attention head. This is necessary to reshape the output of the head. The input
to attention xi and the output from attention ai both have the same dimensionality
[1 × d]. We often call d the model dimensionality, and indeed as we’ll discuss in
Section 8.2 the output hi of each transformer block, as well as the intermediate vec-
tors inside the transformer block also have the same dimensionality [1 × d]. Having
everything be the same dimensionality makes the transformer very modular.
So let’s talk shapes. How do we get from [1 × d] at the input to [1 × d] at the
output? Let’s look at all the internal shapes. We’ll have a dimension dk for the
query and key vectors. The query vector and the key vector are both dimensionality
[1 × dk ], so we can take their dot product qi · k j to produce a scalar. We’ll have a
separate dimension dv for the value vectors. The transform matrix WQ has shape
[d × dk ], WK is [d × dk ], and WV is [d × dv ]. So the output of headi in equation
Eq. 8.13 is of shape [1 × dv ]. To get the desired output shape [1 × d] we’ll need to
reshape the head output, and so WO is of shape [dv × d]. In the original transformer
work (Vaswani et al., 2017), d was 512, dk and dv were both 64.
Multi-head Attention Equations 8.11-8.13 describe a single attention head. But
actually, transformers use multiple attention heads. The intuition is that each head
might be attending to the context for different purposes: heads might be special-
ized to represent different linguistic relationships between context elements and the
current token, or to look for particular kinds of patterns in the context.
multi-head So in multi-head attention we have A separate attention heads that reside in
attention
parallel layers at the same depth in a model, each with its own set of parameters that
allows the head to model different aspects of the relationships among inputs. Thus
 8.2 • T RANSFORMER B LOCKS 7

each head i in a self-attention layer has its own set of query, key, and value matrices:
WQi , WKi , and WVi . These are used to project the inputs into separate query, key,
and value embeddings for each head.
When using multiple heads the model dimension d is still used for the input
and output, the query and key embeddings have dimensionality dk , and the value
embeddings are of dimensionality dv (again, in the original transformer paper dk =
dv = 64, A = 8, and d = 512). Thus for each head i, we have weight layers WQi of
shape [d × dk ], WKi of shape [d × dk ], and WVi of shape [d × dv ].
Below are the equations for attention augmented with multiple heads; Fig. 8.5
shows an intuition.
qci = xi WQc ; kcj = x j WKc ; vcj = x j WVc ; ∀ c 1 ≤ c ≤ A (8.15)
qci · kcj
scorec (xi , x j ) = √ (8.16)
dk
αicj = softmax(scorec (xi , x j )) ∀ j ≤ i (8.17)

headci = αicj vcj

X
(8.18)
j≤i

ai = (head1 ⊕ head2 ... ⊕ headA )WO (8.19)

MultiHeadAttention(xi , [x1 , · · · , xi−1 ]) = ai (8.20)

Note in Eq. 8.20 that MultiHeadAttention is a function of the current input xi , as

well as all the other inputs. For the causal or left-to-right attention that we use in
this chapter, the other inputs are only to the left, but we’ll also see a version of
attention in Chapter 10 where attention is a function of the tokens to the right as
well. We’ll return to this idea about causal inputs in Eq. 8.34 when we introduce the
idea of masking the right context.
The output of each of the A heads is of shape [1 × dv ], and so the output of the
multi-head layer with A heads consists of A vectors of shape [1 × dv ]. These are
concatenated to produce a single output with dimensionality [1 × Adv ]. Then we use
yet another linear projection WO ∈ RAdv ×d to reshape it, resulting in the multi-head
attention vector ai with the correct output shape [1 × d] at each input i.

8.2 Transformer Blocks

The self-attention calculation lies at the core of what’s called a transformer block,
which, in addition to the self-attention layer, includes three other kinds of layers: (1)
a feedforward layer, (2) residual connections, and (3) normalizing layers (colloqui-
ally called “layer norm”).
Fig. 8.6 illustrates a transformer block, sketching a common way of thinking
residual stream about the block that is called the residual stream (Elhage et al., 2021). In the resid-
ual stream viewpoint, we consider the processing of an individual token i through
the transformer block as a single stream of d-dimensional representations for token
position i. This residual stream starts with the original input vector, and the various
components read their input from the residual stream and add their output back into
the stream.
The input at the bottom of the stream is an embedding for a token, which has
dimensionality d. This initial embedding gets passed up (by residual connections),
and is progressively added to by the other components of the transformer: the at-
8 C HAPTER 8 • T RANSFORMERS

ai
[1 x d]
Project down to d WO [Adv x d]
… [1 x Adv ]
Concatenate Outputs

[1 x dv ] [1 x dv ]
Each head Head 1 Head 2 Head 8
attends differently …
WK1 WV1 WQ1 WK2 WV2 WQ2 WK8 WV8 WQ8
to context

… xi-3 xi-2 xi-1 xi [1 ax d]

i
Figure 8.5 The multi-head attention computation for input xi , producing output ai . A multi-head attention
layer has A heads, each with its own query, key, and value weight matrices. The outputs from each of the heads
are concatenated and then projected down to d, thus producing an output of the same size as the input.

hi-1 hi hi+1
Residual
Stream

Feedforward

Layer Norm
… …
+
MultiHead
Attention

Layer Norm

xi-1 xi xi+1

Figure 8.6 The architecture of a transformer block showing the residual stream. This
figure shows the prenorm version of the architecture, in which the layer norms happen before
the attention and feedforward layers rather than after.

tention layer that we have seen, and the feedforward layer that we will introduce.
Before the attention and feedforward layer is a computation called the layer norm.
Thus the initial vector is passed through a layer norm and attention layer, and
the result is added back into the stream, in this case to the original input vector
xi . And then this summed vector is again passed through another layer norm and a
feedforward layer, and the output of those is added back into the residual, and we’ll
use hi to refer to the resulting output of the transformer block for token i. (In earlier
descriptions the residual stream was often described using a different metaphor as
residual connections that add the input of a component to its output, but the residual
stream is a more perspicuous way of visualizing the transformer.)
 8.2 • T RANSFORMER B LOCKS 9

We’ve already seen the attention layer, so let’s now introduce the feedforward
and layer norm computations in the context of processing a single input xi at token
position i.

Feedforward layer The feedforward layer is a fully-connected 2-layer network,

i.e., one hidden layer, two weight matrices, as introduced in Chapter 6. The weights
are the same for each token position i, but are different from layer to layer. It is com-
mon to make the dimensionality dff of the hidden layer of the feedforward network
be larger than the model dimensionality d. (For example in the original transformer
model, d = 512 and dff = 2048.)

FFN(xi ) = ReLU(xi W1 + b1 )W2 + b2 (8.21)

Layer Norm At two stages in the transformer block we normalize the vector (Ba
layer norm et al., 2016). This process, called layer norm (short for layer normalization), is one
of many forms of normalization that can be used to improve training performance
in deep neural networks by keeping the values of a hidden layer in a range that
facilitates gradient-based training.
Layer norm is a variation of the z-score from statistics, applied to a single vec-
tor in a hidden layer. That is, the term layer norm is a bit confusing; layer norm
is not applied to an entire transformer layer, but just to the embedding vector of a
single token. Thus the input to layer norm is a single vector of dimensionality d
and the output is that vector normalized, again of dimensionality d. The first step in
layer normalization is to calculate the mean, µ, and standard deviation, σ , over the
elements of the vector to be normalized. Given an embedding vector x of dimen-
sionality d, these values are calculated as follows.

d
1X
µ = xi (8.22)
d
i=1
v
u d
u1 X
σ = t (xi − µ)2 (8.23)
d
i=1

Given these values, the vector components are normalized by subtracting the mean
from each and dividing by the standard deviation. The result of this computation is
a new vector with zero mean and a standard deviation of one.

(x − µ)
x̂ = (8.24)
σ

Finally, in the standard implementation of layer normalization, two learnable param-

eters, γ and β , representing gain and offset values, are introduced.

(x − µ)
LayerNorm(x) = γ +β (8.25)
σ

Putting it all together The function computed by a transformer block can be ex-
pressed by breaking it down with one equation for each component computation,
using t (of shape [1 × d]) to stand for transformer and superscripts to demarcate
10 C HAPTER 8 • T RANSFORMERS

each computation inside the block:

t1i = LayerNorm(xi ) (8.26)

2
= MultiHeadAttention(ti , t11 , · · · , t1N )
1
 
ti (8.27)
3 2
ti = ti + xi (8.28)
ti = LayerNorm(ti )
4 3
(8.29)
t5i = FFN(t4i ) (8.30)
5 3
hi = ti + ti (8.31)

Notice that the only component that takes as input information from other tokens
(other residual streams) is multi-head attention, which (as we see from Eq. 8.27)
looks at all the neighboring tokens in the context. The output from attention, how-
ever, is then added into this token’s embedding stream. In fact, Elhage et al. (2021)
show that we can view attention heads as literally moving information from the
residual stream of a neighboring token into the current stream. The high-dimensional
embedding space at each position thus contains information about the current to-
ken and about neighboring tokens, albeit in different subspaces of the vector space.
Fig. 8.7 shows a visualization of this movement.

Token A Token B
residual residual
stream stream

Figure 8.7 An attention head can move information from token A’s residual stream into
token B’s residual stream.

Crucially, the input and output dimensions of transformer blocks are matched so
they can be stacked. Each token vector xi at the input to the block has dimensionality
d, and the output hi also has dimensionality d. Transformers for large language
models stack many of these blocks, from 12 layers (used for the T5 or GPT-3-small
language models) to 96 layers (used for GPT-3 large), to even more for more recent
models. We’ll come back to this issue of stacking in a bit.
Equation 8.26 and following are just the equation for a single transformer block,
but the residual stream metaphor goes through all the transformer layers, from the
first transformer blocks to the 12th, in a 12-layer transformer. At the earlier trans-
former blocks, the residual stream is representing the current token. At the highest
transformer blocks, the residual stream is usually representing the following token,
since at the very end it’s being trained to predict the next token.
Once we stack many blocks, there is one more requirement: at the very end of
the last (highest) transformer block, there is a single extra layer norm that is run on
the last hi of each token stream (just below the language model head layer that we
will define soon). 3
3 Note that we are using the most common current transformer architecture, which is called the prenorm
 8.3 • PARALLELIZING COMPUTATION USING A SINGLE MATRIX X 11

8.3 Parallelizing computation using a single matrix X

This description of multi-head attention and the rest of the transformer block has
been from the perspective of computing a single output at a single time step i in
a single residual stream. But as we pointed out earlier, the attention computation
performed for each token to compute ai is independent of the computation for each
other token, and that’s also true for all the computation in the transformer block
computing hi from the input xi . That means we can easily parallelize the entire
computation, taking advantage of efficient matrix multiplication routines.
We do this by packing the input embeddings for the N tokens of the input se-
quence into a single matrix X of size [N × d]. Each row of X is the embedding of
one token of the input. Transformers for large language models commonly have an
input length N from 1K to 32K; much longer contexts of 128K or even up to millions
of tokens can also be achieved with architectural changes like special long-context
mechanisms that we don’t discuss here. So for vanilla transformers, we can think of
X having between 1K and 32K rows, each of the dimensionality of the embedding
d (the model dimension).
Parallelizing attention Let’s first see this for a single attention head and then turn
to multiple heads, and then add in the rest of the components in the transformer
block. For one head we multiply X by the query, key, and value matrices WQ of
shape [d × dk ], WK of shape [d × dk ], and WV of shape [d × dv ], to produce matrices
Q of shape [N × dk ], K of shape [N × dk ], and V of shape [N × dv ], containing all the
key, query, and value vectors:

Q = XWQ ; K = XWK ; V = XWV (8.32)

Given these matrices we can compute all the requisite query-key comparisons simul-
taneously by multiplying Q and K| in a single matrix multiplication. The product is
of shape N × N, visualized in Fig. 8.8.

q1•k1 q1•k2 q1•k3 q1•k4

q2•k1 q2•k2 q2•k3 q2•k4

N
q3•k1 q3•k2 q3•k3 q3•k4

q4•k1 q4•k2 q4•k3 q4•k4

Figure 8.8 The N × N QK| matrix showing how it computes all qi · k j comparisons in a
single matrix multiple.

Once we have this QK| matrix, we can very efficiently scale these scores, take
the softmax, and then multiply the result by V resulting in a matrix of shape N × d:
a vector embedding representation for each token in the input. We’ve reduced the
entire self-attention step for an entire sequence of N tokens for one head to the
architecture. The original definition of the transformer in Vaswani et al. (2017) used an alternative archi-
tecture called the postnorm transformer in which the layer norm happens after the attention and FFN
layers; it turns out moving the layer norm beforehand works better, but does require this one extra layer
at the end.
12 C HAPTER 8 • T RANSFORMERS

following computation:

QK|
  
head = softmax mask √ V (8.33)
dk
A = head WO (8.34)

Masking out the future You may have noticed that we introduced a mask function
in Eq. 8.34 above. This is because the self-attention computation as we’ve described
it has a problem: the calculation of QK| results in a score for each query value to
every key value, including those that follow the query. This is inappropriate in the
setting of language modeling: guessing the next word is pretty simple if you already
know it! To fix this, the elements in the upper-triangular portion of the matrix are set
to −∞, which the softmax will turn to zero, thus eliminating any knowledge of words
that follow in the sequence. This is done in practice by adding a mask matrix M in
which Mi j = −∞ ∀ j > i (i.e. for the upper-triangular portion) and Mi j = 0 otherwise.
Fig. 8.9 shows the resulting masked QK| matrix. (we’ll see in Chapter 10 how to
make use of words in the future for tasks that need it).

q1•k1 −∞ −∞ −∞

q2•k1 q2•k2 −∞ −∞
N
q3•k1 q3•k2 q3•k3 −∞

q4•k1 q4•k2 q4•k3 q4•k4

Figure 8.9 The N × N QK| matrix showing the qi · k j values, with the upper-triangle por-
tion of the comparisons matrix zeroed out (set to −∞, which the softmax will turn to zero).

Fig. 8.10 shows a schematic of all the computations for a single attention head
parallelized in matrix form.
Fig. 8.8 and Fig. 8.9 also make it clear that attention is quadratic in the length
of the input, since at each layer we need to compute dot products between each pair
of tokens in the input. This makes it expensive to compute attention over very long
documents (like entire novels). Nonetheless modern large language models manage
to use quite long contexts of thousands or tens of thousands of tokens.

Parallelizing multi-head attention In multi-head attention, as with self-attention,

the input and output have the model dimension d, the key and query embeddings
have dimensionality dk , and the value embeddings are of dimensionality dv (again,
in the original transformer paper dk = dv = 64, A = 8, and d = 512). Thus for
each head c, we have weight layers WQ c of shape [d × dk ], WK c of shape [d × dk ],
and WV c of shape [d × dv ], and these get multiplied by the inputs packed into X to
produce Q of shape [N × dk ], K of shape [N × dk ], and V of shape [N × dv ]. The
output of each of the A heads is of shape [N × dv ], and so the output of the multi-
head layer with A heads consists of A matrices of shape [N × dv ]. To make use
of these matrices in further processing, they are concatenated to produce a single
output with dimensionality [N × Adv ]. Finally, we use a final linear projection WO
of shape [Adv × d], that reshapes it to the original output dimension for each token.
Multiplying the concatenated [N × Adv ] matrix output by WO of shape [Adv × d]
 8.3 • PARALLELIZING COMPUTATION USING A SINGLE MATRIX X 13

X Q X K X V
Input
WQ Query Input WK Key Input WV Value
Token 1 Token 1 Token 1 Token 1 Token 1
Token 1
Input Input Key Input Value
Query
Token 2 Token 2 Token 2 Token 2
Input x =
Token 2
x = Key
x =
Token 2
Query Input Input Value
Token 3 Token 3 Token 3 Token 3 Token 3
Token 3
Input Input Key Input Value
Query
Token 4 Token 4 Token 4 Token 4
Token 4 d x dk d x dv Token 4
d x dk
Nxd N x dk Nxd N x dk N x dv
Nxd

Q KT QKT QKT masked V A

q1
x = −∞ −∞ −∞ v1 a1
k1

k2

k3

k4

q1•k1 q1•k2 q1•k3 q1•k4 q1•k1

q1•k1

mask q2 q2•k1 q2•k2 q2•k3 q2•k4 = q2•k1 q2•k2 −∞ −∞ x v2 = a2

q3 q3•k1 q3•k2 q3•k3 q3•k4 q3•k1 q3•k2 q3•k3 −∞ v3 a3

q4 dk x N q4•k1 q4•k2 q4•k3 q4•k4 q4•k1 q4•k2 q4•k3 q4•k4 v4 a4

N x dk NxN NxN N x dv N x dv

Figure 8.10 Schematic of the attention computation for a single attention head in parallel. The first row shows
the computation of the Q, K, and V matrices. The second row shows the computation of QKT , the masking
(the softmax computation and the normalizing by dimensionality are not shown) and then the weighted sum of
the value vectors to get the final attention vectors.

yields the self-attention output A of shape [N × d].

Qi = XWQi ; Ki = XWKi ; Vi = XWVi (8.35)

  i i | 
QK
headi = SelfAttention(Q , K , V ) = softmax mask √
i i i
Vi (8.36)
dk
MultiHeadAttention(X) = (head1 ⊕ head2 ... ⊕ headA )WO (8.37)

Putting it all together with the parallel input matrix X The function computed
in parallel by an entire layer of N transformer blocks—each block over one of the N
input tokens—can be expressed as:

O = X + MultiHeadAttention(LayerNorm(X)) (8.38)
H = O + FFN(LayerNorm(O)) (8.39)

Note that in Eq. 8.38 we are using X to mean the input to the layer, wherever it
comes from. For the first layer, as we will see in the next section, that input is the
initial word + positional embedding vectors that we have been describing by X. But
for subsequent layers k, the input is the output from the previous layer Hk−1 . We
can also break down the computation performed in a transformer layer, showing one
equation for each component computation. We’ll use T (of shape [N × d]) to stand
for transformer and superscripts to demarcate each computation inside the block,
and again use X to mean the input to the block from the previous layer or the initial
14 C HAPTER 8 • T RANSFORMERS

embedding:

T1 = LayerNorm(X) (8.40)
T 2
= MultiHeadAttention(T )
1
(8.41)
T3 = T2 + X (8.42)
T4 = LayerNorm(T3 ) (8.43)
T 5
= FFN(T )4
(8.44)
5 3
H = T +T (8.45)

Here when we use a notation like FFN(T3 ) we mean that the same FFN is applied
in parallel to each of the N embedding vectors in the window. Similarly, each of the
N tokens is normed in parallel in the LayerNorm. Crucially, the input and output
dimensions of transformer blocks are matched so they can be stacked. Since each
token xi at the input to the block is represented by an embedding of dimensionality
[1 × d], that means the input X and output H are both of shape [N × d].

8.4 The input: embeddings for token and position

Let’s talk about where the input X comes from. Given a sequence of N tokens (N is
embedding the context length in tokens), the matrix X of shape [N × d] has an embedding for
each word in the context. The transformer does this by separately computing two
embeddings: an input token embedding, and an input positional embedding.
A token embedding, introduced in Chapter 6, is a vector of dimension d that will
be our initial representation for the input token. (As we pass vectors up through the
transformer layers in the residual stream, this embedding representation will change
and grow, incorporating context and playing a different role depending on the kind
of language model we are building.) The set of initial embeddings are stored in the
embedding matrix E, which has a row for each of the |V | tokens in the vocabulary.
(Reminder that V here means the vocabulary of tokens, this V is not related to the
value vector.) Thus each word is a row vector of d dimensions, and E has shape
[|V | × d].
Given an input token string like Thanks for all the we first convert the tokens
into vocabulary indices (these were created when we first tokenized the input using
BPE or SentencePiece). So the representation of thanks for all the might be w =
[5, 4000, 10532, 2224]. Next we use indexing to select the corresponding rows from
E, (row 5, row 4000, row 10532, row 2224).
Another way to think about selecting token embeddings from the embedding
matrix is to represent tokens as one-hot vectors of shape [1 × |V |], i.e., with one
one-hot vector dimension for each word in the vocabulary. Recall that in a one-hot vector all the
elements are 0 except one, the element whose dimension is the word’s index in the
vocabulary, which has value 1. So if the word “thanks” has index 5 in the vocabulary,
x5 = 1, and xi = 0 ∀i 6= 5, as shown here:
[0 0 0 0 1 0 0 ... 0 0 0 0]
1 2 3 4 5 6 7 ... ... |V|
Multiplying by a one-hot vector that has only one non-zero element xi = 1 simply
selects out the relevant row vector for word i, resulting in the embedding for word i,
as depicted in Fig. 8.11.
 8.4 • T HE INPUT: EMBEDDINGS FOR TOKEN AND POSITION 15

5 |V| 5 d
1 0000100…0000 ✕ E = 1

|V|

Figure 8.11 Selecting the embedding vector for word V5 by multiplying the embedding
matrix E with a one-hot vector with a 1 in index 5.

We can extend this idea to represent the entire token sequence as a matrix of one-
hot vectors, one for each of the N positions in the transformer’s context window, as
shown in Fig. 8.12.

d
|V| d
0000100…0000
0000000…0010
1000000…0000 ✕ E =
…
N 0000100…0000
N
| V|
Figure 8.12 Selecting the embedding matrix for the input sequence of token ids W by mul-
tiplying a one-hot matrix corresponding to W by the embedding matrix E.

These token embeddings are not position-dependent. To represent the position

of each token in the sequence, we combine these token embeddings with positional
positional
embeddings embeddings specific to each position in an input sequence.
Where do we get these positional embeddings? The simplest method, called
absolute
position absolute position, is to start with randomly initialized embeddings corresponding
to each possible input position up to some maximum length. For example, just as
we have an embedding for the word fish, we’ll have an embedding for the position 3.
As with word embeddings, these positional embeddings are learned along with other
parameters during training. We can store them in a matrix Epos of shape [N × d].
To produce an input embedding that captures positional information, we just
add the word embedding for each input to its corresponding positional embedding.
The individual token and position embeddings are both of size [1×d], so their sum is
also [1×d], This new embedding serves as the input for further processing. Fig. 8.13
shows the idea.

Transformer Block

X = Composite
Embeddings
(word + position)
+
+

Word
Janet

back
will

the

bill

Embeddings
Position
1

Embeddings
Janet will back the bill

Figure 8.13 A simple way to model position: add an embedding of the absolute position to
the token embedding to produce a new embedding of the same dimensionality.
16 C HAPTER 8 • T RANSFORMERS

The final representation of the input, the matrix X, is an [N × d] matrix in which

each row i is the representation of the ith token in the input, computed by adding
E[id(i)]—the embedding of the id of the token that occurred at position i—, to P[i],
the positional embedding of position i.
A potential problem with the simple position embedding approach is that there
will be plenty of training examples for the initial positions in our inputs and corre-
spondingly fewer at the outer length limits. These latter embeddings may be poorly
trained and may not generalize well during testing. An alternative is to choose a
static function that maps integer inputs to real-valued vectors in a way that better
handles sequences of arbitrary length. A combination of sine and cosine functions
with differing frequencies was used in the original transformer work. Sinusoidal po-
sition embeddings may also help in capturing the inherent relationships among the
positions, like the fact that position 4 in an input is more closely related to position
5 than it is to position 17.
A more complex style of positional embedding methods extend this idea of cap-
relative
position turing relationships even further to directly represent relative position instead of
absolute position, often implemented in the attention mechanism at each layer rather
than being added once at the initial input.

8.5 The Language Modeling Head

The last component of the transformer we must introduce is the language modeling
language
modeling head head. Here we are using the word head to mean the additional neural circuitry we
head add on top of the basic transformer architecture when we apply pretrained trans-
former models to various tasks. The language modeling head is the circuitry we
need to do language modeling.
Recall that language models, from the simple n-gram models of Chapter 3 through
the feedforward and RNN language models of Chapter 6 and Chapter 13, are word
predictors. Given a context of words, they assign a probability to each possible next
word. For example, if the preceding context is “Thanks for all the” and we want to
know how likely the next word is “fish” we would compute:

P(fish|Thanks for all the)

Language models give us the ability to assign such a conditional probability to every
possible next word, giving us a distribution over the entire vocabulary. The n-gram
language models of Chapter 3 compute the probability of a word given counts of
its occurrence with the n − 1 prior words. The context is thus of size n − 1. For
transformer language models, the context is the size of the transformer’s context
window, which can be quite large, like 32K tokens for large models (and much larger
contexts of millions of words are possible with special long-context architectures).
The job of the language modeling head is to take the output of the final trans-
former layer from the last token N and use it to predict the upcoming word at posi-
tion N + 1. Fig. 8.14 shows how to accomplish this task, taking the output of the last
token at the last layer (the d-dimensional output embedding of shape [1 × d]) and
producing a probability distribution over words (from which we will choose one to
generate).
The first module in Fig. 8.14 is a linear layer, whose job is to project from the
output hLN , which represents the output token embedding at position N from the final
 8.5 • T HE L ANGUAGE M ODELING H EAD 17

y1 y2 … y|V| Word probabilities 1 x |V|

Language Model Head Softmax over vocabulary V

takes hLN and outputs a u1 u2 … u|V| Logits 1 x |V|
distribution over vocabulary V
Unembedding layer Unembedding layer d x |V|
U = ET

hL1 hL2 hLN 1xd

Layer L
Transformer
Block
…
w1 w2 wN

Figure 8.14 The language modeling head: the circuit at the top of a transformer that maps from the output
embedding for token N from the last transformer layer (hLN ) to a probability distribution over words in the
vocabulary V .

logit block L, (hence of shape [1 × d]) to the logit vector, or score vector, that will have a
single score for each of the |V | possible words in the vocabulary V . The logit vector
u is thus of dimensionality [1 × |V |].
This linear layer can be learned, but more commonly we tie this matrix to (the
weight tying transpose of) the embedding matrix E. Recall that in weight tying, we use the
same weights for two different matrices in the model. Thus at the input stage of the
transformer the embedding matrix (of shape [|V | × d]) is used to map from a one-hot
vector over the vocabulary (of shape [1 × |V |]) to an embedding (of shape [1 × d]).
And then in the language model head, ET , the transpose of the embedding matrix (of
shape [d × |V |]) is used to map back from an embedding (shape [1 × d]) to a vector
over the vocabulary (shape [1×|V |]). In the learning process, E will be optimized to
be good at doing both of these mappings. We therefore sometimes call the transpose
unembedding ET the unembedding layer because it is performing this reverse mapping.
A softmax layer turns the logits u into the probabilities y over the vocabulary.

u = hLN ET (8.46)
y = softmax(u) (8.47)

We can use these probabilities to do things like help assign a probability to a

given text. But the most important usage is to generate text, which we do by sam-
pling a word from these probabilities y. We might sample the highest probability
word (‘greedy’ decoding), or use another of the sampling methods from Section ??
or Section 8.6.
In either case, whatever entry yk we choose from the probability vector y, we
generate the word that has that index k.
Fig. 8.15 shows the total stacked architecture for one token i. Note that the input
to each transformer layer xi` is the same as the output from the preceding layer hi`−1 .
A terminological note before we conclude: You will sometimes see a trans-
former used for this kind of unidirectional causal language model called a decoder-
decoder-only only model. This is because this model constitutes roughly half of the encoder-
model
decoder model for transformers that we’ll see how to apply to machine translation
in Chapter 12. (Confusingly, the original introduction of the transformer had an
encoder-decoder architecture, and it was only later that the standard paradigm for
18 C HAPTER 8 • T RANSFORMERS

Token probabilities y1 y2 … y|V| wi+1

Sample token to
softmax generate at position i+1
Language
Modeling
Head logits u1 u2 … u|V|

hLi
feedforward
layer norm
Layer L
attention
layer norm
hL-1i = xLi
…
h2i = x3i
feedforward
layer norm
Layer 2
attention
layer norm

h1i = x2i
feedforward
layer norm
Layer 1
attention
layer norm

x1i
+ i
Input
Encoding E

Input token wi

Figure 8.15 A transformer language model (decoder-only), stacking transformer blocks

and mapping from an input token wi to to a predicted next token wi+1 .

causal language model was defined by using only the decoder part of this original
architecture).

8.6 More on Sampling

The sampling methods we introduce below each have parameters that enable trad-
ing off two important factors in generation: quality and diversity. Methods that
emphasize the most probable words tend to produce generations that are rated by
people as more accurate, more coherent, and more factual, but also more boring
and more repetitive. Methods that give a bit more weight to the middle-probability
words tend to be more creative and more diverse, but less factual and more likely to
be incoherent or otherwise low-quality.

8.6.1 Top-k sampling

top-k sampling Top-k sampling is a simple generalization of greedy decoding. Instead of choosing
the single most probable word to generate, we first truncate the distribution to the
 8.7 • T RAINING 19

top k most likely words, renormalize to produce a legitimate probability distribution,

and then randomly sample from within these k words according to their renormalized
probabilities. More formally:
1. Choose in advance a number of words k
2. For each word in the vocabulary V , use the language model to compute the
likelihood of this word given the context p(wt |w<t )
3. Sort the words by their likelihood, and throw away any word that is not one of
the top k most probable words.
4. Renormalize the scores of the k words to be a legitimate probability distribu-
tion.
5. Randomly sample a word from within these remaining k most-probable words
according to its probability.
When k = 1, top-k sampling is identical to greedy decoding. Setting k to a larger
number than 1 leads us to sometimes select a word which is not necessarily the most
probable, but is still probable enough, and whose choice results in generating more
diverse but still high-enough-quality text.

8.6.2 Nucleus or top-p sampling

One problem with top-k sampling is that k is fixed, but the shape of the probability
distribution over words differs in different contexts. If we set k = 10, sometimes
the top 10 words will be very likely and include most of the probability mass, but
other times the probability distribution will be flatter and the top 10 words will only
include a small part of the probability mass.
top-p sampling An alternative, called top-p sampling or nucleus sampling (Holtzman et al.,
2020), is to keep not the top k words, but the top p percent of the probability mass.
The goal is the same; to truncate the distribution to remove the very unlikely words.
But by measuring probability rather than the number of words, the hope is that the
measure will be more robust in very different contexts, dynamically increasing and
decreasing the pool of word candidates.
Given a distribution P(wt |w<t ), we sort the distribution from most probable, and
then the top-p vocabulary V (p) is the smallest set of words such that
X
P(w|w<t ) ≥ p. (8.48)
w∈V (p)

8.7 Training
We described the training process for language models in the prior chapter. Re-
call that large language models are trained with cross-entropy loss, also called the
negative log likelihood loss. At time t the cross-entropy loss is the negative log prob-
ability the model assigns to the next word in the training sequence, − log p(wt+1 ).
Fig. 8.16 illustrates the general training approach. At each step, given all the
preceding words, the final transformer layer produces an output distribution over the
entire vocabulary. During training, the probability assigned to the correct word by
the model is used to calculate the cross-entropy loss for each item in the sequence.
The loss for a training sequence is the average cross-entropy loss over the entire
sequence. The weights in the network are adjusted to minimize the average CE loss
over the training sequence via gradient descent.
20 C HAPTER 8 • T RANSFORMERS

Next token long and thanks for all …

log yand … =
<latexit sha1_base64="AovqpaL476UmJ1EU1xZPgDZ70tQ=">AAAB9nicbVDLSsNAFL2pr1pfURcu3AwWwY0lEakui25cVrAPaEqYTCbt0EkmzEzEEvIrbkTcKPgZ/oJ/Y9Jm09YDA4dzznDvPV7MmdKW9WtU1tY3Nreq27Wd3b39A/PwqKtEIgntEMGF7HtYUc4i2tFMc9qPJcWhx2nPm9wXfu+ZSsVE9KSnMR2GeBSxgBGsc8k1Ty4dLkZo6qZOiPVYhimO/CyruWbdalgzoFVil6QOJdqu+eP4giQhjTThWKmBbcV6mGKpGeE0qzmJojEmEzyi6WztDJ3nko8CIfMXaTRTF3I4VGoaenmy2E0te4X4nzdIdHA7TFkUJ5pGZD4oSDjSAhUdIJ9JSjSf5gQTyfINERljiYnOmypOt5cPXSXdq4bdbDQfr+utu7KEKpzCGVyADTfQggdoQwcIZPAGn/BlvBivxrvxMY9WjPLPMSzA+P4DPEiSHA==</latexit>

log ythanks
<latexit sha1_base64="q3ZgXDyG7qtkT7t8hT47RdlwYG4=">AAAB+XicbVDLSsNAFJ3UV62vWHe6GVsEN5bERXUlBUVcVrAPaEqYTCft0MlMmJkIIQT8AT/CTRE3Cv6Ev+DfmLTdtPXAwOGcM9x7jxcyqrRl/RqFtfWNza3idmlnd2//wDwst5WIJCYtLJiQXQ8pwignLU01I91QEhR4jHS88W3ud56JVFTwJx2HpB+gIac+xUhnkmseXzhMDGHsJk6A9EgGiR4hPlZpWnLNqlWzpoCrxJ6TauP0tXw3qdw0XfPHGQgcBYRrzJBSPdsKdT9BUlPMSFpyIkVChMdoSJLp5ik8y6QB9IXMHtdwqi7kUKBUHHhZMl9PLXu5+J/Xi7R/3U8oDyNNOJ4N8iMGtYB5DXBAJcGaxRlBWNJsQ4hHSCKss7Ly0+3lQ1dJ+7Jm12v1x6yDezBDEZyACjgHNrgCDfAAmqAFMHgBE/AJvozEeDPejY9ZtGDM/xyBBRjff79pldo=</latexit>

Loss

Language
Modeling
logits logits logits logits logits …
Head U U U U U

Stacked
Transformer
… … … … … …
Blocks

x1 x2 x3 x4 x5 …
+ 1 + 2 + 3 + 4 + 5
Input
Encoding E E E E E
…

Input tokens So long and thanks for

Figure 8.16 Training a transformer as a language model.

With transformers, each training item can be processed in parallel since the out-
put for each element in the sequence is computed separately.
Large models are generally trained by filling the full context window (for exam-
ple 4096 tokens for GPT4 or 8192 for Llama 3) with text. If documents are shorter
than this, multiple documents are packed into the window with a special end-of-text
token between them. The batch size for gradient descent is usually quite large (the
largest GPT-3 model uses a batch size of 3.2 million tokens).

8.8 Dealing with Scale

Large language models are large. For example the Llama 3.1 405B Instruct model
from Meta has 405 billion parameters (L=126 layers, a model dimensionality of
d=16,384, A=128 attention heads) and was trained on 15.6 terabytes of text tokens
(Llama Team, 2024), using a vocabulary of 128K tokens. So there is a lot of research
on understanding how LLMs scale, and especially how to implement them given
limited resources. In the next few sections we discuss how to think about scale (the
concept of scaling laws), and important techniques for getting language models to
work efficiently, such as the KV cache and parameter-efficient fine tuning.

8.8.1 Scaling laws

The performance of large language models has shown to be mainly determined by
3 factors: model size (the number of parameters not counting embeddings), dataset
size (the amount of training data), and the amount of compute used for training. That
is, we can improve a model by adding parameters (adding more layers or having
wider contexts or both), by training on more data, or by training for more iterations.
The relationships between these factors and performance are known as scaling
scaling laws laws. Roughly speaking, the performance of a large language model (the loss) scales
 8.8 • D EALING WITH S CALE 21

as a power-law with each of these three properties of model training.

For example, Kaplan et al. (2020) found the following three relationships for
loss L as a function of the number of non-embedding parameters N, the dataset size
D, and the compute budget C, for models training with limited parameters, dataset,
or compute budget, if in each case the other two properties are held constant:

Nc
 αN
L(N) = (8.49)
N
Dc
  αD
L(D) = (8.50)
D
Cc
 αC
L(C) = (8.51)
C

The number of (non-embedding) parameters N can be roughly computed as fol-

lows (ignoring biases, and with d as the input and output dimensionality of the
model, dattn as the self-attention layer size, and dff the size of the feedforward layer):

N ≈ 2 d nlayer (2 dattn + dff )

≈ 12 nlayer d 2 (8.52)
(assuming dattn = dff /4 = d)

Thus GPT-3, with n = 96 layers and dimensionality d = 12288, has 12 × 96 ×

122882 ≈ 175 billion parameters.
The values of Nc , Dc , Cc , αN , αD , and αC depend on the exact transformer
architecture, tokenization, and vocabulary size, so rather than all the precise values,
scaling laws focus on the relationship with loss.4
Scaling laws can be useful in deciding how to train a model to a particular per-
formance, for example by looking at early in the training curve, or performance with
smaller amounts of data, to predict what the loss would be if we were to add more
data or increase model size. Other aspects of scaling laws can also tell us how much
data we need to add when scaling up a model.

8.8.2 KV Cache
We saw in Fig. 8.10 and in Eq. 8.34 (repeated below) how the attention vector can
be very efficiently computed in parallel for training, via two matrix multiplications:

QK|
 
A = softmax √ V (8.53)
dk

Unfortunately we can’t do quite the same efficient computation in inference as

in training. That’s because at inference time, we iteratively generate the next tokens
one at a time. For a new token that we have just generated, call it xi , we need to
compute its query, key, and values by multiplying by WQ , WK , and WV respec-
tively. But it would be a waste of computation time to recompute the key and value
vectors for all the prior tokens x<i ; at prior steps we already computed these key
and value vectors! So instead of recomputing these, whenever we compute the key
KV cache and value vectors we store them in memory in the KV cache, and then we can just
grab them from the cache when we need them. Fig. 8.17 modifies Fig. 8.10 to show
22 C HAPTER 8 • T RANSFORMERS

Q QKT V A
KT v1

x = x v2

k1

k2

k3

k4
=
v3

dk x N v4
q4 q4•k1 q4•k2 q4•k3 q4•k4 a4

1 x dk 1xN N x dv 1 x dv

Figure 8.17 Parts of the attention computation (extracted from Fig. 8.10) showing, in black,
the vectors that can be stored in the cache rather than recomputed when computing the atten-
tion score for the 4th token.

the computation that takes place for a single new token, showing which values we
can take from the cache rather than recompute.

8.8.3 Parameter Efficient Fine Tuning

As we mentioned above, it’s very common to take a language model and give it more
information about a new domain by finetuning it (continuing to train it to predict
upcoming words) on some additional data.
Fine-tuning can be very difficult with very large language models, because there
are enormous numbers of parameters to train; each pass of batch gradient descent
has to backpropagate through many many huge layers. This makes finetuning huge
language models extremely expensive in processing power, in memory, and in time.
For this reason, there are alternative methods that allow a model to be finetuned
parameter-
without changing all the parameters. Such methods are called parameter-efficient
efficient fine fine tuning or sometimes PEFT, because we efficiently select a subset of parameters
tuning
PEFT to update when finetuning. For example we freeze some of the parameters (don’t
change them), and only update some particular subset of parameters.
LoRA Here we describe one such model, called LoRA, for Low-Rank Adaptation. The
intuition of LoRA is that transformers have many dense layers which perform matrix
multiplication (for example the WQ , WK , WV , WO layers in the attention computa-
tion). Instead of updating these layers during finetuning, with LoRA we freeze these
layers and instead update a low-rank approximation that has fewer parameters.
Consider a matrix W of dimensionality [N × d] that needs to be updated during
finetuning via gradient descent. Normally this matrix would get updates ∆W of
dimensionality [N × d], for updating the N × d parameters after gradient descent. In
LoRA, we freeze W and update instead a low-rank decomposition of W. We create
two matrices A and B, where A has size [N ×r] and B has size [r ×d], and we choose
r to be quite small, r << min(d, N). During finetuning we update A and B instead
of W. That is, we replace W + ∆W with W + AB. Fig. 8.18 shows the intuition.
For replacing the forward pass h = xW, the new forward pass is instead:

h = xW + xAB (8.54)

LoRA has a number of advantages. It dramatically reduces hardware requirements,

since gradients don’t have to be calculated for most parameters. The weight updates
can be simply added in to the pretrained weights, since AB is the same size as W).
4 For the initial experiment in Kaplan et al. (2020) the precise values were αN = 0.076, Nc = 8.8 ×1013
(parameters), αD = 0.095, Dc = 5.4 ×1013 (tokens), αC = 0.050, Cc = 3.1 ×108 (petaflop-days).
 8.9 • I NTERPRETING THE T RANSFORMER 23

d
h 1

d
× r B
Pretrained
Weights
N N A
W

d r

x 1
d

Figure 8.18 The intuition of LoRA. We freeze W to its pretrained values, and instead fine-
tune by training a pair of matrices A and B, updating those instead of W, and just sum W and
the updated AB.

That means it doesn’t add any time during inference. And it also means it’s possible
to build LoRA modules for different domains and just swap them in and out by
adding them in or subtracting them from W.
In its original version LoRA was applied just to the matrices in the attention
computation (the WQ , WK , WV , and WO layers). Many variants of LoRA exist.

8.9 Interpreting the Transformer

How does a transformer-based language model manage to do so well at language
interpretability tasks? The subfield of interpretability, sometimes called mechanistic interpretabil-
ity, focuses on ways to understand mechanistically what is going on inside the
transformer. In the next two subsections we discuss two well-studied aspects of
transformer interpretability.

8.9.1 In-Context Learning and Induction Heads

As a way of getting a model to do what we want, we can think of prompting as being
fundamentally different than pretraining. Learning via pretraining means updating
the model’s parameters by using gradient descent according to some loss function.
But prompting with demonstrations can teach a model to do a new task. The model
is learning something about the task from those demonstrations as it processes the
prompt.
Even without demonstrations, we can think of the process of prompting as a kind
of learning. For example, the further a model gets in a prompt, the better it tends
to get at predicting the upcoming tokens. The information in the context is helping
give the model more predictive power.
in-context
learning The term in-context learning was first proposed by Brown et al. (2020) in their
introduction of the GPT3 system, to refer to either of these kinds of learning that lan-
24 C HAPTER 8 • T RANSFORMERS

guage models do from their prompts. In-context learning means language models
learning to do new tasks, better predict tokens, or generally reduce their loss dur-
ing the forward-pass at inference-time, without any gradient-based updates to the
model’s parameters.
How does in-context learning work? While we don’t know for sure, there are
induction heads some intriguing ideas. One hypothesis is based on the idea of induction heads
(Elhage et al., 2021; Olsson et al., 2022). Induction heads are the name for a circuit,
which is a kind of abstract component of a network. The induction head circuit
is part of the attention computation in transformers, discovered by looking at mini
language models with only 1-2 attention heads.
The function of the induction head is to predict repeated sequences. For example
if it sees the pattern AB...A in an input sequence, it predicts that B will follow,
instantiating the pattern completion rule AB...A→ B. It does this by having a prefix
matching component of the attention computation that, when looking at the current
token A, searches back over the context to find a prior instance of A. If it finds one,
the induction head has a copying mechanism that “copies” the token B that followed
the earlier A, by increasing the probability the B will occur next. Fig. 8.19 shows an
example.

Figure
Figure 1:8.19 An induction
In the sequence head
“...vintage looking
cars ... vintage”,atanvintage uses
induction head the prefix
identifies matching
the initial mechanism
occurrence of “vintage”,to
find a prior
attends to theinstance
subsequentofword “cars” forand
vintage, prefixthe copying
matching, and mechanism
predicts “cars” to predict
as the next wordthatthrough will
cars the occur
copying
mechanism.
again. Figure from Crosbie and Shutova (2022).

determines each head’s independent output for the 4.2 Identifying Induction Heads
Olsson et al. (2022) propose that a generalized fuzzy version of this pattern com-
current token. To identify
pletion rule, implementing a rule like A*B*...A→
Leveraging this decomposition, Elhage et al. sure the ability B,induction
where heads
A* ≈within
A and models,
B* ≈ weB mea-
(by
of all attention heads to perform
we mean
≈(2021) theyathey
discovered arebehaviour
distinct semantically
in certainsimilar in some way), might be responsible
prefix matching on random input sequences. We 4
forattention
in-context learning.
heads, which Suggestive
they named induction evidence for the
heads. follow their hypothesis
task-agnostic comes
approach from Cros-
to computing pre-
ablating This
bie andbehaviour
Shutova emerges when who
(2022), these heads
showprocess
that ablating induction
fix matching headsbycauses
scores outlined Bansal etin-context
al. (2023).
sequences of the form "[A] [B] ... [A] → ". In Weis argue that focusing solely on term
prefix matching
learning performance to decrease. Ablation originally a medical
these heads, the QK circuit directs attention to- scores is sufficient for our analysis, as high pre-
meaning
the removal
wards [B], whichofappears
something. Wetheuse
directly after it in NLP
previous interpretability studies as a tool for
fix matching cores specifically indicate induction
testing
occurrencecausal
of theeffects; if we
current token [A].knock out a hypothesized
This behaviour heads, while less cause,
relevantwe would
heads tend toexpect
show high the
is termed
effect prefix matching.
to disappear. The OV
Crosbie andcircuit subse- (2022)
Shutova copyingablate induction
capabilities (Bansalheads by first
et al., 2023). We find-
gen-
quently increases the output logit of the [B] token, erate a sequence of 50 random tokens, excluding
ing attention heads that perform as induction heads on random input sequences, and
termed copying. An overview of this mechanism is the 4% most common and least common tokens.
then
shown zeroing
in Figureout
1. the output of these heads by setting certain terms of the output ma-
This sequence is repeated four times to form the
trix WO to zero. Indeed they find that ablated
inputmodels are The
to the model. much worse
prefix at in-context
matching score is cal-
4 Methods
learning: they have much worse performance at learning
culated from
by averaging the demonstrations
attention values fromin the
each
prompts. token to the tokens that directly followed the same
4.1 Models
token in earlier repeats. The final prefix matching
We utilise two recently developed open-source scores are averaged over five random sequences.
8.9.2 Logit
models, namely Lens 2 and InternLM2-20B
Llama-3-8B The prefix matching scores for Llama-3-8B are
(Cai et al., 2024), both of which are based on the shown in Figure 2. For IntermLM2-20B, we refer
logit lens original Llama
Another useful(Touvron et al., 2023a)
interpretability tool, the logit
architec- lens 8(Nostalgebraist,
to Figure in Appendix A.1. Both 2020),
modelsoffers
exhibit a
ture. These models feature grouped-query atten- heads with notably high prefix matching scores,
way to visualize what the internal layers of the transformer might be representing.
tion mechanisms (Ainslie et al., 2023) to enhance distributed across various layers. In the Llama-3-
The idea
efficiency. is that we
Llama-3-8B, take any
comprises vector
32 layers, eachfrom8Bany layer
model, ~3% of theheads
of the transformer
have a prefixand, pre-
matching
tending that it is
with 32 attention theand
heads prefinal
it uses aembedding,
query group simply
score of multiply by the unembedding
it indicating
0.3 or higher, a degree of spe-
layer
size ofto4 get
attention heads.
logits, andIt compute
has shown a superior
softmaxcialisation
to see the distribution
in prefix matching, andover
somewords that
heads have
performance compared to its predecessors, even high scores of up to 0.98.
that vector might be
the larger Llama-2 models.
representing. This can be a useful window into the internal
representations
InternLM2-20B,offeaturing
the model. Since
48 layers with the
48 at-network wasn’t
4.3 Head trained to make the internal
Ablations
tention heads each, uses a query group size of 6 To investigate the significance of induction heads
attention heads. We selected InternLM2-20B for for a specific ICL task, we conduct zero-ablations
its exemplary performance on the Needle-in-the- of 1% and 3% of the heads with the highest prefix
Haystack3 task, which assesses LLMs’ ability to matching scores. This ablation process involves
retrieve a single critical piece of information em- masking the corresponding partition of the output
bedded within a lengthy text. This mirrors the matrix, denoted as Woh in Eq. 1, by setting it to
functionality of induction heads, which scan the zero. This effectively renders the heads inactive
 8.10 • S UMMARY 25

representations function in this way, the logit lens doesn’t always work perfectly, but
this can still be a useful trick to help us visualize the internal layers of a transformer.

8.10 Summary
This chapter has introduced the transformer and its components for the language
modeling task introduced in the previous chapter. Here’s a summary of the main
points that we covered:
• Transformers are non-recurrent networks based on multi-head attention, a
kind of self-attention. A multi-head attention computation takes an input
vector xi and maps it to an output ai by adding in vectors from prior tokens,
weighted by how relevant they are for the processing of the current word.
• A transformer block consists of a residual stream in which the input from
the prior layer is passed up to the next layer, with the output of different com-
ponents added to it. These components include a multi-head attention layer
followed by a feedforward layer, each preceded by layer normalizations.
Transformer blocks are stacked to make deeper and more powerful networks.
• The input to a transformer is computed by adding an embedding (computed
with an embedding matrix) to a positional encoding that represents the se-
quential position of the token in the window.
• Language models can be built out of stacks of transformer blocks, with a
language model head at the top, which applies an unembedding matrix to
the output H of the top layer to generate the logits, which are then passed
through a softmax to generate word probabilities.
• Transformer-based language models have a wide context window (200K to-
kens or even more for very large models with special mechanisms) allowing
them to draw on enormous amounts of context to predict upcoming words.
• There are various computational tricks for making large language models
more efficient, such as the KV cache and parameter-efficient finetuning.

Historical Notes
The transformer (Vaswani et al., 2017) was developed drawing on two lines of prior
research: self-attention and memory networks.
Encoder-decoder attention, the idea of using a soft weighting over the encodings
of input words to inform a generative decoder (see Chapter 12) was developed by
Graves (2013) in the context of handwriting generation, and Bahdanau et al. (2015)
for MT. This idea was extended to self-attention by dropping the need for separate
encoding and decoding sequences and instead seeing attention as a way of weighting
the tokens in collecting information passed from lower layers to higher layers (Ling
et al., 2015; Cheng et al., 2016; Liu et al., 2016).
Other aspects of the transformer, including the terminology of key, query, and
value, came from memory networks, a mechanism for adding an external read-
write memory to networks, by using an embedding of a query to match keys rep-
26 C HAPTER 8 • T RANSFORMERS

resenting content in an associative memory (Sukhbaatar et al., 2015; Weston et al.,

2015; Graves et al., 2014).
MORE HISTORY TBD IN NEXT DRAFT.
 Historical Notes 27

Ba, J. L., J. R. Kiros, and G. E. Hinton. 2016. Layer normal-

ization. NeurIPS workshop.
Bahdanau, D., K. H. Cho, and Y. Bengio. 2015. Neural ma-
chine translation by jointly learning to align and translate.
ICLR 2015.
Brown, T., B. Mann, N. Ryder, M. Subbiah, J. Kaplan,
P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry,
A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger,
T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu,
C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin,
S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish,
A. Radford, I. Sutskever, and D. Amodei. 2020. Language
models are few-shot learners. NeurIPS, volume 33.
Cheng, J., L. Dong, and M. Lapata. 2016. Long short-term
memory-networks for machine reading. EMNLP.
Crosbie, J. and E. Shutova. 2022. Induction heads as an
essential mechanism for pattern matching in in-context
learning. ArXiv preprint.
Elhage, N., N. Nanda, C. Olsson, T. Henighan, N. Joseph,
B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, N. Das-
Sarma, D. Drain, D. Ganguli, Z. Hatfield-Dodds, D. Her-
nandez, A. Jones, J. Kernion, L. Lovitt, K. Ndousse,
D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCan-
dlish, and C. Olah. 2021. A mathematical framework for
transformer circuits. White paper.
Graves, A. 2013. Generating sequences with recurrent neural
networks. ArXiv.
Graves, A., G. Wayne, and I. Danihelka. 2014. Neural Tur-
ing machines. ArXiv.
Holtzman, A., J. Buys, L. Du, M. Forbes, and Y. Choi. 2020.
The curious case of neural text degeneration. ICLR.
Kaplan, J., S. McCandlish, T. Henighan, T. B. Brown,
B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and
D. Amodei. 2020. Scaling laws for neural language mod-
els. ArXiv preprint.
Ling, W., C. Dyer, A. W. Black, I. Trancoso, R. Fermandez,
S. Amir, L. Marujo, and T. Luı́s. 2015. Finding function
in form: Compositional character models for open vocab-
ulary word representation. EMNLP.
Liu, Y., C. Sun, L. Lin, and X. Wang. 2016. Learning natural
language inference using bidirectional LSTM model and
inner-attention. ArXiv.
Llama Team. 2024. The llama 3 herd of models.
Nostalgebraist. 2020. Interpreting gpt: the logit lens. White
paper.
Olsson, C., N. Elhage, N. Nanda, N. Joseph, N. DasSarma,
T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, et al.
2022. In-context learning and induction heads. ArXiv
preprint.
Sukhbaatar, S., A. Szlam, J. Weston, and R. Fergus. 2015.
End-to-end memory networks. NeurIPS.
Uszkoreit, J. 2017. Transformer: A novel neural network ar-
chitecture for language understanding. Google Research
blog post, Thursday August 31, 2017.
Vaswani, A., N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones,
A. N. Gomez, Ł. Kaiser, and I. Polosukhin. 2017. Atten-
tion is all you need. NeurIPS.
Weston, J., S. Chopra, and A. Bordes. 2015. Memory net-
works. ICLR 2015.