How Llama 3.1 Generates the Next Token

Introduction

I’ve recently been reading a lot about mechanistic interpretability. Like a significant fraction of AI researchers in the LLM era, I found the idea that scaling up compute and training models to predict the next token could eventually produce systems capable of chipping away at math problems on the frontierof human research both astounding and utterly counter-intuitive.

As a first step towards understanding these models better, I wanted to understand exactly how a relatively conventional, dense LLM works from end to end. No mixture-of-experts (MoE), state-space models like Mamba, or other fancy architectural tweaks. Llama 3.1 seemed like a good place to start since It’s one of the last generations of prominent open-weight models from before everyone pivoted to architectures like MoE. It follows a relatively straightforward, decoder-only transformer architecture, is well documented, and, at least in its 8B variant, can be run locally without too much difficulty.

So how does Llama 3.1 actually generate text?

At the highest level, the process is exactly what accessible explanation says. If you give them model some input text, it processes the sequence and produces a set of scores representing how likely each token in its vocabulary is to come next. These scores are converted into a probability distribution, from which we select the next token. We then append this newly generated token to the input sequence and repeat the process, generating text one token at a time. This continues until the model generates a special token indicating that it’s finished, or until some other stopping condition is reached. This process is called autoregressive generation, and it’s the basic mechanism underlying how most modern LLMs generate text.

For the rest of this writeup, I will try to break down exactly what happens between the text input and the decision of what token comes next. The figures on this page use the Llama 3.1 8B configuration: 32 layers, a residual stream of width 4,096, and a vocabulary of 128,000 tokens.¹ The walkthrough follows one token from characters on the screen to that draw.²

A bird’s eye view of the architecture

1. 01

   Tokenization

   The prompt is split into subword token ids from a vocabulary of 128,256.
2. 02

   Embedding

   Each id is looked up as a vector of 4,096 numbers, with position from RoPE.
3. 03

   Transformer blocks

   The sequence passes through 32 blocks.
4. 04

   Unembedding

   The last position becomes a score for every token in the vocabulary.
5. 05

   Next token

   A sampling rule picks one token from that distribution.

If you zoom in a bit and focus on just the essentials, the model’s processing pipeline can be divided into five stages. First, the input text is converted into tokens through a process called tokenization. Tokens can be thought of as units of text — not necessarily whole words — that make up the model’s vocabulary. I’ll speak more on this in later sections. These tokens then pass through an embedding layer, which converts each token into a high-dimensional vector.

These vectors are what get fed into the transformer blocks, where most of the model’s computation takes place. Llama 3.1 uses different number of transformer blocks depending on the model size. The 8 billion parameter variant of Llama 3.1 has 32 transformer blocks, the 70 billion parameter variant has 80 blocks and the 405 billion parameter variant has 126 blocks. The blocks are arranged sequentially, with each block progressively transforming the representations of the input tokens.

The output from the final transformer block is then normalized and passed through an unembedding layer, which projects each token’s representation back into a vector with one value for every token in the vocabulary. These values, called logits, represent the model’s predictions about which token should come next. Finally, a process called sampling is used to select the next token. Typically, the logits are converted into probabilities using a function called softmax, and a token is selected according to a sampling strategy. This could involve choosing the most probable token or randomly sampling from the probability distribution. The selected token is then appended to the input sequence, and the process repeats to generate text one token at a time.

Tokenization

Tokenization is the first transformation that the LLM’s input undergoes withing our model’s pipeline. Tokenization takes text and converts it into an abstract vocabulary of symbols that represent little chunks of text, parts of words. As the example below shows, the word hippopotamus may be transformed into three tokens “hipp”, “opot” and “amus”, with each token, or symbol being represented by an integer.

Why an integer? Because the set of all possible tokens form our finite vocabulary from which all inputs are recognized and all ouputs and generated. So if our vocabulary contains 128256 tokens, we assing each token an arbitrary number. In our example, “hipp” is token number 44,656, “opot” is token number 89,490 and “amus” is token number 56,455. Since all models are matrix math in the end, we use numbers to represent our vocabulary.

Why tokenize?

As we’ll see in a bit the tokenization process is quite complex. A natural question to ask after seeing all this complexity is why do we tokenize our input this way rather than taking each character in our alphabet as a token or each word in our dictionary as a token? Why this specific way of forming the vocabulary of our model?

The answer lies in an engineering trade-off between compression, vocabulary size, computational efficiency, and the ability to represent arbitrary text.

Consider the situation if we tokenize individual characters. Each character becomes a token, giving us a relatively small vocabulary while retaining the flexibility to represent any text composed of those characters. However, this leads to extremely long input sequences. The transformer must then learn to compose characters into words and higher-level structures before it can reason about their meaning. This becomes particularly expensive because of how transformers work. In a standard transformer, the self-attention mechanism compares tokens against one another, producing an attention matrix whose size scales quadratically with the sequence length $n$.

So doubling the number of tokens approximately quadruples the computation required by the attention operation when processing the full sequence. This is particularly relevant during training and when processing a new prompt. During autoregressive generation, techniques such as KV caching reduce the attention computation required for each additional token to approximately $O(n)$, although the amount of context still matters.

In other words, using individual characters as tokens would significantly increase the amount of computation needed to process the same text.

Alternatively, consider the situation where we tokenize entire words. This produces much shorter sequences, but introduces a different set of problems. First, it reduces the model’s flexibility. If every token corresponds to a complete word, the model can only generate words that exist in its vocabulary. Ask it to invent a new word, and it would have no way of representing that word using the existing tokens. Modern LLMs, in contrast, can invent new words quite easily. Similarly, what happens when we input a word that isn’t in the vocabulary? A traditional word-level tokenizer might have to replace it with an UNKNOWN token, unnecessarily discarding information about the original word. Second, word-level tokenization dramatically increases the required vocabulary size. The vocabulary would need to account for different languages, names, technical terminology, variations of words, and so on.

