INSIDE THE
TRANSFORMER

Before the vectors: tokenization

A transformer block can only take matrices as input, so the first real question is what representation of raw text actually conveys its essence. We start by dividing text into chunks called tokens. The obvious first guess is to split based on words — older NLP pipelines did exactly that — but that captures less nuance than we'd like, and it isn't free: it's a trade-off between vocabulary size and sequence length, and both heavily affect a model's cost and capability. Wait a minute …. Before we jump ahead of ourselves let us establish some definitions first so that we are all on the same page. Definition alert !!! Jump to Definitions

Split word-by-word and the vocabulary explodes the moment you account for misspellings, inflected forms, and out-of-vocabulary words. Drop all the way to character- or byte-level and the vocabulary shrinks to a tidy 256 bytes — but the sequence length balloons to roughly four times the word-level length, and any single character carries too little information on its own, pushing that burden onto later layers.

GPT-3's 50,257-token vocabulary is chosen to sit precisely at that sweet spot: the model can process roughly 4× more text in the same context window than a character-level model would allow, while still representing any possible Unicode string without ever needing an <UNK> token — and it also helps with reconstructing word boundaries. I have built a simple tool you can play around with which tells you how different sentences are split and which token represents a particular number in the vocab: Tokenization Tool

Byte‑pair encoding – mathematical formulation

GPT-3 uses a byte-level byte-pair encoding (BPE) it operates on raw bytes (UTF-8), which naturally handles any Unicode text, and it pre‑splits text with a regex before merging. The tokenization process has two stages. The tokenizer is defined by a fixed regular expression pre‑splitter, an ordered list of merge rules \(\mathcal{M}\), and a vocabulary \(\mathcal{V}\).

Training

1. Pre‑tokenization. The corpus is split into a multiset of chunks \(C = \{c_1, c_2, \dots\}\) using the regex pattern.
   Each chunk \(c\) is converted to a byte sequence \(\mathbf{b} = (b_1,\dots,b_\ell),\ b_i \in \{0,\dots,255\}\).
2. Initial vocabulary. \(\mathcal{V}_0 = \{0,1,\dots,255\}\). Every token is a single byte.
3. For \(k = 1\) to \(N\) (with \(N = 50\,000\)):
   • Count frequencies of all adjacent token pairs \((p,q)\) across all current token sequences.
   • Choose the most frequent pair: \[ (p_k, q_k) = \arg\max_{(p,q)} \text{freq}(p,q). \]
   • Create a new token \(r_k = 255 + k\) and update the vocabulary: \[ \mathcal{V}_k = \mathcal{V}_{k-1} \cup \{r_k\}. \]
   • Apply the merge: in every token sequence, replace every occurrence of the contiguous pair \((p_k,q_k)\) by \(r_k\).
4. Final vocabulary. \(\mathcal{V} = \mathcal{V}_0 \cup \{r_1,\dots,r_N\} \cup \{\texttt{<|endoftext|>}\}\), with \(|\mathcal{V}| = 50\,257\).
5. Merge rule list. The ordered list \[ \mathcal{M} = \big((p_1,q_1)\to r_1,\; (p_2,q_2)\to r_2,\; \dots,\; (p_N,q_N)\to r_N\big) \] is the sole parameter that defines encoding and decoding.

Encoding a new text string \(s\)

1. Regex split. \(s \mapsto (c_1, c_2, \dots, c_m)\).
2. Byte encode each chunk. \(c_j \mapsto \mathbf{b}^{(j)} = (b_1^{(j)},\dots,b_{\ell_j}^{(j)}),\ b_i^{(j)} \in \{0,\dots,255\}\).
3. Apply merge list. For each chunk’s byte sequence \(\mathbf{t} \gets \mathbf{b}^{(j)}\), iterate through \(\mathcal{M}\) in order: \[ \text{for } (p,q)\to r \text{ in } \mathcal{M},\ \text{replace every leftmost occurrence of } (p,q) \text{ in } \mathbf{t} \text{ by } r. \] This yields a sequence of tokens from \(\mathcal{V}\).
4. Concatenate chunk token sequences. The final token list is the input to the transformer.

