12 Training A Transformer to Generate Text

This chapter walks through building and training a GPT-style Transformer from scratch in PyTorch, focusing on a practical, lightweight model inspired by GPT-2 that can be trained on a regular GPU. Beyond showing how to generate text, it explains why training (or fine-tuning) your own model matters: tailoring behavior to a domain, improving relevance, and retaining control over data and privacy. The workflow spans designing a scaled-down architecture, preparing data, training with appropriate safeguards against overfitting, and deploying the model for autoregressive text generation.

The model mirrors GPT-2’s decoder-only design but is deliberately small: about 5.12 million parameters with 3 decoder blocks, 256-dimensional embeddings, 4 attention heads, learned positional embeddings, GELU activations, layer normalization, residual connections, and dropout. Training data is curated from three Ernest Hemingway novels to balance stylistic coverage and computational cost. The text is cleaned (including quote normalization), lowercased, and tokenized at the word level with punctuation separated, yielding a vocabulary of roughly 10,600 tokens plus an UNK token for robustness. Text is converted to integer sequences, batched into 128-token contexts, and trained via next-token prediction by shifting targets one step to the right.

Training uses Adam with cross-entropy loss, gradient norm clipping, checkpointing every 10 epochs up to 40, and GPU acceleration where available. After training, text is generated autoregressively from a prompt using temperature and optional top-k sampling. Results illustrate the trade-off between coherence and memorization: the 20-epoch model produces mostly coherent output with fewer verbatim passages, while the 40-epoch model is more fluent but more likely to copy training text, especially under unconditional prompts. The chapter emphasizes selecting suitable training length or using validation to curb overfitting, and it equips readers with an end-to-end recipe—from data choice and tokenization to architecture, training, and controlled text generation.

The architecture of a decoder-only Transformer, designed to generate text. The text from three Hemingway novels is tokenized and then converted into indexes. We arrange 128 indexes into a sequence, and each batch contains 32 such sequences. The input first undergoes word embedding and positional encoding, with the input embedding being the sum of these two components. This input embedding is then processed through three decoder layers. Following this, the output undergoes layer normalization and passes through a linear layer, resulting in an output size of 10,600, which corresponds to the number of unique tokens in the vocabulary.

The training process for a decoder-only Transformer to generate text, Hemingway-style.

• The style of the generated text from a GPT model will be heavily influenced by the training data. For effective text generation, it's important to have a balance of text length and variation in the training material. The training dataset should be sufficiently large for the model to learn and emulate a specific writing style accurately. However, if the dataset lacks diversity, the model might end up reproducing passages directly from the training text. Conversely, overly long training datasets can require excessive computational resources for training.
• Choosing the right hyperparameters in the GPT model is crucial for successful model training and text generation. Setting the hyperparameters too large may lead to too many parameters. This results in longer training time and an overfitted model. Setting the hyperparameters too small may hinder the model’s ability to learn effectively and capture the writing style in the training data. This may lead to incoherent generated text.
• The appropriate number of training epochs is important for text generation. While training too few epochs may lead to incoherent text, training for too many epochs may lead to an overfitted model that generates text identical to passages in the training text.

FAQ