GPT-2: a decoder that writes one token at a time

By Hiva Mohammadzadeh · Stanford MSCS, UC Berkeley EECS

What makes GPT-2 different

GPT-2 is a pre-trained deep learning model. It uses unidirectional transformers to generate one token at a time. The word "unidirectional" is the main distinction between GPT-2 and models like BERT.

BERT is a bidirectional encoder. It sees the entire input sequence at once and builds representations by attending to tokens on both sides of every position. That suits understanding tasks such as classification, entity recognition, and question answering, but it cannot generate text. It has no notion of "past" and "future" in a sequence.

GPT-2 works the other way. It is a decoder-only model that generates output autoregressively. It learns to predict the next word in a sequence from the previous words. At each step the model looks only backward and never at tokens that come later. A causal mask in the attention mechanism enforces this constraint, and that constraint is what makes GPT-2 a generative model.

The core architecture has 12 Transformer Decoder blocks. Each block has two sub-layers: a masked multi-head self-attention mechanism and a position-wise fully connected feed-forward network. For GPT-2 Small, the variant I focus on here, the configuration is 128 tokens, a 768 embedding dimension, and a 3072 feed-forward filter size.

Training: next-token prediction

GPT-2 training happens in two stages.

Pre-training. The model is pre-trained on a large, diverse text corpus with an unsupervised learning approach. The objective is to predict the next token from a sequence. The model sees tokens 1 through t and learns to predict token t+1. Training runs across billions of tokens, and it gives the model a broad understanding of language structure, grammar, facts, and some reasoning ability.

Fine-tuning. After pre-training, the model is fine-tuned on specific downstream tasks by adding a task-specific output layer. The pre-trained weights provide a strong initialization, and fine-tuning adapts the model to the distribution and format of the target task.

This two-stage approach is what "pre-trained" means in the name. The model arrives at your fine-tuning task already knowing a lot about language. You do not train from scratch; you specialize an existing model.

The architecture

The GPT-2 architecture has two parts: the Transformer Decoder stack and a task-specific output layer on top.

Each decoder block contains two sub-layers in sequence: a masked multi-head self-attention layer, then a position-wise feed-forward network. Both sub-layers use residual connections and layer normalization.

The input tokens are embedded into 768-dimensional vectors, and those embeddings flow upward through all 12 decoder blocks. The output at the top is a distribution over the vocabulary for the next-token prediction.

The masked attention mechanism

This masking is central to GPT-2, and it is where the difference from BERT lives.

Each block has two components: the attention computation and the feed-forward network that follows it. This section breaks down the attention computation.

Masked self-attention lets the model attend to different parts of the input sequence, with one constraint. Each position can attend only to positions at or before it. A lower-triangular mask enforces this by setting all future attention scores to negative infinity before the softmax. After softmax those entries become zero, so the model cannot attend to future tokens.

A few design choices in the decoder block matter:

• Residual connections avoid the vanishing gradient problem. The input to each sub-layer is added back to the output, so gradients always have a direct path backward through the network.
• Layer normalization improves the model's convergence speed by normalizing activations within each layer.
• GELU is the nonlinearity in the feed-forward network. It performs better than other activation functions in transformer architectures, because smoother gradients give smoother training. I covered this in my activation functions post.

Dimensions of the weight matrices

The concrete dimensions matter when you implement this or count compute. For GPT-2 Small:

• Input embedding: (batch_size, 128, 768), which can be reshaped as (batch_size, 128, 64, 12) for parallel processing across the 12 attention heads.
• Projection weight matrices (Wq, Wk, Wv, Wo): each is (768, 768). These project the input into the query, key, value, and output spaces.
• Per-head Q, K, V matrices: Query is (batch_size, 1 token, 768), or equivalently (batch_size, 128, 64, 12). Key and Value are (batch_size, sequence_length, 768), or (batch_size, 128, 64, 12).
• Feed-forward weights: the first dense layer is (768, 3072) and the second dense layer is (3072, 768).

The computation flow

The attention block runs the following steps.

The input X has shape (L, Ed) = (128, 768). It is projected through three weight matrices to produce Q, K, and V. Then:

1. Compute Q * K^T to get raw attention scores.
2. Apply the causal mask, which sets upper-triangular entries to negative infinity.
3. Apply softmax to get normalized attention weights (shape: L x L per head).
4. Multiply the attention weights by V to get per-head outputs z_h.
5. Concatenate all 12 heads to get z with shape (L, Ed).
6. Project through Wo (shape Ed x Ed) to get the attention output (L, Ed).
7. Add the residual connection (the original input X).
8. Apply layer normalization.

The output is (128, 768), the same shape as the input, ready for the feed-forward layer.

The feed-forward network

The FFN in GPT-2 follows the same structure as BERT's feed-forward layer. It is a two-layer network with a GELU activation in between.

The computation is:

1. Take the input y (the output of the attention sub-layer), shape (128, 768).
2. Apply a linear projection: (Ed, Fl) = (768, 3072). This expands the representation to 4x the embedding dimension.
3. Apply GELU activation.
4. Apply a second linear projection: (Fl, Ed) = (3072, 768). This compresses back to the embedding dimension.
5. Add the residual connection (the input y).
6. Apply layer normalization.

The output is again (128, 768). The expansion to 3072 and back to 768 gives the network a wider internal representation to work with before compressing it back down. This bottleneck structure is standard across transformer architectures.