Formally, the encoder is a deterministic function \(\operatorname{encode}: \Sigma^* \to \mathcal{V}^*\) where \(\Sigma\) is the set of Unicode characters.

Decoding

Given a sequence \((t_1,\dots,t_k),\ t_i \in \mathcal{V}\):

1. Recursively expand each token \(t_i\) by reversing the merge rules (if \(t_i = r_k\), replace it by \((p_k,q_k)\)) until only bytes remain.
2. Concatenate all resulting bytes and decode the UTF‑8 byte stream: \(\operatorname{decode}: \mathcal{V}^* \to \Sigma^*\).

Other tokenization schemes in use

Choosing a vocabulary and a tokenization scheme is still a tricky, mostly iterative and intuition-driven craft. It's also worth noting that tokenization has been used as a method for jailbreaking LLMs via adversarial inputs, and that certain anomalous "glitch tokens" cause undocumented failures in GPT-style models — some break determinism even at temperature zero. Further reading here. As we'll see later, the architecture is built around maximizing the probability of the next token rather than chasing a single ground truth — so the model is always giving its best approximation for the next possible word

hidden_state = self.embed_dropout(tok_vec + pos_vec)  # h_0, shape (B, T, d_model)

token_ids = x                                    # (batch, seq_len)
tok_vec = self.token_embedding(token_ids)        # lookup & shape (B, T, d_model)
tok_vec = tok_vec * math.sqrt(self.d_model)  # scaling

positions = torch.arange(seq_len, device=x.device)   # (T,)
pos_vec = self.position_embedding(positions)     # (T, d_model), broadcasts over batch

Alternatives to sinusoidal position

# Embedding lookup — the very first step of the forward pass
# token_embedding:    (vocab_size, d_model)   — one row per vocabulary entry
# position_embedding: (max_seq_len, d_model)  — one row per sequence position
token_ids = x                                        # (batch, seq_len) integer ids
positions = torch.arange(seq_len, device=x.device)   # (seq_len,)
tok_vec = self.token_embedding(token_ids)            # (batch, seq_len, d_model)
pos_vec = self.position_embedding(positions)         # (seq_len, d_model), broadcasts over batch
hidden_state = self.embed_dropout(tok_vec + pos_vec)  # this is h_0

Attention mechanism

GPT‑2 uses multi‑head causal self‑attention in every decoder block. For a given layer, the input is the hidden state h ∈ ℝ^B×T×d, where B is batch size, T is sequence length, and d is the model dimension. The output is a context‑aware tensor of the same shape, computed as follows.

Multi‑head scaled dot‑product attention with causal masking