Since the number of parameters in the model’s embedding layer scales linearly with vocabulary size, a very large vocabulary can significantly increase the model’s parameter count and memory requirements.

For a vocabulary of size $V$ and an embedding dimension $d$, the embedding matrix contains:

$N_{\textnormal{embedding}} = V \times d$

So neither character-level nor word-level tokenization is an ideal solution.

This brings us to subword tokenization, which provides a useful middle ground between these two extremes.

How to Tokenize?

In Llama 3, this is achieved using an algorithm called Byte Pair Encoding (BPE).

The process begins with an initial split of the input text using a regular expression (regex). Rather than immediately tokenizing individual characters, the regex first divides the text into chunks based on simple patterns. For example, it groups consecutive letters, separates punctuation, handles common English contractions, and splits sequences of digits into groups of at most three.

Consider the following input:

The elephant drinks 123456 litres.

The regex would initially divide it into

"The" | " elephant" | " drinks" | " " | "123" | "456" | " litres" | "."

Notice that spaces are sometimes grouped with the following word. This is intentional: common sequences such as ” elephant” can then be represented by a single token, including the preceding space. The regex also prevents the BPE algorithm from merging arbitrary sequences across these initial boundaries.

Once the text has been divided into chunks, the Byte Pair Encoding algorithm takes over.

The tokenizer itself is trained independently of the language model on a representative corpus of text. Initially, its vocabulary contains only the 256 possible byte values. This is sufficient to represent arbitrary UTF-8 text, although doing so one byte at a time would be inefficient. The BPE training algorithm then examines the corpus and identifies the most frequently occurring pair of adjacent symbols. It merges this pair into a new symbol and adds that symbol to the vocabulary.

For example, suppose our training corpus frequently contains the word lower. Initially, it is represented as individual bytes:

l | o | w | e | r

If l and o are a frequently occurring pair, the algorithm might merge them:

lo | w | e | r

It then recalculates the frequencies of adjacent pairs throughout the training corpus and repeats the process. Over successive iterations, we might obtain:

l | o | w | e | r
       ↓
lo | w | e | r
       ↓
lo | w | er
       ↓
low | er
       ↓
lower

Each merge creates a new entry in the vocabulary. The process continues until the desired vocabulary size is reached. In Llama 3, this results in 128,000 ordinary text tokens, with additional reserved token IDs for special purposes such as marking message boundaries. Once training is complete, the vocabulary and the ordering of these merges are fixed. When the tokenizer encounters new text, it applies the same regex-based splitting and then repeatedly merges adjacent byte sequences according to their learned priorities.

The result is that frequently occurring sequences can be represented by a single token, while rarer sequences are broken down into smaller pieces. This achieves a form of lossless compression: common sequences require fewer tokens, but no information about the original text needs to be discarded. The algorithm is described in OpenAI’s tiktoken implementation, which also provides a simple Python implementation for understanding how BPE is trained and applied.

There is one more crucial detail that makes this approach especially powerful: Llama 3 operates on bytes rather than directly on characters. Almost all ordinary text encountered on the internet is represented using character encodings such as UTF-8. In UTF-8, each character is represented by a sequence of one or more bytes. For example:

A  →  41
é  →  C3 A9
🦊 →  F0 9F A6 8A

Since a byte can take only 256 possible values, a tokenizer that includes all 256 individual bytes can represent any valid UTF-8 text, even if it has never encountered a particular word, symbol, or emoji before. A common emoji might be represented by a single token if its byte sequence appears frequently enough in the training data. An uncommon character might require several tokens, but it can still be represented without losing information. This means we don’t need to reserve a separate vocabulary entry for every possible character in every language, or every word anyone might invent. This is why modern LLMs are so flexible. They have a vocabulary carefully chosen so that they can construct any word from these tokens.

There are limitations, of course. Just because a tokenizer can represent a particular word or language doesn’t mean the model understands it. That depends on what the model has learned during training. Moreover, text from languages that are poorly represented in the tokenizer’s training corpus may require many more tokens to express the same amount of information. Nevertheless, byte-level BPE provides an elegant engineering compromise: a manageable vocabulary, relatively compact input sequences, and the flexibility to represent essentially arbitrary text.

Below I’ve embedded a visualization of how this tokenization happens so that you can enter your own text inside and see how Llama 3.1’s tokenizer splits it up.

Tokenizer

The tokenizer this model uses.

VocabularyLlama 3.1cl100k_baseo200k_base

PromptA hippopotamus is drinking water.

SubwordNumberContraction

Loading the Llama 3.1 vocabulary…

Llama 3.1 ranks are Meta’s tokenizer vocabulary, used under the Llama 3 Community License. cl100k_base and o200k_base come from the tiktoken package.

TL; DR

Our tokenizer breaks up input text into chunks based on statistical patterns in a way gives us a good tradeoff between sequence length, computational complexity, vocabulary size, and model flexibility.

String

A hippopotamus is drinking water.

Tokenizer

Token ids

[32, 44656, 89490, 56455, 374, 16558, 3090, 13]

Embeddings

OK great, now our tokenizer has converted our input text into a nice integer representation that the rest of our model can use. But our transformer won’t ingest this number directly. It first gets transformed into something called an embedding vector. In the case of Llama. The embedding vector is of dimension 4096. So each token gets mapped to a 4096 dimensional vector by a giant 4096x128k matrix.

Why Embed?

