Transformers from scratch

From text to the next token

Start with GPT-2's job

Transformers like GPT-2 exist to model text. Given a sequence of tokens, it gives every token in its vocabulary a score for what should come next. To generate text, we choose one token from those scores, add it to the sequence, and run the model again. This repeated process is autoregressive generation.

One helpful picture: each position reads the words to its left and guesses the word to its right. A position can look back, but it cannot look ahead.

Image from ARENA 3.0 by Callum McDougall.

Image from ARENA 3.0 by Callum McDougall.

This page follows one short prompt the whole way through:

"The cat sat on the mat"

We will follow this prompt from its original text to GPT-2's next-token prediction.

GPT-2 Small dimensions

This week's model is gpt2-small. Keep these dimensions nearby while you work through the page and notebook.

The notebook's Config also includes debug, layer_norm_eps, and init_range. Those settings control optional assertions, numerical stability, and the scale used to initialize weights.

Turn text into token IDs

GPT-2 cannot calculate with raw text. We use a tokenizer to split the text into chunks called tokens and assign each token an integer ID from a fixed vocabulary.

A vocabulary containing only complete words could not handle every name, URL, spelling mistake, or piece of punctuation. Using one byte for every token would handle any text, but common words would require many tokens. GPT-2 uses byte pair encoding to get the benefits of both approaches.

Byte pair encoding starts with 256 possible byte values. During tokenizer training, it repeatedly joins common neighboring sequences. Common text chunks can become single tokens, while uncommon text can still be represented with smaller pieces. GPT-2 ends up with 50,257 vocabulary entries, numbered 0 through 50,256.

The ARENA notebook uses TransformerLens, a Python library that loads transformer models and exposes their internal weights and activations. Its HookedTransformer class wraps GPT-2 and provides helper methods such as to_tokens.

Two details can cause errors:

• A token can include a leading space. The token ' Ralph' differs from the two tokens 'R' and 'alph', so "Ralph" and " Ralph" tokenize differently.
• HookedTransformer.to_tokens prepends a beginning of sequence token by default. This token is called BOS. GPT-2 writes it as <|endoftext|> and assigns it ID 50256.

With BOS included, our prompt becomes 7 tokens:

Here every word is a single token. Quoted strings show a leading space as a literal space. The diagram marks the same space with a dot, and raw tokenizer output marks it with Ġ. Joining the token strings reconstructs the original text, including its spaces.

The ARENA notebook's to_tokens method adds BOS automatically. The raw Hugging Face tokenizer does not, so this example prepends the BOS ID explicitly before decoding each token back to a string.

from transformers import GPT2TokenizerFast

tokenizer = GPT2TokenizerFast.from_pretrained("gpt2", local_files_only=True)
prompt = "The cat sat on the mat"
ids = [tokenizer.bos_token_id] + tokenizer.encode(prompt)
pieces = [tokenizer.decode([token_id]) for token_id in ids]

round_trip = "".join(pieces[1:])
assert round_trip == prompt
print("BOS-prefixed prompt tokens:")
print(" | ".join(f"{piece!r}:{token_id}" for piece, token_id in zip(pieces, ids)))
print(f"Text round trip: {round_trip!r} ({round_trip == prompt})")
for text in ("Ralph", " Ralph"):
    print(f"{text!r} -> {tokenizer.tokenize(text)}")

Tokenization changes the units that GPT-2 reads, but it preserves the original text. The transformer later uses each token ID to select a learned vector from its embedding table.

One sequence gives the model many predictions

During training, GPT-2 predicts the next token after every prefix at the same time. With BOS included, our prompt has 7 positions. The model therefore returns 7 rows of scores in one forward pass.

The row at position j predicts the token at position j + 1. At position j, causal attention lets the model use positions 0 through j and blocks every later position. The model cannot see the token it is supposed to predict.

The prompt supplies a known target for the first 6 rows. The final row predicts the first new token after the prompt.

Turn token IDs into logits

The model receives token IDs with shape (1, 7) and returns logits with shape (1, 7, 50257).

• 1 is the number of sequences in the batch.
• 7 is the number of token positions.
• 50257 is the number of possible next tokens.

A logit is an unrestricted score, not a probability. Logits can be positive or negative, and they do not have to add to one. Softmax turns one row of 50,257 logits into positive probabilities that add to one.