1. Project to queries, keys, and values. A single weight matrix W_{c_attn} ∈ ℝ^d×3d and bias b_{c_attn} ∈ ℝ^3d produce a combined projection: \[ \text{qkv} = h\,W_{\text{c\_attn}} + b_{\text{c\_attn}} \quad\in \mathbb{R}^{B\times T\times 3d}. \] Split the result into three equally‑sized chunks along the last dimension: \[ Q = \text{qkv}[:,:,\,0:d],\quad K = \text{qkv}[:,:,\,d:2d],\quad V = \text{qkv}[:,:,\,2d:3d]. \]
2. Divide into heads. Let n_heads be the number of attention heads (e.g. 12 for GPT‑2 small) and d_k = d / n_heads. Reshape and transpose so that each head processes a d_k‑dimensional subspace independently: \[ Q \to (B, n_{\text{heads}}, T, d_k),\quad K \to (B, n_{\text{heads}}, T, d_k),\quad V \to (B, n_{\text{heads}}, T, d_k). \]
3. Compute raw attention scores. For each head, the unnormalised attention logits are the scaled dot products between queries and keys: \[ \text{scores}_{h,i,j} = \frac{1}{\sqrt{d_k}} \sum_{m=1}^{d_k} Q_{h,i,m}\, K_{h,j,m}, \] where h indexes the head, i the query position, and j the key position. In matrix form: \[ S_h = \frac{Q_h K_h^\top}{\sqrt{d_k}} \quad\in \mathbb{R}^{T\times T}. \]
4. Apply the causal mask. GPT‑2 is an autoregressive model; the i‑th query must only attend to positions j ≤ i. A lower‑triangular mask M is added to the scores: \[ M_{i,j} = \begin{cases} 0, & \text{if } i \geq j,\\ -\infty, & \text{if } i < j. \end{cases} \] The masked scores are: \[ \tilde{S}_h = S_h + M. \]
5. Softmax normalisation. Convert masked scores to attention probabilities: \[ A_h = \operatorname{softmax}(\tilde{S}_h),\qquad (A_h)_{i,j} = \frac{\exp\big((\tilde{S}_h)_{i,j}\big)}{\sum_{k=0}^{T-1}\exp\big((\tilde{S}_h)_{i,k}\big)}. \] Because of the mask, all probabilities for future positions (j > i) are zero.
6. Weighted sum of values. For each head, the output is a convex combination of the value vectors: \[ O_h = A_h\, V_h \quad\in \mathbb{R}^{T\times d_k}. \] This is repeated for every head, yielding n_heads output tensors.
7. Concatenate heads and project. The per‑head outputs are concatenated along the d_k dimension, restoring the original model size d: \[ O_{\text{concat}} = [O_1; O_2; \dots; O_{n_{\text{heads}}}] \in \mathbb{R}^{B\times T\times d}. \] Finally, a linear projection (with weights W_{c_proj} ∈ ℝ^d×d and bias b_{c_proj} ∈ ℝ^d) produces the attention sub‑layer output: \[ \text{Attn}(h) = O_{\text{concat}}\, W_{\text{c\_proj}} + b_{\text{c\_proj}}. \]

Integration into the decoder block

The attention equation

Every decoder block projects its normalized hidden state into queries, keys, and values with one fused weight matrix, then splits the result across heads:

# qkv_proj: (d_model, 3 * d_model) — one combined matrix for Q, K, and V
q, k, v = qkv_proj(normed_hidden).split(d_model, dim=-1)
# reshape so each of the H heads attends independently
q = q.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)
k = k.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)
v = v.view(batch, seq_len, num_heads, head_dim).transpose(1, 2)

Scaled dot-product attention with the causal mask:

attention_weights = softmax((q @ k.transpose(-2, -1)) / sqrt(head_dim) + causal_mask)
context = attention_weights @ v

The 1/√d_h scaling keeps dot products from growing too large as head dimension increases, which would otherwise push softmax toward a near one-hot, hard-to-train distribution. The causal mask sets every position j > i to −∞ before the softmax, so token i can only ever attend to itself and whatever came before it — the mechanism that makes autoregressive generation well-defined at all.

Attention module, renamed and commented

Setup happens once, in __init__ — the fused projection, the output projection, dropout, and (only as a fallback) an explicit causal mask buffer:

class MultiHeadSelfAttention(nn.Module):
    """Causal multi-head self-attention with a single fused QKV projection."""
    def __init__(self, cfg):
        super().__init__()
        assert cfg.d_model % cfg.num_heads == 0
        # one matrix produces queries, keys, and values together — cheaper than three separate ones
        self.qkv_proj = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=cfg.use_bias)
        self.attn_out_proj = nn.Linear(cfg.d_model, cfg.d_model, bias=cfg.use_bias)
        self.attn_drop = nn.Dropout(cfg.drop_prob)
        self.resid_drop = nn.Dropout(cfg.drop_prob)
        self.num_heads = cfg.num_heads
        self.d_model = cfg.d_model
        self.drop_prob = cfg.drop_prob
        # prefer a fused, hardware-optimized attention kernel when available
        self.use_flash = hasattr(F, "scaled_dot_product_attention")
        if not self.use_flash:
            # fallback: an explicit lower-triangular mask, built once and reused
            causal = torch.tril(torch.ones(cfg.max_seq_len, cfg.max_seq_len))
            self.register_buffer("causal_mask", causal.view(1, 1, cfg.max_seq_len, cfg.max_seq_len))