But why not feed the integers directly to our transformer?

The first thing to understand is that the integers assigned by our tokenizer are essentially arbitrary identifiers. They do not encode any information about the relationships between tokens or their position within the sentence. For example, consider the words eat, drink, and hallucinate. We intuitively understand that all three are verbs, and that eat and drink are more closely related in meaning than either is to hallucinate. However, there is no reason to expect their token IDs to reflect these relationships. The difference between two token IDs tells us nothing about how similar their meanings are.

Feeding these integers directly into a neural network as numerical values would therefore introduce artificial relationships between tokens. A token with ID 1000 would appear numerically closer to token 1001 than to token 5000, even though those numbers have no linguistic significance.

One way to avoid this problem is through one-hot encoding. Suppose our vocabulary contains 10,000 tokens. We could represent each token as a vector containing 10,000 elements, all zeros except for a single 1 at the position corresponding to its token ID. This gives every token a unique representation without introducing any arbitrary numerical ordering. However, it creates two problems.

First, these vectors are extremely large and mostly empty. Llama 3.1, for instance, has a vocabulary of 128,256 tokens. Representing every token with a vector of 128,256 elements and maintaining representations of this size throughout the transformer would be computationally expensive.

Second, one-hot vectors contain also contain no information about the relationships between tokens, and no way for the model itself to learn about these relationships except in the transformer layers. Every pair of distinct one-hot vectors is orthogonal, regardless of whether the corresponding tokens represent closely related words or completely unrelated concepts.

Enter embeddings.

How to Embed

Instead of representing each token in a space with as many dimensions as our vocabulary, we can learn a mapping into a much smaller, dense vector space. In Llama 3.1 8B, for instance, each of the 128,256 tokens is assigned a vector containing just 4,096 numbers. Mathematically, this mapping can be represented by an embedding matrix:

$$E \in \mathbb{R}^{128256 \times 4096}$$

Each row of this matrix corresponds to a token, and each column represents one dimension of the embedding space. If we represent a token as a one-hot vector $\vec{x}$, obtaining its embedding is equivalent to multiplying it by this matrix:

$$\vec{e} = \vec{x}E$$

In practice, we don’t need to perform this multiplication. Since the one-hot vector contains only a single nonzero element, we can simply look up the corresponding row of the embedding matrix. For a sequence of $n$ tokens, this produces a matrix of shape $n \times 4096$, which becomes the initial representation processed by the transformer. This is also the beginning of what is commonly called the residual stream: a sequence of vectors that the subsequent transformer layers progressively update. So the embedding lookup is not a complex operation. The property that is useful to the network is that the values in the embedding matrix are learnable weights of the model.

TL; DR

Embeddings convert our token ids into 4096 dimensional vectors which our transformer blocks can work with.

Token ids

1. 32
2. 44656
3. 89490
4. 56455
5. 374
6. 16558
7. 3090
8. 13

Embedding

Vectors

A 32 ·hipp 44656 opot 89490 amus 56455

0 +0.0103 +0.0117 -0.0036 +0.0042

1 -0.0028 -0.0081 -0.0065 -0.0053

2 -0.0065 +0.0250 +0.0026 +0.0006

3 -0.0078 -0.0053 +0.0125 +0.0107

⋮ ⋮⋮⋮⋮

4,096

·is 374 ·drinking 16558 ·water 3090 . 13

0 -0.0029 +0.0042 -0.0162 -0.0007

1 +0.0013 +0.0038 -0.0090 -0.0007

2 -0.0003 -0.0070 +0.0048 +0.0004

3 +0.0042 -0.0104 -0.0005 +0.0017

⋮ ⋮⋮⋮⋮

4,096

A 32 ·hipp 44656 opot 89490 amus 56455 ·is 374 ·drinking 16558 ·water 3090 . 13

0 +0.0103 +0.0117 -0.0036 +0.0042 -0.0029 +0.0042 -0.0162 -0.0007

1 -0.0028 -0.0081 -0.0065 -0.0053 +0.0013 +0.0038 -0.0090 -0.0007

2 -0.0065 +0.0250 +0.0026 +0.0006 -0.0003 -0.0070 +0.0048 +0.0004

3 -0.0078 -0.0053 +0.0125 +0.0107 +0.0042 -0.0104 -0.0005 +0.0017

⋮ ⋮⋮⋮⋮⋮⋮⋮⋮

4,096

The transformer blocks

The transfomer blocks are where the bulk of the model’s computation happens. The network stacks the same block 32 times. The block can be broken down into two sub-blocks: the Grouped Query Attention and SWIGLU feedforward blocks. Each reads from the residual stream performs computations and adds back to the residual stream.

Architecture / Llama 3.1 8B

What happens inside a single transformer block.

Two residual updates are applied to the residual stream. One from Grouped Query Attention (GQA) and one from a SWIGLU FeedForward block.

Scroll horizontally to inspect the full diagram →

NormalizationRMSNorm(v) = γ ⊙ v / √(mean(v²) + ε)Per token, over 4,096 features; each norm has its own learned γ. ε = 10⁻⁵.

Attention notationQ and K in the attention equation include RoPE. g(i) = ⌈i / 4⌉ selects the shared K/V head. M is 0 for allowed positions and −∞ for future positions; softmax runs over keys.

Weight dimensions

Rows × columns = input × output, for the row-vector convention uW.

W_q, W_o

4,096 × 4,096

W_k, W_v

4,096 × 1,024

W₁, W₃

4,096 × 14,336

W₂

14,336 × 4,096

γ₁, γ₂

4,096 each

Each W packs all its heads into one matrix. W_q contains 32 column slices of 4,096 × 128; W_k and W_v each contain 8.