For each position, we can inspect the model's highest-scoring guess for the next token.

import torch
from transformers import GPT2LMHeadModel, GPT2TokenizerFast

tokenizer = GPT2TokenizerFast.from_pretrained("gpt2", local_files_only=True)
model = GPT2LMHeadModel.from_pretrained("gpt2", local_files_only=True)
model.eval()

prompt = "The cat sat on the mat"
input_ids = torch.tensor([[tokenizer.bos_token_id] + tokenizer.encode(prompt)])
with torch.inference_mode():
    logits = model(input_ids).logits[0]

tokens = [tokenizer.decode([i]) for i in input_ids[0].tolist()]
predictions = logits.argmax(dim=-1).tolist()
print("position       argmax next token")
for position in range(3, len(tokens)):
    predicted = tokenizer.decode([predictions[position]])
    print(f"after {tokens[position]!r:6} -> {predicted!r}")

The output predicts ' on' after ' sat' and ' the' after ' on', so those guesses match the prompt. It predicts ' floor' after ' the' instead of the prompt's ' mat'. Both can follow "sat on the". At the final ' mat' position, GPT-2 predicts a comma. To make these guesses, the model reads the earlier words through attention.

Choose a token and repeat

To continue the whole prompt, we select logits[0, -1]. The first index selects the only sequence in the batch. The second index selects its final position. This leaves one score for every token in the vocabulary. Greedy generation appends the token with the highest score, gives the model an 8 token input, and runs the same process again. Other generation methods sample from the probabilities instead. We will compare those methods later on this page.

Follow one prompt through GPT-2

At a high level, token and position embeddings start the residual stream. Each transformer block uses attention to move information between positions and an MLP to update each position. LayerNorm prepares the stream before these operations. After all 12 blocks, a final LayerNorm and unembedding produce the logits. The diagram below reads from bottom to top.

Image from ARENA 3.0 by Callum McDougall.

The first axis is the batch. The second axis lists the seven token positions. The last axis holds the 768 numbers that represent each position. Every transformer block preserves (1, 7, 768). Unembedding changes the last axis to the vocabulary size.

Add token and position embeddings

An embedding is a lookup table. GPT-2's token embedding matrix, W_E, has shape (50257, 768). Each token ID selects one row, so looking up all seven tokens produces a tensor with shape (1, 7, 768).

Attention cannot infer input order on its own, so GPT-2 also looks up a learned vector for each position. The positional embedding matrix, W_pos, has shape (1024, 768). The model adds the token and position vectors:

residual = W_E[token id] + W_pos[position]

The result starts the residual stream. Each position now has one 768 entry vector that represents its token and position.

Update the residual stream

Every later component reads from the residual stream and adds its output back. First, LayerNorm prepares each vector for attention, and the attention output is added to the stream. Then LayerNorm prepares the updated stream for the MLP, and the MLP output is added as well. Earlier information can still affect later layers because these outputs update the stream instead of replacing it. Every update must have the same shape as the stream.

Image from ARENA 3.0 by Callum McDougall.

Attention moves information between positions

Attention is the only part of the model that moves information between positions. The MLP updates each position independently.

Queries, keys, and values

Consider a sequence in which the name Mary appears earlier and the current token is is. To predict Mary next, the current position needs to find an earlier token that contains a name and copy the name information from it.

The current position's query describes the information it needs, such as "who has a name?" Each earlier position's key describes what kind of information is available there. The key at Mary can answer "I have a name," so a strong query and key match gives that position a high attention weight. The value at Mary contains the information that is moved back to the current position, namely that the name is Mary.

Image from ARENA 3.0 by Callum McDougall.

Each attention head learns its own query, key, and value projections. For every position, each projection turns the normalized 768 entry vector into a 64 entry vector. GPT-2 Small computes 12 heads at once, so each of the query, key, and value tensors has shape (1, 7, 12, 64).

Image from ARENA 3.0 by Callum McDougall. Select the image to open it at full resolution.

The (7, 7) grid contains one score for every query position and key position. The code below uses einops.einsum with full axis names. An axis that appears in the inputs but is omitted from the output is summed over. For example, the score calculation omits d_head, so it takes the dot product of each query and key. Dividing by √d_head, which is 8, keeps the scores from growing with vector width.

