How LLMs Work: Next-Token Prediction in Llama 3.1

Introduction

Over the last six months or so, I’ve been reading a lot about large language and vision models in an effort to understand their inner workings better. Like a significant fraction of AI researchers in the LLM era, the idea that scaling up compute and trying to predict the next token could eventually produce systems capable of chipping away at math problems on the frontier of human research felt both astounding and utterly counter-intuitive to me.

As a first step towards understanding LLMs better, I wanted to understand exactly how an LLM works from end to end. Llama 3.1 seemed like a good place to start. It’s a text only model, so I can focus on one thing at a time, and it’s one of the last generations of open-weight models from before everyone started pivoting to MoEs and fancy state space models like Mamba. Technically, it’s quite simple — using a standard decoder only architecture with minor tweaks and it has an 8b variant that runs on my laptop (24 GB Macbook Air) quite well.

So how does Llama 3.1 actually generate text?

At the highest level, the popular description that they just predict the next token is reasonably accurate. You feed some text into the model, and at the output it produces a set of probabilities for which bit of text comes next. Then you tack this next bit on at the end of the text you initially fed in and repeat the process again until the model says it’s done. 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’ll try to give you my best understanding of exactly what happens between the text input and the decision of what token comes next. I’ll be using the Llama 3.1 8B variant which has 32 layers of transformer blocks, a residual stream of width 4,096, and a vocabulary of 128,256 entries (128,000 ordinary text tokens and 256 special or reserved entries).¹²³

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 just a little, the model’s processing pipeline can be split into five stages. First, the input text is converted into something called tokens. These tokens then pass through an embedding layer, which converts each token into a high-dimensional vector.

These vectors are processed by several 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 within our model’s pipeline. It takes text that’s fed into the model and converts it into a sequence of symbols from a fixed vocabulary. These tokens represent chunks of text, which may be whole words or parts of words. In our example sentence, “A hippopotamus is drinking water.”, the text hippopotamus, including its leading space, becomes three tokens: hipp, opot, and amus. Without that space, hippopotamus splits into four tokens: hip, pop, ot, and amus. Each of these tokens are represented by an integer id.

Why an integer? Because each token is part of a finite vocabulary, we can identify it by a number. Llama 3.1 has 128,256 vocabulary entries, each with its own ID. In our example, hipp (with the leading space) is token number 44,656, opot is token number 89,490, and amus is token number 56,455. These IDs tell the model which tokens it has received; the numbers themselves don’t encode their meanings.

Why tokenize?

As we’ll see in a bit the tokenization process is quite complex. A natural question to ask is why tokenize our input this specific way? Why not use character in our alphabet as tokens or dictionary words as tokens? Why do we use specific, complex 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 use individual characters such as ‘c’ or ‘a’ as tokens. Our vocabulary becomes relatively small while retaining the flexibility to represent any text composed of those characters. However, this leads to extremely long input sequences. The transformer must also 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 (explained in Hugging Face’s caching guide) 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 use entire words as tokens. 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.1, tokenization uses an algorithm called Byte Pair Encoding (BPE).

First, a regular expression splits up our input 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" | "."

You might notice that spaces are sometimes grouped with the following word. 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 begins its work.

The tokenizer’s vocabulary and merge rules are learned separately from the language model’s weights. To understand the idea, let’s look at a simplified example of training a byte-level BPE tokenizer from scratch. 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 (hence, byte pair encoding). 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. In this simplified training procedure, the process continues until the desired vocabulary size is reached. Meta didn’t build Llama 3’s entire vocabulary from scratch this way though: its 128,000 ordinary text tokens combine 100,000 tokens from the tiktoken library with 28,000 additional tokens to better support non-English languages.¹ Another 256 entries are set aside for special tokens and reserved IDs (like stopping generation for instance), bringing the total to 128,256.² Once the tokenizer is built, its vocabulary and merge priorities are fixed. When it encounters new text, it applies the regex-based splitting and then repeatedly merges adjacent byte sequences according to those priorities.

The result is that frequently occurring sequences can be represented by a single token, while rarer sequences are broken down into smaller pieces. You can think of this as a form of lossless compression similar to a Huffman Code. 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.

There are limitations, of course. For instance, 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 than overrepresented languages. Nevertheless, byte-level BPE provides an elegant engineering compromise.

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, our tokenizer has converted our input text into a nice integer representation that the rest of our model can use. But our transformer doesn’t ingest this number directly. It first gets mapped into something called an embedding vector. In the case of Llama. The embedding vector is of dimension 4096.

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. We could represent each token as a vector containing VOCAB_SIZE 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.

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

