Pytorch implementation of paper: GPT-1 (2018):”Improving Language Understanding by Generative Pre-Training”

Written by

in

The 2018 paper “Improving Language Understanding by Generative Pre-Training” (GPT-1) by Radford et al. (OpenAI) introduced the concept of Generative Pre-Training. While BERT (released around the same time) was a “bidirectional” encoder, GPT-1 was a unidirectional (left-to-right) decoder.

GPT-1 proved that a model pre-trained on a massive corpus of unlabeled text to predict the next word could be “fine-tuned” on specific tasks (like classification or entailment) with very little data.

1. The Causal Multi-Head Attention

The critical difference between GPT-1 and the standard Transformer encoder is the Causal Mask. To ensure the model only looks at the past when predicting the next word, we mask all “future” tokens by setting their attention scores to $-\infty$.

Python

import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads, max_len):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        
        self.w_qkv = nn.Linear(d_model, 3 * d_model)
        self.fc = nn.Linear(d_model, d_model)
        
        # Causal mask: lower triangular matrix of ones
        self.register_buffer("mask", torch.tril(torch.ones(max_len, max_len))
                                     .view(1, 1, max_len, max_len))

    def forward(self, x):
        batch_size, seq_len, d_model = x.size()
        
        # Generate Queries, Keys, and Values
        q, k, v = self.w_qkv(x).chunk(3, dim=-1)
        
        # Reshape for multi-head: [batch, heads, seq_len, d_k]
        q = q.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        
        # Scaled Dot-Product Attention
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5)
        
        # Apply Causal Mask: prevent attending to future tokens
        scores = scores.masked_fill(self.mask[:, :, :seq_len, :seq_len] == 0, float('-inf'))
        
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        
        # Re-assemble heads
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        return self.fc(out)


2. The GPT-1 Decoder Block

GPT-1 uses the “Transformer Decoder” but removes the Encoder-Decoder attention layer (since there is no encoder). It consists only of Masked Self-Attention and a Feed-Forward Network.

Python

class GPTBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, max_len, dropout=0.1):
        super().__init__()
        self.ln_1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads, max_len)
        self.ln_2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(), # GPT-1 used GELU
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        # Pre-LayerNorm (standard in GPT models)
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


3. Training Objective: Next Token Prediction

During pre-training, GPT-1 minimizes the negative log-likelihood of the next token given the previous ones:

$$L_1(\mathcal{U}) = \sum_i \log P(u_i | u_{i-k}, \dots, u_{i-1}; \Theta)$$

For fine-tuning, GPT-1 adds a task-specific linear layer but continues to use the language modeling objective as an auxiliary loss to improve generalization.

Python

class GPT1(nn.Module):
    def __init__(self, vocab_size, d_model, n_layers, n_heads, d_ff, max_len):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_len, d_model)
        
        self.blocks = nn.ModuleList([
            GPTBlock(d_model, n_heads, d_ff, max_len) for _ in range(n_layers)
        ])
        
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, x):
        batch_size, seq_len = x.size()
        pos = torch.arange(seq_len).unsqueeze(0).to(x.device)
        
        x = self.token_emb(x) + self.pos_emb(pos)
        
        for block in self.blocks:
            x = block(x)
            
        x = self.ln_f(x)
        return self.head(x) # Logits for the entire vocab


Key Technical Details

  • Byte-Pair Encoding (BPE): GPT-1 used a 40,000-merge BPE vocabulary.
  • GELU Activation: Like BERT, GPT-1 used the Gaussian Error Linear Unit for smoother gradients.
  • Task-Specific Input Transformations: To handle different tasks (Multiple Choice, Entailment, etc.) with the same decoder, GPT-1 introduced special delimiters like [Extract], [Start], and [Delim].

Why it was a breakthrough

Before GPT-1, most models were trained from scratch for each specific task. GPT-1 popularized the “Pre-train then Fine-tune” paradigm, showing that the knowledge gained by simply reading a large library (BooksCorpus) could be transferred to help a model understand logic, sentiment, and grammar.

Comments

Leave a Reply

Your email address will not be published. Required fields are marked *