And forward does the actual reshaping, dispatch to the fused kernel when it's available, and the manual fallback path when it isn't:

    def forward(self, x):
        batch, seq_len, d_model = x.size()
        q, k, v = self.qkv_proj(x).split(self.d_model, dim=2)
        head_dim = d_model // self.num_heads
        q = q.view(batch, seq_len, self.num_heads, head_dim).transpose(1, 2)
        k = k.view(batch, seq_len, self.num_heads, head_dim).transpose(1, 2)
        v = v.view(batch, seq_len, self.num_heads, head_dim).transpose(1, 2)
        if self.use_flash:
            # fused kernel — see FORGE-1's spec panel (FlashAttention-3)
            out = F.scaled_dot_product_attention(
                q, k, v, attn_mask=None,
                dropout_p=self.drop_prob if self.training else 0,
                is_causal=True,
            )
        else:
            scores = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(head_dim))
            scores = scores.masked_fill(self.causal_mask[:, :, :seq_len, :seq_len] == 0, float("-inf"))
            weights = self.attn_drop(F.softmax(scores, dim=-1))
            out = weights @ v
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, d_model)
        return self.resid_drop(self.attn_out_proj(out))

Feed‑forward network – mathematical formulation

The feed‑forward network (FFN) is a position‑wise sublayer. It processes each token independently, introduces non‑linearity, and contains the bulk of the model's capacity. The operation is identical for every sequence position.

Standard ReLU FFN (Vaswani et al., 2017)

Two affine transformations with a ReLU activation in between.

Parameters
\(W_1 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}\), \(b_1 \in \mathbb{R}^{d_{\text{ff}}}\)
\(W_2 \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}\), \(b_2 \in \mathbb{R}^{d_{\text{model}}}\)
\(d_{\text{ff}} = 4 \cdot d_{\text{model}}\) in the original paper.

Per‑token operation (for a single vector \(x \in \mathbb{R}^{d_{\text{model}}}\)):

\[ \operatorname{FFN}(x) = \max(0,\; xW_1 + b_1)\,W_2 + b_2 \]

Full sequence (input matrix \(X \in \mathbb{R}^{n \times d_{\text{model}}}\)):

\[ \operatorname{FFN}(X) = \max(0,\; XW_1 + \mathbf{1}_n b_1^T)\,W_2 + \mathbf{1}_n b_2^T \]

This is equivalent to two independent \(1\!\times\!1\) convolutions applied to every token vector.

Variant: GELU activation

Used in BERT, GPT‑2, and many encoder models. Smooth alternative to ReLU:

\[ \operatorname{GELU}(x) = x \cdot \Phi(x) \approx 0.5x\left(1 + \tanh\left[\sqrt{2/\pi}\left(x + 0.044715 x^3\right)\right]\right) \] \[ \operatorname{FFN}_{\text{GELU}}(x) = \operatorname{GELU}(xW_1 + b_1)\,W_2 + b_2 \]

Variant: SwiGLU (gated linear unit)

Modern large language models (PaLM, LLaMA, etc.) use a gated variant with a Swish activation, improving performance while maintaining comparable computational cost.

Parameters
\(W_1, W_3 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}\), \(W_2 \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}\)
(Biases are typically omitted.)

Operation (element‑wise gating):

\[ \operatorname{SwiGLU}(x) = \big(\operatorname{Swish}(xW_1) \;\odot\; xW_3\big)\,W_2 \]

where \(\operatorname{Swish}(z) = z \cdot \sigma(z) = z \cdot \operatorname{silu}(z)\). The inner dimension \(d_{\text{ff}}\) is often rescaled (e.g., \(\frac{8}{3}d_{\text{model}}\)) to keep the parameter count similar to the standard FFN.

Forward pass of the FFN sublayer

Inside a Transformer block, the FFN is wrapped with a residual connection and layer normalisation. Given input \(X\) (after attention and residual):