Enter embeddings.

What is Embedding?

king − man + woman ≈ queen

An embedding represents a word or token as a learned, dense vector whose relationships to other vectors can capture patterns in language. Over two decades ago, Bengio and colleagues showed in 2003 how learning word vectors alongside a next-word predictor lets a model generalize across similar words, rather than treating every word sequence as unrelated.⁴ In 2013, Mikolov and colleagues made learning these representations much cheaper with with their landmark word2vec paper, which predicts a word from its neighbours or neighbours from a word. The resulting vectors captured useful semantic and grammatical relationships.⁵ In the years since, neural network based language models largely converged on these vector representations as a useful or even standard part of language modeling pipeleines. Modern transformers, including Llama, continue that approach with token embeddings learned jointly with the rest of the model.⁶¹

How to Embed

So embeddings map from our vocabulary to a small, dense vector space. In Llama 3.1 8B, 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 actually perform this multiplication since most of these results would be discarded. 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 called the residual stream.

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. This is an optimization problem, distinct from vanishing or exploding gradients.⁷

Residual connections address this by allowing a block of layers to learn a modification to its input rather than an entirely new representation. Mathematically, instead of learning $H(x)$ directly, the block learns a residual function $F(x)$, producing $H(x)=x+F(x)$. The identity shortcut carries the input directly to the addition, and this formulation makes very deep architectures easier to optimize.⁷

Subsequent work showed that residual connections also mitigate the shattered gradients problem: in the networks studied at initialization, gradients at different inputs become increasingly uncorrelated as depth increases, while skip connections help preserve those correlations.⁸

In transformers the residual stream behaves as a shared representation that successive attention and feedforward layers read from and update additively, rather than completely replacing at every layer.⁹ 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 “attend to” 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 thing to keep in mind 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.

Why Transformers?

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

Before transformers, it was common to use recurrent neural networks to processes sequences. Each update depends on the previous one, so the computation must move through the sequence in order. So we encounter our old enemy, the Vanishing Gradients Problem. Information from an early token 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 mdoern GPUs handle very efficiently. These were some of the central motivations for the original transformer architecture.⁶.

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 very long sequences remain expensive.

What’s Attention?

The attention block is the heart of the transformer. Attention is the mechanism that allows each token position to gather a mixture of information from the positions it can read. Let’s stay with our example, “A hippopotamus is drinking water.” Its eight tokens give us eight rows in the residual stream. Each row contains 4,096 numbers, so the input to this part of a block is a matrix $X$ with shape $N \times 4096$, where $N = 8$.

First, Llama normalizes each row with RMSNorm, giving a matrix $Z$ of the same shape. Three different learned projections then produce queries, keys, and values. For one query head and its associated key/value head:

$$Q = \operatorname{RoPE}(ZW_Q), \qquad K = \operatorname{RoPE}(ZW_K), \qquad V = ZW_V$$

Each weight matrix here has shape $4096 \times 128$, so Q, K, and V each have shape $N \times 128$. Their learned weights extract different features from the same residual vector. Intuitively, a query describes what a position is looking for; a key gives it something to match against; a value carries the information it can contribute. RoPE adds position information to Q and K. We’ll unpack RMSNorm and RoPE below.

This is mathematically represented first by performing a matrix multiplication between the Query and Key matrices:

$$S = \frac{QK^{\mathsf T}}{\sqrt{128}}, \qquad (N \times 128)(128 \times N) \longrightarrow N \times N$$

Transposing K makes its token rows into columns. Each entry $S_{ij}$ is the dot product of query row $i$ and key row $j$, divided by $\sqrt{128}$. We have turned a table of features per token into a table of scores between token positions.⁶

The causal mask sets future-position scores to $-\infty$. Softmax is applied across each row, producing a weight matrix A of the same $N \times N$ shape. Every row sums to one, so positions after that row’s query position get zero weight. For example, the row for ” drinking” can read itself and the five tokens before it, but not ” water” or the period.

Attention / Llama 3.1 8B / one query head and its K/V head

Attention maps the residual stream to a mixture of context.

“A hippopotamus is drinking water.” gives us N = 8 token rows. Follow those same rows through the calculation. Shapes below are rows × columns.

Scroll each diagram sideways to inspect its matrices →

01 / Project the same input three ways

The projections change the features within each row. No token rows have been mixed yet. Q and K go into the comparison below; V is saved for step 4.