Causal attention prevents a position from reading tokens that come after it. The model replaces future scores with negative infinity before softmax. Softmax then gives those entries a weight of zero and makes each row sum to one.

import math

import einops
import torch
from transformers import GPT2LMHeadModel, GPT2TokenizerFast

tokenizer = GPT2TokenizerFast.from_pretrained("gpt2", local_files_only=True)
model = GPT2LMHeadModel.from_pretrained("gpt2", local_files_only=True)
model.eval()

prompt = "The cat sat on the mat"
input_ids = torch.tensor([[tokenizer.bos_token_id] + tokenizer.encode(prompt)])
tokens = [tokenizer.decode([i]) for i in input_ids[0].tolist()]
layer = model.transformer.h[0]
n_heads = model.config.n_head
d_head = model.config.n_embd // n_heads

with torch.inference_mode():
    positions = torch.arange(input_ids.shape[1]).unsqueeze(0)
    residual = model.transformer.wte(input_ids) + model.transformer.wpe(positions)
    normalized = layer.ln_1(residual)
    # HF Conv1D stores c_attn as [input, 3 * output]. Convert to ARENA axes.
    qkv_weights = layer.attn.c_attn.weight.split(model.config.n_embd, dim=1)
    qkv_biases = layer.attn.c_attn.bias.split(model.config.n_embd)
    W_Q, W_K = [einops.rearrange(w, "d_model (head d_head) -> head d_model d_head", head=n_heads) for w in qkv_weights[:2]]
    b_Q, b_K = [einops.rearrange(b, "(head d_head) -> head d_head", head=n_heads) for b in qkv_biases[:2]]
    q = einops.einsum(normalized, W_Q, "batch query d_model, head d_model d_head -> batch query head d_head") + b_Q
    k = einops.einsum(normalized, W_K, "batch key d_model, head d_model d_head -> batch key head d_head") + b_K
    scores = einops.einsum(q, k, "batch query head d_head, batch key head d_head -> batch head query key") / math.sqrt(d_head)
    future = torch.triu(torch.ones_like(scores, dtype=torch.bool), diagonal=1)
    pattern = scores.masked_fill(future, -torch.inf).softmax(dim=-1)

head = 0
query = tokens.index(" sat")
print(f"Head {head}, raw scores from {tokens[query]!r} (before masking):")
print(" | ".join(f"{token!r}:{scores[0, head, query, key]:.3f}" for key, token in enumerate(tokens)))
print(f"Head {head}, attention from {tokens[query]!r}:")
print(" | ".join(f"{token!r}:{pattern[0, head, query, key]:.3f}" for key, token in enumerate(tokens)))
future_zero = bool((pattern.masked_select(future) == 0).all())
rows_normalized = bool(torch.allclose(pattern.sum(-1), torch.ones_like(pattern.sum(-1))))
assert future_zero and rows_normalized
print(f"Future weights zero: {future_zero}; rows sum to one: {rows_normalized}")

For layer 0, head 0, the ' sat' query starts with finite scores for every token, including the later tokens ' on', ' the', and ' mat'. After the mask and softmax, those three future positions have weight 0.000. The head puts weight 0.638 on the beginning token and 0.175 on ' sat' itself. These numbers come from GPT-2's learned weights.

The code also verifies every row and every head. It prints Future weights zero: True; rows sum to one: True, which confirms that all future weights are zero and every row sums to one.

The attention pattern supplies weights for a sum of the value vectors. W_O maps each head's 64 entry result to 768 entries. The model sums the results and adds the combined update to the residual stream.

Use the MLP to update each position

The MLP applies the same computation to each position independently. It projects 768 entries to 3,072, applies GELU, and projects the result back to 768. GELU is a smooth activation function. Attention can bring earlier information into a position, and the MLP can then transform the information already at that position.

Apply LayerNorm before attention and the MLP

LayerNorm subtracts the mean of each 768 entry vector, then divides by the square root of its variance plus a small epsilon. This gives the vector mean 0 and variance close to 1 before a learned scale and offset. GPT-2 applies LayerNorm before each attention module, before each MLP, and once before unembedding. The notebook asks you to implement LayerNorm first because the later modules use it.

Assemble one block