What’s the residual stream?

One of the major challenges in training deep neural networks was that simply adding more layers could increase training error rather than improve performance as the gradients degraded in the deeper layers. He et al. (2016) addressed this with residual connections, which allow each layer to learn a modification to its input rather than an entirely new representation. Mathematically, instead of learning $H(x)$ directly, a layer learns a residual function $F(x)$, producing $H(x)=x+F(x)$. These identity shortcuts provide a direct path for information and gradients to propagate through the network, making very deep architectures easier to optimize (He et al., 2016). Subsequent work showed that residual connections also mitigate the shattered gradients problem, where gradients become increasingly uncorrelated as network depth increases (Balduzzi et al., 2017). In transformers, this architecture gives rise to what is commonly called the residual stream: a shared representation that successive attention and feedforward layers read from and update additively, rather than completely replacing at every layer (Elhage et al., 2021). This also provides a useful interpretation of the residual stream as a form of computational workspace, where information can be accumulated, preserved, and refined across layers.

What’s a transformer?

So far, each token has been assigned a vector from our embedding table. But this lookup has an important limitation: the same token always gets the same initial vector, regardless of the sentence it appears in. Consider the word “bank” in “I deposited money at the bank” and “I sat on the river bank”. Assuming it is represented by the same token in both sentences, its embedding is identical. The surrounding words are what tell us which meaning is relevant.

A transformer is a neural network architecture that turns these initial vectors into representations that incorporate context. Its basic block combines two operations: self-attention, which lets token positions exchange information, and a feedforward network, also called a multilayer perceptron or MLP, which processes the resulting vector at each position.³

In our example, attention gives the representation at “bank” a way to draw on information from “money” or “river”. The MLP then applies a learned, nonlinear transformation to the representation at that position. We’ll look at both operations in detail below. For now, the useful distinction is that attention moves information between positions, while the MLP computes on the information available at each position.

Both operations write their results back into the residual stream we just discussed. Repeating this across many blocks allows later layers to work with features constructed by earlier ones. The blocks share the same architecture, but each has its own learned weights.

Llama 3.1 is a decoder-only transformer. It has no separate encoder that first reads the prompt. Instead, the prompt and generated text pass through the same stack, with each position allowed to attend only to itself and earlier positions. This restriction is what lets us use the resulting representations to predict the next token without giving the model access to the answer.¹

Why Transformers?

Why use this particular architecture to process a sequence of tokens?

Before transformers, a common approach was a recurrent neural network, which processes a sequence by repeatedly updating a hidden state. Each update depends on the previous one, so the computation must move through the sequence in order. Information from an early token also has to survive those successive updates to influence a much later one.

Self-attention gives a later position a direct route to an earlier position within the available context. It also allows the positions within a layer to be processed in parallel when the input sequence is already known, as it is during training. Much of this work can be expressed as large matrix multiplications, which GPUs handle efficiently. These were central motivations for the original transformer architecture.³

There is an important distinction here. During training, we already have the text and can calculate next-token predictions at many positions simultaneously, masking out future tokens at each position. During generation, we don’t yet know the next token. We still have to select it before we can generate the one after it.

There is also a cost to letting positions read one another. As we saw in the tokenization section, full attention requires a number of pairwise comparisons that grows quadratically with sequence length. So transformers give us flexible access to context and efficient parallel computation, but long sequences remain expensive.

What’s RMSNorm?

Before either attention or the MLP reads the residual stream, Llama applies an operation called Root Mean Square Normalization, or RMSNorm. The idea is to control the scale of the vector that the next operation receives.

To see what this means, imagine two vectors: $[1, 2, 3]$ and $[10, 20, 30]$. They point in the same direction, but the second is ten times larger. If we feed them directly into a matrix multiplication, the second produces an output ten times larger too. As layers repeatedly add updates to the residual stream, its magnitude can change. Normalization makes the computations that read it less sensitive to that overall scale.

For a token’s vector $\vec{x}$ containing $d$ numbers, the root mean square is exactly what the name suggests: square the numbers, take their mean, and then take the square root. RMSNorm divides by this scale and applies a learned multiplier to each coordinate:⁴

$$\operatorname{RMSNorm}(\vec{x})_i = g_i\frac{x_i}{\sqrt{\frac{1}{d}\sum_{j=1}^{d}x_j^2 + \epsilon}}$$

Here $d = 4096$ for Llama 3.1 8B, $g_i$ is a learned weight for coordinate $i$, and $\epsilon$ is a small positive constant that keeps the denominator away from zero. Each token is normalized independently across its own 4,096 coordinates; we aren’t averaging across tokens.

For our example, the RMS of $[1, 2, 3]$ is approximately $2.16$, while that of $[10, 20, 30]$ is approximately $21.6$. Ignoring $\epsilon$, dividing by these values gives the same vector in both cases: approximately $[0.46, 0.93, 1.39]$. The learned weights can then adjust the scale of individual coordinates.

Unlike LayerNorm, RMSNorm does not subtract the vector’s mean before rescaling it.⁴ In Llama, the normalized vector is passed into attention or the MLP, while the residual connection carries the original vector forward. The resulting update is added to that original stream. There is also one final RMSNorm before the output projection.¹

What’s the SWIGLU block?

After attention has gathered information from other positions, what do we do with it? This is where the feedforward part of the transformer comes in. It takes the representation at each position and computes an update to that position’s residual stream. In Llama 3.1, this network uses a structure called SwiGLU, short for Swish Gated Linear Unit.⁵

The name sounds more complicated than the operation. We take the normalized vector, send it through two separate learned projections, and multiply their outputs together after applying a nonlinear function to one of them. A third projection brings the result back to the width of the residual stream.