02 / Compare every query row with every key row

The 128 feature coordinates are summed over. What remains is one score for every pair of token positions: 8 × 8 = 64 scores.

03 / Turn each score row into a set of mixing weights

Each row of A is a recipe for mixing value vectors. The diagonal is allowed: a position can read its own value. Shading marks the permitted cells, not measured weight magnitudes.

04 / Use those weights to mix the value rows

This multiplication mixes information across token positions. Each output row still belongs to its original query position, but now contains a weighted mixture of the value vectors it could read.

Back to the full attention block

32 head outputs N × 128 each → concatenate N × 4,096 → output projection 4,096 × 4,096 → add to X N × 4,096

The second multiplication uses those weights to combine the values:

$$H = AV, \qquad (N \times N)(N \times 128) \longrightarrow N \times 128$$

Each row of H is a weighted sum of the rows of V. The weights in $A_{ij}$ determine how much value row $j$ contributes to output row $i$, with the same weight applied to all 128 coordinates of that value vector. This is where information moves between token positions.

So the sequence keeps its $N$ rows throughout. Q and K determine the mixing weights; V supplies the information being mixed. One head returns 128 numbers per position. Combining all 32 head outputs gives $N \times 4096$ numbers, which pass through an output projection before being added to the original residual stream. The following GQA section explains how Llama shares keys and values across heads.

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, let’s imagine two vectors: $[1, 2, 3]$ and $[10, 20, 30]$. Both point in the same direction, but [10, 20, 30] is ten times larger than [1, 2, 3]. If we feed them directly into a matrix multiplication, the higher magnitude vector produces a higher magnitude output. As layers repeatedly add updates to the residual stream, its magnitude can change. Normalization reduces the network’s sensitivity to this scale.

For a token’s vector $\vec{x}$ containing $d$ numbers, RMSNorm divides by the square root of the sum of squares of the d numbers 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.¹

How does grouped-query attention work?

The model performs several of these attention computations in parallel, using what are called attention heads. Llama 3.1 8B has 32 query heads, each working with 128-dimensional vectors, but only eight key heads and eight value heads. Each group of four query heads shares one key head and one value head. This is grouped-query attention, or GQA: the queries can produce different attention weights while sharing the keys and values they read from, reducing the amount of data stored in the KV cache during generation. The 32 head outputs are then concatenated into 4,096 numbers per position, passed through a learned output projection, and added back to the residual stream. That’s the attention update: each position gathers a learned mixture of information from itself and earlier positions before the feedforward network processes it.¹¹¹

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.

Feedforward / Llama 3.1 8B

Two MLPs with one behaving as an information gate.

Start with the MLP you know: a learned projection followed by an activation. SwiGLU adds another projection of the same input, then multiplies the two paths together before projecting back down.

The familiar MLP

Input → expand → activate → project back

SwiGLU keeps this flow, adding a second branch at the multiplication below. Llama’s projections have no bias terms.

Scroll each diagram sideways to inspect its matrices →

01 / Expand the same input along two paths

Each token’s 4,096 numbers become two sets of 14,336 numbers. One supplies values; the activated branch supplies a multiplier for each value.

02 / Multiply matching cells — this is the gate

Pair up the cells and multiply; don’t sum them. The gate is a set of continuous multipliers, not on/off switches or probabilities. Its numbers can exceed one or be negative.

03 / Project back to the residual stream’s width

The wider space is temporary. Every token gets a 4,096-number update, ready to add back to its residual stream.

Back to the transformer block

H after attention + M feedforward update = next residual stream N × 4,096

The original H bypasses this calculation. Each token is processed independently, using the same three weight matrices within this layer.

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: Rotary Position Embedding

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. In Llama 3.1 8B, RoPE is applied across the 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

The intuition is that when attention takes a rotated query’s dot product with a rotated key, their shared absolute rotation cancels, and leaves behind a positional contribution that depends on their relative offset. The score still depends on their content.

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

Where does position information 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.

For Llama 3.1 something called a causal mask which is applied during training which does introduce some information about order. When the model is trained, a token can read earlier tokens, but not later ones, so moving a token does change which context it can access. But this does not give the model an explicit representation of the relative position of the token within the input prompt.

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 a newer technique called 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 can be 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.

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

Seeing the rotation pattern

Let’s lay out the angles for one head over 128 token positions. Each row is a position $m$, and each column is one of the 64 coordinate pairs $i$. The color records the rotation angle $m\theta^{-i/64}$, wrapped into the range $[0, 2\pi)$: after a full turn, the color repeats. This gives us a $128 \times 64$ map of angles for the head’s $128 \times 128$ coordinates.

