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

All
rights reserved. Draft of January 6, 2026.

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.

output probabilities (over tokens)

Softmax

Unembedding

+
Feedforward

Layer Norm
residual N times
stream +
MultiHead
Attention

Layer Norm

+ Positional

Embedding

input token

Figure 8.1 A transformer decoder for language modeling, showing the residual stream for
processing an input token. A single token is embedded and passed forward in the network,
with the feedforward and attention components adding information. The multihead attention
layer takes inputs (not shown in detail) from the neighboring token streams. This is thus one
column of an autoregressive transformer language model, taking an input token and outputting
a distribution over next tokens.

Fig. 8.1 sketches the transformer architecture following a single token as it is

passes up through the layers of the network. Each token is first converted to an
embedding from the embedding matrix E. Recall from Chapter 6 in Section ??
that E is a linear layer that maps a token id to a vector embedding representing that
token. Each token in the vocabulary has an initial embedding representation in E.
2 C HAPTER 8 • T RANSFORMERS

Transformers also have a special mechanism for encoding the position/index of the
token in the input string, which is simply added to the embedding. The resulting
embedding represents both the word and its position. and is then passed through a
set of N transformer blocks.
It’s common to think of each of these transformer blocks as part of a stream in
which the input embedding is directly passed up to the output, while simultaneously
being enriched by the application of various processing modules: the multi-head
attention layer, feedforward networks and the layer normalization. The value of the
stream at any layer is the sum of the original embedding and all the outputs from all
the previous layers and blocks.
The core intuition of the transformer, and the component that distinguishes it
from the feedforward layers we saw in Chapter 6, is this multi-head attention layer,
also called a self-attention layer. 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. It can also be thought of as a way to move information
from one residual stream to another, augmenting the stream at one token position
with information from another token position.
After the N transformer blocks we take the output embedding that is produced
by the final transformer block, pass it through an linear unembedding matrix U
and then a softmax over the vocabulary to generate a distribution over possible next
tokens. These last two components (the unembedding matrix and the softmax) are
sometimes called the language modeling head. In the rest of this chatper we’ll
introduce attention and the rest of these modules in more detail.
Fig. ?? shows the transformer architecture applied to a context window with
the words So long and thanks for, showing at each token position what is the
most likely token to be generated. In this full figure, 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 arrows in the figure shows how information from the hidden
representations of preceding tokens is incorporated into the transformer block.
Transformer-based language models are complex, and so the details will unfold
over this chapter and the next few chapters. Chapter 7 already discussed how lan-
guage models are pretrained, and how tokens are generated via sampling. In the
rest of this chapter we’ll introduce multi-head attention, the rest of the transformer
block, and the input encoding and language modeling head components of the trans-
former. Chapter 9 introduces masked language modeling and the BERT family of
bidirectional transformer encoder models. Chapter 10 shows how to instruction-
tune language models to perform NLP tasks, and how to align the model with hu-
man preferences. Chapter 12 will introduce machine translation with the encoder-
decoder architecture. And we’ll see application of the transformer to speech recog-
nition, as well as 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
 8.1 • ATTENTION 3

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.2 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. The embeddings at each token position in the residual
stream are passed up the stack, and the arrows in the figure shows how information from the
hidden representations of preceding tokens are also incorporated.

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).1 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.
(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 9.)
1 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.
4 C HAPTER 8 • T RANSFORMERS

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.3 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.3 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.

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
 8.1 • ATTENTION 5

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 9 we’ll generalize attention
so it can also look ahead to future words.)
Fig. 8.4 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.4 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.4 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,
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
6 C HAPTER 8 • T RANSFORMERS

into a probability distribution used to weight the sum of the prior vectors. 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)

We illustrate this in Fig. 8.5 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 left-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
 8.1 • ATTENTION 7

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.5 Calculating the value of a3 , the third element of a sequence using causal (left-
to-right) self-attention.

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
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 =
8 C HAPTER 8 • T RANSFORMERS

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.6
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 9 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.7 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-
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
 8.2 • T RANSFORMER B LOCKS 9

ai
[1 x d]

[Adv x d]
Project to final representation WO usually dv=d/A, hence [d x d]

Concatenate Outputs [1 x Adv ] here A=4

[1 x dv ] [1 x dv ] [1 x dv ] [1 x dv ]
Each head
attends differently Head 1 Head 2 Head 3 Head 4
K4
to context WK1 WV1 WQ1 WK2 WV2 WQ2 WK3 WV3W
WQ3 WK3 WV4 WQ4

… xi-3 xi-2 xi-1 xi

[1 x d]
Figure 8.6 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. In this figure, we show A = 4, a
smaller value than is usually used, just to fit on the page. The outputs from each of the heads are of shape [1×dv ]
and are concatenated and then projected into a different space by the WO matrix. Usually the dimensionality
dv of the heads is set so that dv = d/A, with the result that WO is a square matrix of shape [Adv × d] = [d × d].
usually of the same size. then projected d, thus producing an output of the same size as the input.