Let’s follow one token through this process. Call its vector after attention $\vec{h}$, and its normalized vector $\vec{z} = \operatorname{RMSNorm}(\vec{h})$. In Llama 3.1 8B, $\vec{z}$ contains 4,096 numbers. Both of the first two projections expand it into a vector containing 14,336 numbers:⁶

$$\vec{a} = \vec{z}W_1, \qquad \vec{b} = \vec{z}W_3$$

Here we’re treating our vectors as rows, so $W_1$ and $W_3$ each have dimensions $4096 \times 14336$. Although both projections read the same input, they have different learned weights and can extract different combinations of the information it contains.

The first branch passes through an activation function called SiLU, also known as Swish with its parameter fixed at one:

$$\operatorname{SiLU}(a) = a\sigma(a), \qquad \sigma(a) = \frac{1}{1 + e^{-a}}$$

The sigmoid $\sigma(a)$ lies between zero and one. For a large positive input, SiLU therefore returns approximately the input itself. For a very negative input, it returns a small negative value close to zero. Between these extremes, it changes smoothly.

We then multiply this activated vector by the second branch, coordinate by coordinate:

$$\vec{u} = \operatorname{SiLU}(\vec{a}) \odot \vec{b}$$

The symbol $\odot$ means elementwise multiplication, rather than a matrix multiplication or a dot product. If two illustrative branch outputs were $[0.1, 2, 0]$ and $[5, 3, 7]$, their product would be $[0.5, 6, 0]$. Each number in the first branch scales the corresponding number in the second. This is the gate in SwiGLU. It is a continuous multiplier, not a binary switch; the SiLU output can also be negative or greater than one.

Finally, a matrix $W_2$ of dimensions $14336 \times 4096$ projects this wider vector back into 4,096 numbers. Putting the operations together gives the feedforward update:⁷

$$\vec{m} = \left[\operatorname{SiLU}(\vec{z}W_1) \odot (\vec{z}W_3)\right]W_2$$

Llama uses no bias terms in these three projections. The update is then added to the residual stream, giving $\vec{y} = \vec{h} + \vec{m}$. The temporary expansion to 14,336 dimensions happens inside the feedforward network; the residual stream still has width 4,096.

The same operation, with the same weights within a layer, is applied independently at every token position. SwiGLU doesn’t read neighbouring positions directly. Any information about those positions must already be present in the vector it receives, through attention and earlier blocks.

Why SWIGLU rather than MLP?

First, a small clarification: SwiGLU is the particular form of MLP used in Llama. The comparison is with a conventional MLP that expands the input through one projection, applies an activation such as ReLU or GELU, and projects the result back down. Ignoring bias terms, that looks like:

$$\operatorname{MLP}(\vec{z}) = \phi(\vec{z}W_{\text{up}})W_{\text{down}}$$

Here $\phi$ is the activation function. Why do we need a nonlinearity at all? Without it, the two matrix multiplications could be combined into one matrix. Expanding the vector and then shrinking it would still produce only a linear transformation. The activation allows the network’s response to depend on the input in more complex ways.

SwiGLU adds another way for that response to depend on the input: multiplication between two learned branches. One branch supplies values, while the other modulates how much each value contributes. As an intuition, this gives the network a way to make a feature’s contribution depend on another combination of features in the same representation. Both branches are learned together; we haven’t assigned one to detect “river” and the other to detect “bank”. The model has to discover useful computations during training.

That intuition doesn’t establish that SwiGLU will work better. The evidence is experimental. Shazeer (2020) compared feedforward variants in a T5-style transformer with matched parameter and computation budgets. SwiGLU and GEGLU achieved better held-out perplexity than the ReLU and GELU baselines in those experiments.⁵ This supports using gated variants, but doesn’t mean SwiGLU must outperform every alternative in every training setup.

There is also an engineering trade-off. At the same intermediate width, SwiGLU has three large weight matrices instead of two, so it requires more parameters and matrix-multiplication work. For a residual width $d$ and intermediate width $f$, the bias-free parameter counts are $2df$ for the conventional MLP and $3df$ for SwiGLU. Choosing a narrower intermediate layer can balance this extra cost.

For our walkthrough, the role of this part of the block remains the same: attention brings information to a position, and the feedforward network computes with it. SwiGLU gives that computation an input-dependent gate before writing the resulting update back into the residual stream.

TL; DR

The network stacks the same block 32 times (for the 8b variant), the larger parameter variants use more blocks and more attention heads. Each block reads from the residual stream performs some computations and adds back to the residual stream.

Architecture / Llama 3.1 8B

One block, in outline

N × 4,096 in and out · repeated 32 times

RoPE

TL;DR

RoPE gives attention a way to use position by rotating the queries and keys before comparing them. Take one head’s $N \times 128$ matrix: each row is a token, and its 128 coordinates form 64 pairs. Rotate each pair by $m\omega_i$, where $m$ is the token’s position and $\omega_i$ is that pair’s rotation speed. Different pairs turn at different speeds; each pair keeps its length, and the output is still $N \times 128$.

The same position-and-pair angle schedule applies to every query and key head, so one head is enough to illustrate the operation. Their learned vectors differ, even when their rotation angles match. In Llama 3.1 8B, this happens independently across 32 query heads and 8 key heads; values aren’t rotated.⁸

RoPE application for One query or key head

RoPE Rotates Pairs of elements in the Query and Key vectors.

1 / Input

2 / Rotate each pair

3 / Output

When attention takes a rotated query’s dot product with a rotated key, their shared absolute rotation cancels, leaving a positional contribution that depends on their relative offset. The score still depends on their content. Here’s why this works.