Move the slider to change the RoPE base $\theta$, starting at Llama’s 500,000. Increasing it slows the later pairs, while pair zero always turns by one radian per position. Notice the repeating bands on the left and the almost unchanged colors on the right. This diagram uses the base schedule, before the long-context adjustment described above; the same angle pattern is applied to every query and key head.

One head · 128 positions × 64 coordinate pairs

A map of the rotation angles in one attention head.

RoPE base θ

500,000

1,00010,000,000

Logarithmic scale · base schedule Reset to 500,000

Token position m

0326496127

A 128 by 64 map of angles: (m × θ^(−i/64)) modulo 2π. Position zero has zero rotation; higher pair indices turn more slowly.

016324863

Coordinate pair i

Rotation angle (radians, wrapped)

0π/2π3π/22π ≡ 0

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. RoPE also does not guarantee that attention decreases steadily with distance. The score also depends on the learned query and key vectors.

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. 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 final RMSNorm and unembedding, we have our logits. Before turning them into probabilities, we can divide them by a parameter called temperature. A temperature of one leaves the logits unchanged. A temperature between zero and one makes the resulting probability distribution more concentrated around the highest-scoring tokens. Intuitively, this makes the model more likely to stick with its strongest predictions, which can produce more consistent and sometimes “boring” outputs.

A temperature greater than one has the opposite effect. It makes the distribution flatter, giving lower-scoring tokens more of a chance relative to the highest-scoring ones. This can produce more varied or “crazier” outputs, though those outputs may also be less coherent. Temperature changes the probabilities we sample from; it doesn’t change the model’s weights or the ordering of its logits.

What about a temperature of zero? We can’t actually divide by zero. Instead, generation code conventionally treats this setting as an instruction to pick the highest-scoring token directly. This is called greedy decoding.

Sampling

Softmax

For a positive temperature, we calculate the token probabilities by applying softmax to the scaled logits. This gives us a probability distribution over the vocabulary.

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 original, unscaled logit corresponding to token $i$, and $T > 0$ is the temperature. The division by $T$ in this formula is the temperature scaling we just discussed; we only apply it once.

The exponential ensures that all values are positive, while dividing by the sum ensures that the resulting probabilities add up to one. The figure below follows four scores through those two steps.

Softmax

Scores become a probability distribution.

An illustration of softmax on 4 logits. Each logit is divided by temperature T, then exponentiated. Dividing by the sum gives the probabilities.

Temperature

1.0

0.1 · more concentrated10 · more even

Softmax: pᵢ = exp(zᵢ / T) / Σ exp(zⱼ / T). Exponential bars share a scale that adjusts for large values.

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 choose the highest-probability token at each step. That’s greedy decoding, and with the same prompt and unchanged scores, it produces the same continuation. Alternatively, we can randomly sample according to the probabilities. A token with probability 0.2 then has a 20% chance of being selected on that step. This gives the same prompt a chance to produce different continuations. We can sample from the full distribution, or first restrict the candidates using top-k or top-p.

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 from highest to lowest probability 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 when the temperature is positive, and greedy decoding otherwise. For continuous text generation, the selected token is appended to the sequence, and generation continues from there.

The interactive figure below uses the real next-token logits for “A hippopotamus is drinking water.” Adjust temperature to change the full-vocabulary softmax, then use top-$k$ to keep at most 15 candidates and top-$p$ to narrow that pool further. Both filters renormalize the probabilities. Only the 15 largest logits are displayed. Click Draw next token repeatedly to sample from the final distribution; each draw uses the same prompt. The interval below the bars shows how a uniform random number selects a token: each token occupies a segment as wide as its final probability. At temperature zero, selection is greedy and always returns the highest-scoring token.

Sampling / Llama 3.1 8B

How sampling works in Llama 3.1 8b

A visual explainer of how temperature, top-k and top-p sampling work to generate the next token.

“A hippopotamus is drinking water.”→Llama 3.1 8B→128,256 logits

Temperature

1.00

0 = greedy · 2 = flatter Top-k

10

Keep 1–15 candidates Top-p

0.90

Cutoff after top-k

01 / Model outputToken & logit15 largest scores

Token / idLogit

·It 110214.2500

·The 57814.1875

·A 36213.8750

·What 363913.3750

·Which 1629912.7500

·He 128312.6875

space 22012.2500