1. Normalise. \(\bar{X} = \operatorname{LayerNorm}(X)\)
2. Compute FFN. \(Y = \operatorname{FFN}(\bar{X})\) using one of the formulations above.
3. Residual addition. \(X_{\text{out}} = X + Y\)

Key properties

• No token mixing – the FFN preserves the sequence length \(n\); it operates purely as a per‑token non‑linear feature extractor.
• Linear complexity in \(n\) – \(O(n \cdot d_{\text{model}} \cdot d_{\text{ff}})\), unlike the quadratic self‑attention. For long sequences the FFN can dominate the per‑token computation.
• Memory interpretation – the large weight matrices can be viewed as key‑value memories that store linguistic and factual knowledge.

Position-wise, in code

Every position runs through the identical MLP independently — no mixing across positions happens here, that already happened in attention. The expand-then-contract shape gives the network room to compute something more elaborate than a single linear pass would allow before compressing back down to the model's working width.

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.mlp_up_proj = nn.Linear(cfg.d_model, 4 * cfg.d_model, bias=cfg.use_bias)
        self.activation = nn.GELU()
        self.mlp_down_proj = nn.Linear(4 * cfg.d_model, cfg.d_model, bias=cfg.use_bias)
        self.drop = nn.Dropout(cfg.drop_prob)
    def forward(self, x):
        x = self.mlp_up_proj(x)
        x = self.activation(x)
        x = self.mlp_down_proj(x)
        return self.drop(x)

The forward pass

GPT — Generative Pretrained Transformer — is the specific decoder-only recipe this walkthrough builds towards. There are several different types of transformer varaints.

layer normalization

attention block

normalization

feed forward layer

norm and gets added to residual stream + repeat

Assembling the decoder block

class DecoderBlock(nn.Module):
    """One pre-norm attention sub-layer + one pre-norm feed-forward sub-layer,
    each wrapped in its own residual connection."""
    def __init__(self, cfg):
        super().__init__()
        self.pre_attn_norm = LayerNorm(cfg.d_model, cfg.use_bias)
        self.attn = MultiHeadSelfAttention(cfg)
        self.pre_mlp_norm = LayerNorm(cfg.d_model, cfg.use_bias)
        self.mlp = FeedForward(cfg)
    def forward(self, x):
        x = x + self.attn(self.pre_attn_norm(x))   # attention sub-layer, residual add
        x = x + self.mlp(self.pre_mlp_norm(x))      # feed-forward sub-layer, residual add
        return x

The full model

Everything upstream — token embedding, position embedding, the stack of decoder blocks, and the tied unembedding head — is assembled once in __init__:

class ForgeTransformer(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.backbone = nn.ModuleDict(dict(
            token_embedding=nn.Embedding(cfg.vocab_size, cfg.d_model),
            position_embedding=nn.Embedding(cfg.max_seq_len, cfg.d_model),
            embed_dropout=nn.Dropout(cfg.drop_prob),
            layers=nn.ModuleList([DecoderBlock(cfg) for _ in range(cfg.num_layers)]),
            final_norm=LayerNorm(cfg.d_model, cfg.use_bias),
        ))
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        # weight tying: the unembedding is literally the embedding matrix, transposed
        self.backbone.token_embedding.weight = self.lm_head.weight

And forward runs the actual pass — embed, loop through every decoder block, normalize, then either compute a training loss or return only the last position's logits for inference:

    def forward(self, token_ids, targets=None):
        batch, seq_len = token_ids.size()
        positions = torch.arange(seq_len, device=token_ids.device)
        tok_vec = self.backbone.token_embedding(token_ids)
        pos_vec = self.backbone.position_embedding(positions)
        x = self.backbone.embed_dropout(tok_vec + pos_vec)
        for layer in self.backbone.layers:
            x = layer(x)
        x = self.backbone.final_norm(x)
        if targets is not None:
            logits = self.lm_head(x)
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
        else:
            # inference: only the last position's logits are actually needed
            logits = self.lm_head(x[:, [-1], :])
            loss = None
        return logits, loss

From hidden state to token: the unembedding pipeline

At the final layer of the transformer, each position in the sequence holds a hidden state vector \(\mathbf{h} \in \mathbb{R}^{d_{\text{model}}}\). For autoregressive generation, the hidden state at the last position is used to predict the next token. The unembedding matrix \(W_U \in \mathbb{R}^{|\mathcal{V}| \times d_{\text{model}}}\) maps this hidden state to a score for every token in the vocabulary:

\[ \mathbf{z} = W_U \, \mathbf{h} \qquad \text{(logits)} \]

Applying the softmax function yields a proper probability distribution:

\[ P(t_i \mid \text{context}) = \frac{\exp(z_i / \tau)}{\sum_{j=1}^{|\mathcal{V}|} \exp(z_j / \tau)} \]

where \(\tau\) is the temperature parameter — lower values sharpen the distribution (more deterministic), higher values flatten it (more exploratory). At \(\tau=1\) this is the standard softmax. The model then either selects the token with the highest probability (greedy decoding) or samples from the distribution. Finally, the chosen token ID is decoded back to text by reversing the BPE merge rules — recursively expanding merged tokens until only raw bytes remain, then decoding those bytes via UTF-8.

The interactive tool below demonstrates this entire pipeline for a single position, using the same 4‑dimensional toy embedding space and small vocabulary introduced in the BPE encoding example above.

Mathematical formulation of unembedding & decoding

Step 1 — Logit computation

Given hidden state \(\mathbf{h} \in \mathbb{R}^{d}\) and unembedding matrix \(W_U \in \mathbb{R}^{|\mathcal{V}|\times d}\):

\[ z_i = \sum_{k=1}^{d} (W_U)_{i,k} \cdot h_k \quad \text{for } i = 1,\dots,|\mathcal{V}| \]

Equivalently, \(\mathbf{z} = W_U \mathbf{h}\). Each logit \(z_i\) is the dot product between the hidden state and the learned vector for token \(i\).

Step 2 — Softmax with temperature

\[ P(t_i) = \frac{\exp(z_i / \tau)}{\sum_{j=1}^{|\mathcal{V}|} \exp(z_j / \tau)} \]

As \(\tau \to 0\), the distribution approaches a one-hot encoding of the argmax token (greedy). As \(\tau \to \infty\), it approaches a uniform distribution.

Step 3 — Token selection

\[ t^* = \arg\max_i P(t_i) \quad \text{(greedy)} \quad \text{or} \quad t^* \sim P(t) \quad \text{(sampling)} \]

Step 4 — BPE decoding

Given token \(t^*\), recursively apply the inverse of the merge list \(\mathcal{M}\):

1. If \(t^* \leq 255\), it is a raw byte — emit the corresponding UTF-8 byte.
2. Otherwise, look up the merge rule \((p,q) \to t^*\) in \(\mathcal{M}\) and recursively decode \(p\) then \(q\).
3. Concatenate all emitted bytes and decode the UTF-8 byte stream to obtain the final string.

This guarantees a lossless round-trip: decode(encode(s)) = s for any Unicode string \(s\).

Weight tying: why \(W_U\) often equals \(W_E^\top\)

In many transformer architectures (including GPT-2/3, LLaMA, and others), the unembedding matrix is not independently learned — instead, it's set to the transpose of the embedding matrix: \(W_U = W_E^\top\). This is called weight tying. The intuition is elegant: if the embedding for "run" is the vector that represents that token in the input space, then the same vector should serve as the detector for that token in the output space. The dot product \(\mathbf{h} \cdot \mathbf{e}_{\text{token}}\) measures how closely the hidden state aligns with that token's learned representation.

Weight tying reduces the parameter count by \(|\mathcal{V}| \times d_{\text{model}}\) (roughly 50,257 × 768 ≈ 38.6M parameters for a small GPT-style model) and has been shown to improve generalization, especially for rare tokens whose output embeddings would otherwise receive very little gradient signal during training.