Full architecture walkthrough

The complete GPT-2 block chains the attention sub-layer and the FFN sub-layer in sequence, each with its own residual connection and layer normalization.

The flow for one decoder block:

Input (128, 768) → Masked Multi-Head Self-Attention → + Residual → Layer Norm → FFN → + Residual → Layer Norm → Output (128, 768)

This block is repeated 12 times. The output of block 1 feeds into block 2, and so on. After all 12 blocks, the final representation is projected to the vocabulary size for next-token prediction.

Counting FLOPs

Compute cost matters when you choose model sizes, plan training runs, or compare architectures. The FLOP breakdown for one forward pass through GPT-2 Small follows.

Attention block FLOPs

Weight matrix projections (Q, K, V):

3 x 128 x 768 x 768 x 2 = 452,984,832 FLOPs

These are the three matrix multiplications that project the input into query, key, and value spaces. The factor of 2 accounts for the multiply-accumulate operations.

Masked multi-head self-attention (for 12 heads):

• Q * K^T: 2 x 12 x 128 x 64 x 128 = 25,165,824 FLOPs
• Scores * V: 2 x 12 x 128 x 128 x 64 = 25,165,824 FLOPs
• Wo projection: 2 x 128 x 768 x 768 = 150,994,944 FLOPs

Total MHA: 201,326,592 FLOPs

Note: the causal mask zeroes out future positions but does not reduce the FLOP count in practice, because the full matrix multiplications are computed and then masked.

Feed-forward network FLOPs

128 x 768 x 3072 x 2 x 2 + 128 = 1,207,959,552 FLOPs

The two factors of 2 account for the two linear layers (up-projection and down-projection) and the multiply-accumulate operations. The FFN dominates the compute; it accounts for roughly 65% of the total FLOPs per block.

Total for 12 transformer blocks

Multiply by 12 blocks for the full model forward pass, which gives roughly 22.3 billion FLOPs. This matches BERT Base, which has the same layer configuration and dimensions. The causal mask changes what the model can see, but not the compute cost.

Comparison with BERT

BERT Base has the same layer configuration (12 layers, 768 embedding, 12 heads, 3072 FFN), so the per-block FLOP counts are nearly identical. The difference is architectural rather than computational. BERT uses full bidirectional attention with no mask and processes the whole sequence at once for understanding, while GPT-2 uses causal masked attention and generates one token at a time. The mask barely changes the compute, because it is only a comparison and assignment on the attention scores before softmax.

The bigger picture: from GPT-2 to GPT-3 to GPT-4

GPT-2's architecture matters because later models reuse it. The decoder-only, autoregressive design with the causal mask scales to much larger sizes.

GPT-3: few-shot learning at scale

GPT-3 took the same GPT-2 architecture and scaled it to 175 billion parameters, 10x more than any previous non-sparse language model. At this scale the model does few-shot learning. It can generate high-quality text for translation, question answering, and even programming tasks with little or no task-specific training data.

GPT-3 also showed that scaling up language models greatly improves task-agnostic, few-shot performance, sometimes reaching competitiveness with prior state-of-the-art fine-tuning approaches. It can even do zero-shot learning, generating quality output for tasks it was never explicitly trained on.

GPT-4: multimodal and human-level

GPT-4 extended the decoder-only lineage into a large-scale multimodal model that accepts both image and text inputs. It shows human-level performance on various professional and academic benchmarks, including passing a simulated bar exam around the median of human test-takers. GPT-4 is also highly steerable. Users can instruct how they want the model to respond, which improves prompt engineering.

Attention types: forward, causal, and triangle

In my independent study, I also surveyed the attention patterns that underpin these models:

• Forward (self-) attention: each query attends to all keys and values. Used in BERT. Bidirectional.
• Causal attention: each query can attend only to keys and values at or before its position. Used in GPT-2/3/4. The triangular mask we discussed.
• Triangle attention: each query attends to a subset of keys and values based on a maximum distance, forming a triangular pattern. Used to reduce computation while capturing long-range dependencies.

This section draws from my CS199 Supervised Independent Study at UC Berkeley.

Key takeaways

The main points:

• GPT-2 is decoder-only and autoregressive. It generates one token at a time and always looks backward. This choice separates it from encoder models like BERT.
• The causal mask enforces unidirectionality. A lower-triangular mask with negative infinity in the upper triangle zeros out attention to future tokens after softmax.
• The FFN dominates compute. In GPT-2 Small, the feed-forward network accounts for roughly 65% of the FLOPs per transformer block, with the attention projections and computation making up the rest.
• Residual connections and layer normalization keep training stable. They make it possible to train a 12-layer transformer without vanishing gradients and with stable convergence.
• GELU is the activation of choice. Its smooth gradient profile outperforms ReLU in transformer architectures.
• Two-stage training (pre-train then fine-tune) is what makes these models effective. Pre-training on massive text gives broad language understanding; fine-tuning specializes it.

The GPT-2 architecture is a stack of attention and feed-forward blocks with residual connections and normalization. What makes it work is the scale of pre-training, the causal mask that enables generation, and the choice of components (GELU, layer norm, residual paths) that keep training stable at depth.

Based on my CS199 Supervised Independent Study at UC Berkeley and the presentations I created in 2023.