Decoder-Only Transformers

Decoder-only transformers are the architecture behind GPT style language models. They stack causal Self-Attention blocks and are trained to predict the next token, so a single model handles generation, completion, and in-context learning.

Masked self-attention

• Masked self-attention is a variant of self-attention used in decoder-only transformers.
• The tokens that follow a given token within the sequence are “masked out” - no attention score is computed between the given token and the following tokens.
  • This prevents looking forward in the sequence during self-attention.

Multi-Head Attention

• Instead of using a single set of $Q,K,V$, we create $h$ independent sets, where each head has $Q_i=Q{W}_i^Q,K_i=K{W}_i^K,V_i=V{W}_i^V$
  • Here, $W^Q,W^K,W^V\in R^{d\times d_h},$ and $d_h=d/h$
• $hea{d}_i=Attention(Q_i,K_i,V_i)\to Multihead(Q,K,V)=Concat(hea{d}_1,...hea{d}_h)W^O$
  • where $W^O\in R^{d\times d}$ is a learned projection back to the original dimension
• We use multi-head attention because it allows for the attending of multiple relevant parts of a sequence simultaneously, with multiple leraned subspaces

Causal Masking

Causal masking enforces autoregression: position $t$ may attend to positions $\le t$ only. It is implemented by adding $-\infty$ to the attention logits above the diagonal before the softmax:

$$\text{Attn}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}+M\right)V,M_{ij}=\{\begin{array}{ll}0& j\le i\\ -\infty & j>i\end{array}$$

The masked entries become zero after softmax, so no gradient or information flows backward in time.

Next-Token Prediction

Training minimizes the cross entropy of the true next token at every position at once (teacher forcing):

$$\mathcal{L}=-\sum _{t=1}^T\log p_{\theta }(x_t\mid x_{<t})$$

Because the mask blocks future tokens, one forward pass yields $T$ supervised predictions, making training highly parallel over the sequence.

The GPT Block in PyTorch

A pre-norm block: LayerNorm, causal attention, residual, then LayerNorm, MLP, residual. einops.rearrange splits heads and torch.einsum computes the attention scores.

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
 
class CausalSelfAttention(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.h = heads
        self.qkv = nn.Linear(dim, 3 * dim)
        self.proj = nn.Linear(dim, dim)
 
    def forward(self, x, kv_cache=None):
        # x: (b, t, d)
        q, k, v = rearrange(self.qkv(x), "b t (three h d) -> three b h t d",
                            three=3, h=self.h)
        if kv_cache is not None:                    # append new k, v at inference
            k = torch.cat([kv_cache["k"], k], dim=2)
            v = torch.cat([kv_cache["v"], v], dim=2)
            kv_cache["k"], kv_cache["v"] = k, v
 
        scale = q.shape[-1] ** -0.5
        scores = torch.einsum("bhqd,bhkd->bhqk", q, k) * scale
        t_q, t_k = scores.shape[-2:]
        mask = torch.ones(t_q, t_k, device=x.device).tril(t_k - t_q).bool()
        scores = scores.masked_fill(~mask, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        out = torch.einsum("bhqk,bhkd->bhqd", attn, v)
        out = rearrange(out, "b h t d -> b t (h d)")
        return self.proj(out)
 
class GPTBlock(nn.Module):
    def __init__(self, dim, heads, mult=4):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
        self.attn = CausalSelfAttention(dim, heads)
        self.mlp = nn.Sequential(nn.Linear(dim, mult * dim), nn.GELU(),
                                 nn.Linear(mult * dim, dim))
 
    def forward(self, x, kv_cache=None):
        x = x + self.attn(self.ln1(x), kv_cache)
        x = x + self.mlp(self.ln2(x))
        return x

KV Cache at Inference

During autoregressive decoding, keys and values for past tokens do not change, so recomputing them each step is wasted work. The KV cache stores $K$ and $V$ per layer and appends only the new token’s entries. See KV Cache.

Prefill vs decode

With a cache the model processes the whole prompt once (prefill), then feeds one token per step (decode) with $q$ of length 1. Note the mask above uses tril(t_k - t_q) so a single query correctly attends to all cached keys. Cache memory grows linearly with sequence length and dominates long context serving; this motivates multi-query and grouped-query attention.