·( 32012.2500

·Its 1169912.0625

·How 265012.0000

·This 111511.9375

·She 300511.6875

·Suddenly 5880111.6875

<|end_of_text|> 12800111.6875

·Then 511211.6250

02 / TemperatureSoftmax(logits / T)Over all 128,256 tokens

Probability0–100%

·It14.57%

·The13.69%

·A10.01%

·What6.07%

·Which3.25%

·He3.05%

space1.97%

·(1.97%

·Its1.63%

·How1.54%

·This1.44%

·She1.12%

·Suddenly1.12%

<|end_of_text|>1.12%

·Then1.06%

03 / FilterNarrow the fieldThen rescale to sum to 1

→Top-k

Keep the 10 highest probabilities.

Renormalize

↓Top-p

Keep tokens until cumulative probability reaches the cutoff.

Meta uses temperature → top-p. Top-k is an extra control here, applied first.

04 / Ready to sampleFinal probability7 tokens · total 100%

Probability0–100%

·It27.69%

·The26.01%

·A19.03%

·What11.54%

·Which6.18%

·He5.80%

space3.75%

·(0.00%

·Its0.00%

·How0.00%

·This0.00%

·She0.00%

·Suddenly0.00%

<|end_of_text|>0.00%

·Then0.00%

Sample next token

Sample a token based on the probability distribution above.

0A uniform random sample picks a segment.1

A hippopotamus is drinking water.

No tokens sampled yet. Each click samples the same fixed prompt.

Reset samples

Putting it all together

So let me summarize everything again. First, here’s another look at the high level pipeline of the model.

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.

So after we give an LLM some text, our text is tokenized, embedded, processed with multiple transformer layers, then normalized and unembedded to produce logits. We then select the next token using greedy decoding or sampling, feed that token back into the model, and repeat the process until a special token tells us that the model is finished.

When I first really understood this whole process end-to-end I found myself thinking, “That’s it?”

Sure, the process has a lot of steps and is really 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 to me.

And more importantly, even after I understood how the model produced the next token, there was a crucial question that I still had no answer to. How does this result in models that can solve IMO problems and tackle Millennium Prize proofs at the frontier of mathematics? Is intelligence just a lot of data together with obscene amounts of compute?

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 very new field of mechanistic interpretability is trying to answer by probing the large models with mathematical tools. The hope is that we can 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.

---

This article is based on my understanding of the Llama 3.1 paper. I have done my best to keep it accurate, but if you notice an error, please let me know through the contact form.

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. Hugging Face. Llama 3.1 8B model configuration. vocab_size 128256, num_key_value_heads 8, and the rest of the 8B config. ↩ ↩² ↩³ ↩⁴

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

4. Bengio, Yoshua, et al. “A Neural Probabilistic Language Model.” Journal of Machine Learning Research 3 (2003): 1137–1155. Learns word representations jointly with a language model to generalize across semantically similar word sequences. ↩

5. Mikolov, Tomas, et al. “Efficient Estimation of Word Representations in Vector Space.” arXiv:1301.3781 (2013). Introduces the continuous bag-of-words and skip-gram architectures for efficiently learning word vectors. ↩

6. 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. ↩ ↩² ↩³ ↩⁴ ↩⁵

7. He, Kaiming, et al. “Deep Residual Learning for Image Recognition.” CVPR 2016. Sections 1 and 3.1–3.2 introduce the degradation problem and residual mapping $H(x)=x+F(x)$. Section 4.1 distinguishes the observed optimization difficulty from vanishing gradients. ↩ ↩²

8. Balduzzi, David, et al. “The Shattered Gradients Problem: If resnets are the answer, then what is the question?” ICML 2017. Sections 1.1 and 3 analyze how gradient correlations across inputs decay with depth and how skip connections, normalization, and rescaling reduce shattering. The analysis applies at initialization, rather than establishing a general result throughout training. ↩

9. Elhage, Nelson, et al. “A Mathematical Framework for Transformer Circuits.” Transformer Circuits Thread (2021). The “High-Level Architecture” and “Virtual Weights and the Residual Stream as a Communication Channel” sections describe attention and MLP layers reading from the residual stream and adding their outputs back into it. The computational-workspace description here is an interpretation of that communication-channel framing. ↩

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

11. 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. ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶

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

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

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

15. Holtzman, Ari, et al. “The Curious Case of Neural Text Degeneration.” ICLR 2020. Nucleus sampling keeps the smallest set of tokens whose probabilities sum to at least the cutoff. ↩