hi-2 hi-1 hi

+
Feedforward

Layer Norm

… +
MultiHead
Attention

Layer Norm

xi-2 xi-1 xi

Figure 8.7 The architecture of a transformer block showing the residual stream, showing
how most information flows up through the residual stream, and only the attention module
is sensitive to information from other streams at prior token positions. In this figure and
throughout the chapter, we use the prenorm version of the architecture, in which the layer
norms happen before the attention and feedforward layers rather than after. The first

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.
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
10 C HAPTER 8 • T RANSFORMERS

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
each computation inside the block:
t1i = LayerNorm(xi ) (8.26)
t2i = MultiHeadAttention(t1i , t11 , · · · , t1N )
 
(8.27)
3 2
ti = ti + xi (8.28)
ti = LayerNorm(ti )
4 3
(8.29)
ti = FFN(ti )
5 4
(8.30)
hi = t5i + t3i (8.31)
 8.2 • T RANSFORMER B LOCKS 11

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.8 shows a visualization of this movement. We therefore call the attention func-
token-mixing tion the token-mixing component of the architecture, because it mixes information
from neighboring token streams into the current stream.

Token A Token B
residual residual
stream stream

Figure 8.8 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). 2

2 Note that we are using the most common current transformer architecture, which is called the prenorm
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

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.9.

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.9 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
following computation:
QK|
  