To understand why this RoPE block sits between the Query / Key transformations and the attention block, we first need to look at what the attention computation knows about the order of our tokens.

Where does position enter the computation?

Consider two sentences: “The dog chased the cat” and “The cat chased the dog”. They contain the same words, but mean quite different things. The order matters. However, looking up a token in the embedding table gives us the same vector wherever it appears in the sequence. The embedding for ” dog” doesn’t come with a label saying “I’m the second token”.

Let’s examine one attention head before adding any positional information. If $\vec{x}_m$ and $\vec{x}_n$ are the normalized vectors at positions $m$ and $n$, its query, key, and attention score are:

$$\vec{q}_m = \vec{x}_m W_Q, \qquad \vec{k}_n = \vec{x}_n W_K, \qquad s_{mn} = \frac{\vec{q}_m \cdot \vec{k}_n}{\sqrt{128}}$$

The subscripts tell us where the vectors came from. But there is no numerical input here representing $m$, $n$, or the distance between them. If we hold the vectors fixed and move them to different positions, their dot product stays the same.

Without positional information or a mask, rearranging the input tokens would just rearrange the output vectors in the same way. Each query would still see the same collection of keys and values, and compute the same weighted sum. The attention operation would have no way to distinguish those two sentences beyond moving the representations around.

There is a qualification here for Llama: the causal mask already introduces some information about order. A token can read earlier tokens, but not later ones, so moving a token changes which context it can access. It would be too strong to say that the entire network knows nothing about order. What the dot product above lacks is an explicit way to use the distance between two positions.

How did the original transformer do this?

The original transformer by Vaswani et al. added a positional vector to each token’s embedding before it entered the transformer stack:³

$$\vec{h}_m = \operatorname{Embedding}(t_m) + \operatorname{PE}(m)$$

Here $t_m$ is the token at position $m$. The positional vector has the same width as the token embedding, so we can add them coordinate by coordinate. Its entries are sine and cosine functions of the position, with different frequencies across coordinate pairs:

$$\operatorname{PE}(m)_{2i} = \sin\left(\frac{m}{10000^{2i/d}}\right), \qquad \operatorname{PE}(m)_{2i+1} = \cos\left(\frac{m}{10000^{2i/d}}\right)$$

Here $d$ is the embedding width and $i$ indexes a pair of coordinates. These functions supply a pattern of numbers that changes with position. So the same token at two different positions enters the network with two different vectors. These positional encodings were calculated from a fixed formula, rather than learned as a separate table.

Llama 3.1 instead uses Rotary Position Embedding, or RoPE. It applies position information to the queries and keys inside each attention layer, after their learned projections.⁸

What exactly are we rotating?

RoPE is a bit tough to get your head around if you imagine rotating a 128-dimensional vector all at once. It becomes much simpler if we take two numbers at a time.

Ignoring the batch dimension, Llama 3.1 8B produces queries with shape $N \times 32 \times 128$ and keys with shape $N \times 8 \times 128$. Each head’s vector contains 128 numbers. In Meta’s implementation, we group these into 64 adjacent pairs: $(x_0, x_1)$, $(x_2, x_3)$, and so on.⁸ Each pair can be treated as a little two-dimensional vector, an arrow on a plane.

To rotate one of these arrows by an angle $\phi$, we use ordinary two-dimensional rotation:

$$\begin{aligned} x'_{2i} &= x_{2i}\cos\phi - x_{2i+1}\sin\phi \\ x'_{2i+1} &= x_{2i}\sin\phi + x_{2i+1}\cos\phi \end{aligned}$$

For example, rotating $[1, 0]$ by $90^\circ$ gives $[0, 1]$. The arrow changes direction, but its length stays the same. RoPE performs this operation on every coordinate pair, with an angle determined by the token’s position and the pair’s rotation speed.

If pair $i$ rotates by $\omega_i$ radians per token position, the angle at position $m$ is:

$$\phi_{m,i} = m\omega_i$$

We count positions starting at zero. At position zero, no rotation is applied. At position one, pair $i$ rotates by $\omega_i$; at position two, by $2\omega_i$; and so on. A query and key at the same position use the same angles for corresponding pairs, even though their actual numbers come from different learned projections.

Why does rotation tell us about relative position?

This is the clever part. Attention compares a query and a key using their dot product. For two arrows, the dot product is the product of their lengths multiplied by the cosine of the angle between them. Rotating both arrows by the same angle leaves that angle between them unchanged.

Suppose, just for illustration, that one coordinate pair in both the query and key starts as $[1, 0]$, and rotates by $30^\circ$ per position. Put the query at position five and the key at position two. Their rotated angles are $150^\circ$ and $60^\circ$, so the angle between them is $90^\circ$ and their dot product is zero.

Now move the query to position eight and the key to position five, holding the original vectors fixed. Their angles become $240^\circ$ and $150^\circ$. The difference is still $90^\circ$. The dot product hasn’t changed, because the two tokens are still three positions apart. Move only the key, however, and the dot product can change.

We can express this for the full head too. Let $R_m$ denote all 64 pairwise rotations for position $m$. For this equation, write the vectors as columns so the rotation matrices multiply on the left:

$$\widetilde{\vec{q}}_m = R_m\vec{q}_m, \qquad \widetilde{\vec{k}}_n = R_n\vec{k}_n$$

Then:

$$\widetilde{\vec{q}}_m^{\,T}\widetilde{\vec{k}}_n = \vec{q}_m^{\,T}R_m^T R_n\vec{k}_n = \vec{q}_m^{\,T}R_{n-m}\vec{k}_n$$