The next example is a working reference for the notebook's Section 2 exercises. You can skip its implementation details if you want to solve the notebook without seeing an answer. It loads the first GPT-2 block's weights into the classes from this section and compares their output with Hugging Face GPT-2.

import math
from dataclasses import dataclass

import einops
import torch
from torch import nn
from transformers import GPT2LMHeadModel, GPT2TokenizerFast


@dataclass
class Config:
    d_model: int = 768
    d_head: int = 64
    d_mlp: int = 3072
    n_heads: int = 12
    layer_norm_eps: float = 1e-5


class LayerNorm(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.w = nn.Parameter(torch.ones(cfg.d_model))
        self.b = nn.Parameter(torch.zeros(cfg.d_model))

    def forward(self, residual):
        mean = residual.mean(dim=-1, keepdim=True)
        std = (residual.var(dim=-1, keepdim=True, unbiased=False) + self.cfg.layer_norm_eps).sqrt()
        return (residual - mean) / std * self.w + self.b


class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.W_Q = nn.Parameter(torch.empty(cfg.n_heads, cfg.d_model, cfg.d_head))
        self.W_K = nn.Parameter(torch.empty(cfg.n_heads, cfg.d_model, cfg.d_head))
        self.W_V = nn.Parameter(torch.empty(cfg.n_heads, cfg.d_model, cfg.d_head))
        self.W_O = nn.Parameter(torch.empty(cfg.n_heads, cfg.d_head, cfg.d_model))
        self.b_Q = nn.Parameter(torch.empty(cfg.n_heads, cfg.d_head))
        self.b_K = nn.Parameter(torch.empty(cfg.n_heads, cfg.d_head))
        self.b_V = nn.Parameter(torch.empty(cfg.n_heads, cfg.d_head))
        self.b_O = nn.Parameter(torch.empty(cfg.d_model))

    def forward(self, x):
        q = einops.einsum(x, self.W_Q, "batch pos d_model, head d_model d_head -> batch pos head d_head") + self.b_Q
        k = einops.einsum(x, self.W_K, "batch pos d_model, head d_model d_head -> batch pos head d_head") + self.b_K
        v = einops.einsum(x, self.W_V, "batch pos d_model, head d_model d_head -> batch pos head d_head") + self.b_V
        scores = einops.einsum(q, k, "batch query head d_head, batch key head d_head -> batch head query key") / math.sqrt(self.cfg.d_head)
        mask = torch.triu(torch.ones_like(scores, dtype=torch.bool), diagonal=1)
        pattern = scores.masked_fill(mask, -torch.inf).softmax(-1)
        z = einops.einsum(pattern, v, "batch head query key, batch key head d_head -> batch query head d_head")
        return einops.einsum(z, self.W_O, "batch pos head d_head, head d_head d_model -> batch pos d_model") + self.b_O


class MLP(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.W_in = nn.Parameter(torch.empty(cfg.d_model, cfg.d_mlp))
        self.W_out = nn.Parameter(torch.empty(cfg.d_mlp, cfg.d_model))
        self.b_in = nn.Parameter(torch.empty(cfg.d_mlp))
        self.b_out = nn.Parameter(torch.empty(cfg.d_model))

    def forward(self, x):
        pre = einops.einsum(x, self.W_in, "batch pos d_model, d_model d_mlp -> batch pos d_mlp") + self.b_in
        post = 0.5 * pre * (1 + torch.tanh(math.sqrt(2 / math.pi) * (pre + 0.044715 * pre**3)))
        return einops.einsum(post, self.W_out, "batch pos d_mlp, d_mlp d_model -> batch pos d_model") + self.b_out


class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.ln1, self.attn = LayerNorm(cfg), Attention(cfg)
        self.ln2, self.mlp = LayerNorm(cfg), MLP(cfg)

    def forward(self, resid_pre):
        resid_mid = resid_pre + self.attn(self.ln1(resid_pre))
        return resid_mid + self.mlp(self.ln2(resid_mid))


tokenizer = GPT2TokenizerFast.from_pretrained("gpt2", local_files_only=True)
model = GPT2LMHeadModel.from_pretrained("gpt2", local_files_only=True)
model.eval()
cfg = Config(layer_norm_eps=model.config.layer_norm_epsilon)
block = TransformerBlock(cfg).eval()
hf = model.transformer.h[0]

with torch.no_grad():
    block.ln1.w.copy_(hf.ln_1.weight); block.ln1.b.copy_(hf.ln_1.bias)
    block.ln2.w.copy_(hf.ln_2.weight); block.ln2.b.copy_(hf.ln_2.bias)
    q, k, v = hf.attn.c_attn.weight.split(cfg.d_model, dim=1)
    bq, bk, bv = hf.attn.c_attn.bias.split(cfg.d_model)
    for target, source in zip((block.attn.W_Q, block.attn.W_K, block.attn.W_V), (q, k, v)):
        target.copy_(einops.rearrange(source, "d_model (head d_head) -> head d_model d_head", head=cfg.n_heads))
    for target, source in zip((block.attn.b_Q, block.attn.b_K, block.attn.b_V), (bq, bk, bv)):
        target.copy_(einops.rearrange(source, "(head d_head) -> head d_head", head=cfg.n_heads))
    block.attn.W_O.copy_(einops.rearrange(hf.attn.c_proj.weight, "(head d_head) d_model -> head d_head d_model", head=cfg.n_heads))
    block.attn.b_O.copy_(hf.attn.c_proj.bias)
    block.mlp.W_in.copy_(hf.mlp.c_fc.weight); block.mlp.b_in.copy_(hf.mlp.c_fc.bias)
    block.mlp.W_out.copy_(hf.mlp.c_proj.weight); block.mlp.b_out.copy_(hf.mlp.c_proj.bias)

input_ids = torch.tensor([[tokenizer.bos_token_id] + tokenizer.encode("The cat sat on the mat")])
positions = torch.arange(input_ids.shape[1]).unsqueeze(0)
with torch.inference_mode():
    embeddings = model.transformer.wte(input_ids) + model.transformer.wpe(positions)
    actual = block(embeddings)
    expected = model(input_ids, output_hidden_states=True).hidden_states[1]
max_diff = (actual - expected).abs().max().item()
passed = max_diff < 1e-4
assert passed
print(f"Layer 0 max absolute difference: {max_diff:.8g}")
print(f"Matches Hugging Face within 1e-4: {passed}")

The verified maximum difference is about 0.0000153, so the final boolean is True. This is the same kind of output comparison that load_gpt2_test performs in the notebook tests. It confirms that the block matches the reference model after loading the same weights.

Map the final stream to logits

The unembedding matrix, W_U, has shape (768, 50257). It maps every final residual vector to one score for every vocabulary entry. The complete output has shape (1, 7, 50257). The final slice, logits[0, -1], contains the scores for the token after the prompt. For this prompt, the highest score belongs to ','.

Train and sample

Train on next-token targets

Section 3 of the notebook trains a small transformer with d_model = 32 and 4 layers on TinyStories, a dataset of short machine-generated children's stories.

The targets are the input tokens shifted one place to the left. The model uses token j + 1 as the label for position j, so a 7-token sequence gives 6 input and target pairs in one forward pass. The loss does not use the final row of logits because the sequence does not contain its next token.

Each training step follows the same pattern as Week 1:

1. Run the model and calculate cross-entropy from the shifted targets.
2. Call loss.backward() to calculate gradients.
3. Call optimizer.step() to update the parameters.
4. Call optimizer.zero_grad() so gradients do not build up across steps.

At the start, the model is close to assigning equal probability to all 50,257 tokens. The cross-entropy of this uniform prediction is ln(50257), which is about 10.8. The loss often falls quickly at first because the model learns that common tokens such as " the" and " and" should receive more probability. It then learns common pairs of adjacent tokens. Later gains are slower because the model must use longer context and learn less common patterns.

Weights and Biases is part of the notebook's core training loop. wandb.init() starts a run with the project name, run name, and training arguments. wandb.log() records training loss at each step and test accuracy after evaluation. wandb.finish() closes the run after training. The trainer also prints a completion for "Once upon a time" before training and after each epoch. The separate exercise that logs several sample completions in a Weights and Biases table is optional.

Watch the loss curve during training. A working run should show falling loss. After training, completions of "Once upon a time" should contain recognizable story fragments instead of random tokens.

Turn logits into text

Section 4 builds a TransformerSampler. Each decoding method chooses the next token differently.

Greedy decoding

Greedy decoding always selects the token with the highest logit. It is deterministic and can repeat itself. For our prompt, the greedy continuation begins:

"The cat sat on the mat, its head resting on the mat"

Greedy decoding repeats "the mat" from the prompt. This is a common weakness of always taking the highest-scoring token.

Temperature and top k

Temperature rescales logits before softmax.

• A temperature below 1 puts more probability on the highest-scoring tokens.
• A temperature above 1 spreads probability over more tokens.

import torch
from transformers import GPT2LMHeadModel, GPT2TokenizerFast

tokenizer = GPT2TokenizerFast.from_pretrained("gpt2", local_files_only=True)
model = GPT2LMHeadModel.from_pretrained("gpt2", local_files_only=True)
model.eval()
input_ids = torch.tensor([[tokenizer.bos_token_id] + tokenizer.encode("The cat sat on the mat")])

with torch.inference_mode():
    logits = model(input_ids).logits[0, -1]

for temperature in (0.5, 1.0, 2.0):
    probabilities = (logits / temperature).softmax(dim=-1)
    values, indices = probabilities.topk(5)
    items = ", ".join(f"{tokenizer.decode([i])!r} {p:.3%}" for p, i in zip(values.tolist(), indices.tolist()))
    print(f"T={temperature}: {items}")

values, indices = logits.topk(5)
top_k_probabilities = values.softmax(dim=-1)
items = ", ".join(f"{tokenizer.decode([i])!r} {p:.2%}" for p, i in zip(top_k_probabilities.tolist(), indices.tolist()))
print(f"top-k (k=5), renormalized: {items}")

For this prompt, GPT-2's five leading tokens are ',', ' and', ' in', ' with', and ' for'. At temperature 0.5, their printed probabilities are about 45.5%, 20.5%, 7.7%, 7.2%, and 6.7%. At temperature 1.0, the first two fall to about 17.9% and 12.0%. At temperature 2.0, all five are close to 1%. Temperature 2 spreads probability over the full vocabulary, so the five printed probabilities do not need to sum to anything near one.

Top k uses the same logits but keeps only the k highest-scoring tokens. It removes every other token and then samples from the kept set. With k = 5, this example renormalizes the five probabilities to about 34.9%, 23.5%, 14.3%, 13.9%, and 13.4%, which add to one.

Frequency penalty, top p, and beam search

A frequency penalty lowers a token's logit according to how many times that token already appears in the input. If a token has appeared three times and the penalty is 2, its logit is reduced by 6. The notebook's apply_frequency_penalty exercise asks you to implement this rule and run its tests.

Top p keeps the smallest set of leading tokens whose total probability reaches p, then samples from that set. Unlike top k, the size of this set can change at every step.

A log probability is the natural logarithm of a probability. It is always zero or negative because probabilities range from zero to one. Beam search keeps several candidate continuations and ranks them by the sum of their token log probabilities. The notebook asks you to implement the beam generation and filtering steps.

Image from ARENA 3.0 by Callum McDougall.

Inspect GPT-2 in the interactive explainer

Session checklist

The core path covers all four main sections of the 1.1 Transformers from Scratch notebook.

Your implementation is ready when:

• Every implementation exercise passes every test listed beneath it. This includes test_layer_norm_epsilon, test_embed, test_pos_embed, test_causal_mask, test_unembed, the random input tests, and load_gpt2_test where the notebook uses them.
• DemoTransformer passes its rand_int_test and load_gpt2_test, and it reproduces the reference GPT-2 outputs with the same weights.
• The demo prompt predicts ' on' after ' sat' and ' the' after ' on' near the end, and ',' after the final ' mat'.
• TinyStories training produces falling loss, and completions of "Once upon a time" contain recognizable story fragments.
• The sampler passes its tests, and temperature behaves as shown in Example 5.

Quiz before you start the notebook

Choose one answer for each question, then mark your quiz.

Stretch work

All four main sections are core. If you finish early, try these optional exercises:

1. Complete the bonus caching exercise, which stores key and value tensors so generation does not recompute the whole prompt for each new token.
2. Complete the cached beam search bonus after that.
3. Watch Neel Nanda's Transformer Circuits walkthrough, which ARENA links for extra detail on this architecture.
4. Read Hugging Face's How to generate text, the post ARENA assigns before the sampling section, if you skipped it.