head = softmax mask √ V (8.33)
dk
A = head WO (8.34)
 8.3 • PARALLELIZING COMPUTATION USING A SINGLE MATRIX X 13

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.10 shows the resulting masked QK| matrix. (we’ll see in Chapter 9 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.10 The N × N QK| matrix showing the qi · k j values, with the upper-triangle
portion of the comparisons matrix zeroed out (set to −∞, which the softmax will turn to
zero).

Fig. 8.11 shows a schematic of all the computations for a single attention head
parallelized in matrix form.

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.11 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.
14 C HAPTER 8 • T RANSFORMERS

Fig. 8.9 and Fig. 8.10 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]
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
embedding:
T1 = LayerNorm(X) (8.40)
T 2
= MultiHeadAttention(T )1
(8.41)
3 2
T = T +X (8.42)
T 4
= LayerNorm(T )3
(8.43)
T5 = FFN(T4 ) (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
 8.4 • T HE INPUT: EMBEDDINGS FOR TOKEN AND POSITION 15

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.12.

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

|V|

Figure 8.12 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.13.
16 C HAPTER 8 • T RANSFORMERS

d
|V| d
0000100…0000
0000000…0010
1000000…0000 ✕ E =
…
N 0000100…0000
N
| V|
Figure 8.13 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.14
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.14 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.

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
 8.5 • T HE L ANGUAGE M ODELING H EAD 17

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.15 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.15 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
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.
18 C HAPTER 8 • T RANSFORMERS

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

Language Model Head Softmax 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 [1 x d]

Layer L
Transformer
Block
…
w1 w2 wN

Figure 8.15 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 .

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.16 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 h`−1
i .
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
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
 8.6 • M ORE ON S AMPLING 19

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

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

hLi
extra layer norm
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
E Encoding

Input token wi
Figure 8.16 A transformer language model (decoder-only), stacking transformer blocks
and mapping from an input token wi to a predicted next token wi+1 .

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
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-
20 C HAPTER 8 • T RANSFORMERS

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.17 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.
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 • D EALING WITH S CALE 21

Next token long and thanks for all …

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

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

log ythanks … =

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.17 Training a transformer as a language model.

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 (it has L=126 layers, model dimensionality
d=16,384, and A=128 attention heads) and was trained on 15.6 terabytes of text
tokens using a vocabulary of 128K tokens (Llama Team, 2024). 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
(PEFT).

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
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,
22 C HAPTER 8 • T RANSFORMERS

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.3
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.11 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.18 modifies Fig. 8.11 to show
the computation that takes place for a single new token, showing which values we
can take from the cache rather than recompute.
3 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.8 • D EALING WITH S CALE 23

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.18 Parts of the attention computation (extracted from Fig. 8.11) 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.

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 [k × d] that needs to be updated during
finetuning via gradient descent. Normally this matrix would get updates ∆W of
dimensionality [k × d], for updating the k × 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 [k × r] and B has size [r × d], and we choose
r to be quite small, r << min(d, k). During finetuning we update A and B instead
of W. That is, we replace W + ∆W with W + AB. Fig. 8.19 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).
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.
24 C HAPTER 8 • T RANSFORMERS

d
h 1

d
× r B
Pretrained
Weights
k k A
W

d r

x 1
d

Figure 8.19 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.

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-
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,
 8.9 • I NTERPRETING THE T RANSFORMER 25

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.20 shows an
example.

Figure
Figure 1:8.20 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
current token.
Olsson et al. (2022) propose that a generalized To identifyfuzzy version
induction heads of thismodels,
within pattern we com-
mea-
Leveraging this decomposition, Elhage et al. sure the ability of all attention heads to perform
pletion rule, implementing a rule
(2021) discovered a distinct behaviour in certain
like A*B*...A→ B, where A* ≈ A and B* ≈B
prefix matching on random input sequences.4 We
(by ≈ weheads,
attention mean they
which theyare
namedsemantically
induction heads. similar in some way), might be responsible
follow the task-agnostic approach to computing pre-
for in-context
This learning.
behaviour emerges whenSuggestive evidencefixfor
these heads process theirscores
matching hypothesis comes
outlined by Bansalfrom Cros-
et al. (2023).
ablating sequences
bie of the form
and Shutova "[A] [B]
(2022), who [A] → that
... show ". Inablating
We argue induction heads
that focusing solelycauses
on prefixin-context
matching
these heads, the QK circuit directs attention to- scores is sufficient for our analysis, as high pre-
learning performance to decrease. Ablation is originally a medical term meaning
wards [B], which appears directly after the previous fix matching cores specifically indicate induction
the removal
occurrence of current
of the something.
token [A].We Thisuse it in NLP
behaviour interpretability studies as a tool for
heads, while less relevant heads tend to show high
testing
is termedcausal
prefix effects;
matching. ifTheweOVknockcircuit out a hypothesized
subse- cause,
copying capabilities we etwould
(Bansal expect
al., 2023). the
We gen-
quently
effect toincreases the output
disappear. logit of
Crosbie and theShutova
[B] token, (2022)
erate aablate
sequenceinduction heads
of 50 random by first
tokens, find-
excluding
termed copying. An overview of this mechanism is
ing attention heads that perform as inductiontheheads
shown in Figure 1.
4% most on common
randomand leastsequences,
input common tokens. and
This sequence is repeated four times to form the
then zeroing out the output of these heads by setting certain terms of the output ma-
input to the model. The prefix matching score is cal-
trix
4 W O
to zero. Indeed they find that ablated
Methods models
culated are much
by averaging worsevalues
the attention at in-context
from each
learning: they have much worse performance token to the tokens that directly followed the in
at learning from demonstrations the
same
4.1 Models
prompts. token in earlier repeats. The final prefix matching
We utilise two recently developed open-source scores are averaged over five random sequences.
models, namely Llama-3-8B2 and InternLM2-20B 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
8.9.2 Logit
original Llama Lenset al., 2023a) architec- to Figure 8 in Appendix A.1. Both models exhibit
(Touvron
ture. These models feature grouped-query atten- heads with notably high prefix matching scores,
tion mechanisms (Ainslie et al., 2023) to enhance distributed across various layers. In the Llama-3-
logit lens Another
[Link] interpretability
Llama-3-8B, the logit
tool, each
comprises 32 layers, lens (Nostalgebraist, 2020), offers a
8B model, ~3% of the heads have a prefix matching
way
withto32visualize whatandthe
attention heads internal
it uses layers
a query group of score
the transformer
of 0.3 or higher, might be representing.
indicating a degree of spe-
The
size of 4idea is that
attention weIttake
heads. any vector
has shown superiorfrom any layer
cialisation of the
in prefix transformer
matching, and,have
and some heads pre-
performance
tending that compared
it is the toprefinal
its predecessors, even simply
embedding, high scores of up to 0.98.
multiply it by the unembedding
the larger Llama-2 models.
layer to get logits, and compute a softmax to see the distribution over words that
InternLM2-20B, featuring 48 layers with 48 at- 4.3 Head Ablations
that vector
tention headsmight be arepresenting.
each, uses query group sizeThis of 6 can be a useful
To investigate window into
the significance the internal
of induction heads
representations
attention heads. We of selected
the model. Since thefornetwork
InternLM2-20B wasn’t
for a specific ICLtrained
task, weto makezero-ablations
conduct the internal
its exemplary performance
representations functiononinthe Needle-in-the-
this way, the logitoflens
1% and doesn’t
3% of thealways worktheperfectly,
heads with highest prefixbut
Haystack3 task, which assesses LLMs’ ability to matching scores. This ablation process involves
this can still be a useful trick to help us visualize the internal layers of a
retrieve a single critical piece of information em- masking the corresponding partition of the output
transformer.
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
context for prior occurrences of a token to extract 4
In this work, the term "induction heads" refers to what
relevant subsequent information. we define as behavioural induction heads, not mechanistic
ones. A true induction head must be verified mechanistically;
2
[Link] however, our analysis employs prefix-matching scores as a
3
[Link] proxy. We will continue to use the term "induction heads" for
NeedleInAHaystack simplicity throughout the rest of the paper.

4
26 C HAPTER 8 • T RANSFORMERS

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-
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.