The transpose $R_m^T$ undoes the query’s rotation. Combining that with the key’s rotation leaves a rotation corresponding to $n-m$, the relative offset.⁹ RoPE uses each token’s absolute position to rotate its vector, but the positional contribution to their comparison depends on the difference between positions.

The score still depends on the content of the query and key. We haven’t replaced attention with a distance calculation. We’ve given the learned projections a way to make their comparisons depend on both content and relative position. And the shared-shift property above holds when we keep those original vectors fixed; moving actual text can also change its context and therefore its vectors.

Why use different rotation speeds?

A single rotating arrow eventually comes back around. If every pair rotated at the same speed, we’d keep repeating the same positional pattern. Using many speeds gives the model a richer combination of patterns across short and long distances.

The base RoPE schedule for a head of width $d_h = 128$ is:

$$\omega_i = \theta^{-2i/d_h}, \qquad i = 0, 1, \ldots, 63$$

Llama uses $\theta = 500000$.¹ The first pair rotates by one radian per position, and successive pairs rotate more slowly. You can think of them as a collection of clock hands turning at different speeds: some change quickly between neighbouring tokens, while others change appreciably only over much longer stretches.

There is one Llama 3.1-specific detail. Its long-context version adjusts this base schedule: the fastest rotations stay unchanged, the slowest are slowed by a factor of eight, and an intermediate range is blended between the two settings.⁸ So the formula above is the starting schedule, not the complete set of rotation speeds used by Llama 3.1. The adjusted schedule and additional long-context training support its 131,072-token context window.¹

Being able to calculate a rotation for an arbitrarily large position doesn’t by itself mean the model will work well at that position. It still needs training that teaches it to use those patterns. Nor does RoPE guarantee that attention decreases steadily with distance: the score also depends on the learned query and key vectors.

For our walkthrough, RoPE leaves the query and key tensor shapes unchanged. The rotated queries and keys go into the attention score calculation, followed by scaling, masking, and softmax. The values aren’t rotated by RoPE in Llama; they supply the information that attention gathers once it has decided how much weight to give each position.⁸

Unembedding

After all the transformer blocks, what we have at the output is a matrix of dimensions $N \times 4096$, where $N$ is the number of tokens in our input sequence. Each row is the vector at one token position. This is the residual stream after all the successive transformations performed by the transformer blocks.

In Llama 3.1 8B, the residual stream first passes through a final RMSNorm layer, which rescales each vector and applies learned weights to its coordinates. It then passes through a learned linear projection called the unembedding, which maps each 4,096-dimensional vector to 128,256 scores, one for every token in Llama 3.1’s vocabulary.⁶

$$Z = \operatorname{RMSNorm}(X)W_{\text{out}}$$

where $X$ has dimensions $N \times 4096$, and $W_{\text{out}}$ is a learned weight matrix of dimensions $4096 \times 128256$.

Projecting all positions gives an $N \times 128256$ matrix of logits. Each row contains a score for every token in the vocabulary, representing the model’s prediction of which token should follow the token at that position. To generate the next token after our input sequence, we use only the last row. We can also compute the unembedding for just the last-position vector when we don’t need predictions at the earlier positions.

Logits are unnormalized scores, which can take any real value. A higher score means that the model assigns a higher relative probability to that token. We need to convert these scores into an actual probability distribution.

Output / Llama 3.1 8B

Llama 3.1 Unembedding

One more RMSNorm layer, then a bias-free projection · N = 8 for this sentence

Normalization û = γ ⊙ X / √(mean(x²) + ε), one row at a time One γ for this final norm: 4,096 learned scales, one per feature. Each of the N rows is scaled on its own. ε = 10⁻⁵. The shape stays N × 4,096.

Projection

Z = û W_out

Entry i, v is row i of û dotted with column v of W_out. The inner length of that dot product is 4,096.

W_out

4,096 × 128,256

Parameters

525,336,576

Bias

none

γ

4,096

Next-token row row N of Z · length 128,256 Every row of Z is a logit vector for the token that follows that position. Generation reads the last row. Here that row is the period after “water”.

Temperature

So after the RMSNorm, a parameter called temperature is used to scale (divide) the logits. This operation causes the logits to either “spread out” the distribution more or make it more narrow and spiky. Intuitively, a low temperature between zero and one makes the probability distribution narrower and more concentrated around the highest-scoring tokens. Such a model may produce more consistent and “boring” outputs.

A temperature greater than one has the opposite effect. It makes the distribution flatter, increasing the probability of selecting less likely tokens. This introduces more randomness into the generated text, potentially producing more varied or “crazier” outputs, but also increasing the likelihood of nonsensical predictions.

Sampling

Softmax

After scaling by the temperature, we finally calculate the probabilities of the tokens. This is done by applying a function called the softmax over our entire scores vector. It produces essentially a probability distribution over the tokens.

The softmax function is defined as:

$$P(i) = \frac{e^{z_i/T}}{\sum_{j=1}^{V}e^{z_j/T}}$$

where $V$ is our vocabulary size, $z_i$ is the logit corresponding to token $i$, and $T$ is the temperature.

The exponential ensures that all values are positive, while dividing by the sum ensures that the resulting probabilities add up to one. After softmax, we have a VOCAB_SIZE × 1 vector containing the probability assigned to every token in the vocabulary.

Top-k and top-p sampling

We’re still not done though. We could just select the highest probability token and predict that at this point but this generally leads to LLMs generating the same thing all the time given the same output. But we see that this generally doesn’t happen. If you give an LLM the same input twice, it tends to produce different answers each time. This variation is achieved using random sampling.

Top-k sampling takes the $k$ tokens with the highest probabilities, discards all other tokens, and then randomly selects a token from this reduced set according to their relative probabilities. For example, with top_k = 5, only the five most likely tokens are considered, regardless of how much probability mass they collectively represent. The remaining probabilities are renormalized so that they sum to one before sampling.

One limitation of this approach is that the number of candidate tokens is fixed. If the model is extremely confident about the next token, it still considers five candidates. Conversely, if the probability distribution is relatively flat, it might discard many reasonable alternatives.

Top-p sampling, also known as nucleus sampling, takes a slightly different approach. Instead of choosing a fixed number of tokens, it sorts the tokens by their probabilities and selects the smallest set whose cumulative probability is at least some threshold $p$.

For example, with top_p = 0.9, we retain the most likely tokens until their combined probability reaches or exceeds 90%. The remaining tokens are discarded, and the retained probabilities are renormalized before sampling.

The advantage of top-p sampling is that the number of candidate tokens adapts to the model’s confidence. When the model is confident, the candidate set may contain only a few tokens. When it is uncertain, the set can be much larger. Top-k and top-p can also be used together. In this case, both restrictions are applied, and the next token is sampled from the remaining candidates.

For Llama 3.1 8B, Meta’s official reference implementation uses temperature scaling followed by top-p sampling. For continuous text generation, the predicted token is appended to the prompt, and the forward pass runs again.

The figure is this sentence’s forward pass through Llama 3.1 8B. Temperature is $1$, and neither top-$k$ nor top-$p$ is applied, so $q = p$ and the softmax runs over all 128,256 logits. Only the eight largest are drawn. The marked bar is one uniform draw from that distribution.

Sampling / Llama 3.1 8B

Predicting the next token.

The model produces one logit per token. Applying Softmax turns the scores into probabilities, and a sampling process selects the next token.

Putting it all together

So to summarize everything again, let’s take another look at the high level pipeline.

1. 01

   Tokenization

   The prompt is split into subword token ids from a vocabulary of 128,256.
2. 02

   Embedding

   Each id is looked up as a vector of 4,096 numbers, with position from RoPE.
3. 03

   Transformer blocks

   The sequence passes through 32 blocks.
4. 04

   Unembedding

   The last position becomes a score for every token in the vocabulary.
5. 05

   Next token

   A sampling rule picks one token from that distribution.

Our text is tokenized, embedded, processed with multiple transfomer layers where most of the magic happens, then unembedded and a random sampling process used to produce a next token. After this, this token is fed back into the model and the process is repeated.

When I first saw this whole process I was thinking, “That’s it?”

Sure, the process has a lot of steps and is tedious to work through. But once you know the steps, it seems rather simple. Most of the operations boil down to matrix multiplication and addition. The proposition that if you scale this up with a big enough model and train it with enough well curated data, you get intelligence felt deeply dissatisfying.

And more importantly, there’s a crucial question that I had when I started out on this article that was still un-answered in my mind. How does this result in models that can solve IMO problems and tackle Millennium Prize proofs at the frontier of mathematics?

We know that attention allows tokens to exchange information, that the feedforward layers perform nonlinear computations, and that the residual stream carries information through successive layers. But what exactly is being represented in those 4,096-dimensional vectors? What computations are the attention heads performing? How does the model learn to represent concepts, retrieve facts, or carry out something resembling reasoning?

These questions (and other similarly hard questions) are what the field of mechanistic interpretability is trying to answer by probing the model with mathematical tools. The idea is to identify the features, circuits, and algorithms that emerge within these networks as they learn, and understand how these mechanisms produce intelligent behaviour. I believe this field will become increasingly important in the near future, alongside AI alignment research, as we try to understand, predict, and ultimately ensure that we can coexist peacefully with the alien minds we’ve created.

Footnotes

1. Llama Team at Meta. “The Llama 3 Herd of Models.” arXiv:2407.21783 (2024). Architecture table for the layer count, the hidden dimensions, grouped-query attention, RoPE $\theta$, and SwiGLU. ↩ ↩² ↩³ ↩⁴ ↩⁵

2. Polo Club, Georgia Tech. Transformer Explainer. The interaction and the narrative structure of a GPT-2 walkthrough; the architecture on this page is Llama 3.1. ↩

3. Vaswani, Ashish, et al. “Attention Is All You Need.” NeurIPS 2017. Introduces the transformer architecture and explains the parallelism and shorter paths between positions provided by self-attention. ↩ ↩² ↩³

4. Zhang, Biao, and Rico Sennrich. “Root Mean Square Layer Normalization.” NeurIPS 2019. Introduces RMSNorm, which rescales activations without the mean subtraction used by LayerNorm. ↩ ↩²

5. Shazeer, Noam. “GLU Variants Improve Transformer.” arXiv:2002.05202 (2020). Defines SwiGLU and compares gated feedforward variants with ReLU and GELU baselines at matched parameter and computation budgets. ↩ ↩²

6. Hugging Face. Llama 3.1 8B model configuration. vocab_size 128256, num_key_value_heads 8, and the rest of the 8B config. ↩ ↩²

7. Meta. Llama reference implementation, FeedForward. Implements the bias-free update as w2(silu(w1(x)) * w3(x)), followed by a residual addition in TransformerBlock. ↩

8. Meta. Llama 3.1 reference implementation. precompute_freqs_cis, apply_scaling, and apply_rotary_emb define the frequency schedule, long-context adjustment, and pairwise query/key rotations. ↩ ↩² ↩³ ↩⁴ ↩⁵

9. Su, Jianlin, et al. “RoFormer: Enhanced Transformer with Rotary Position Embedding.” Introduces RoPE and derives the relative-position dependence of the query–key dot product. ↩