Building and Training Your Own LLM
Implementing a Transformer Model from Scratch (Using PyTorch or TensorFlow)
This guide covers the core architectural implementation of the Transformer model using PyTorch, focusing on the mechanics of self-attention and feed-forward networks. Understanding this structure is essential for anyone aiming to move beyond high-level APIs to achieve true mastery of generative language modeling. Reach for this knowledge when you need to customize model behavior, optimize inference latency, or debug internal weight representations.
The Core Concept: Multi-Head Self-Attention
The self-attention mechanism is the heart of the Transformer, enabling the model to relate different positions of a single sequence to compute a representation of that sequence. We implement this by projecting the input into Query, Key, and Value matrices. The intuition is that the Query represents what a token is looking for, the Key represents what a token contains, and the Value represents the information to be extracted. By taking the dot product of Queries and Keys, we obtain attention scores that indicate how much focus one token should place on others. We scale these scores by the square root of the head dimension to prevent gradients from vanishing during training. Finally, applying the Softmax function ensures these scores sum to one, effectively creating a weighted sum of the Value vectors. This mechanism allows the model to process relationships between distant tokens in constant path lengths, overcoming the limitations of sequential architectures.
import torch
import torch.nn as nn
import torch.nn.functional as F
class SelfAttention(nn.Module):
def __init__(self, embed_dim, num_heads):
super().__init__()
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
self.qkv = nn.Linear(embed_dim, 3 * embed_dim) # Combined Q, K, V projection
self.out = nn.Linear(embed_dim, embed_dim)
def forward(self, x):
B, T, C = x.shape
qkv = self.qkv(x)
q, k, v = qkv.chunk(3, dim=-1)
# Reshape for multi-head splitting: (B, T, H, D) -> (B, H, T, D)
q, k, v = [x.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) for x in (q, k, v)]
# Scaled dot-product attention
attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
attn = F.softmax(attn, dim=-1)
out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)
return self.out(out)Integrating Positional Encoding
Since Transformers lack a recurrent or convolutional structure, they are inherently permutation-invariant, meaning they treat the input as a bag of tokens regardless of order. To introduce the concept of sequence order, we inject positional encodings into the input embeddings. These encodings carry information about the relative or absolute position of tokens in the sequence. While learned embeddings are popular, sinusoidal functions allow the model to potentially extrapolate to sequence lengths longer than those encountered during training. By adding these vectors element-wise to the word embeddings, the model gains the ability to distinguish between 'The dog bit the man' and 'The man bit the dog.' This additive property is powerful because it preserves the semantic information of the tokens while layering the geometric information of their positions, enabling the attention mechanism to attend to tokens based on both their content and their specific sequence context.
import math
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
super().__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
self.register_buffer('pe', pe.unsqueeze(0))
def forward(self, x):
# Add positional info to token embeddings
return x + self.pe[:, :x.size(1)]Constructing the Feed-Forward Network
After the attention mechanism has processed the dependencies between tokens, we apply a position-wise Feed-Forward Network (FFN) to each token independently. This component serves as the primary processing unit that refines the representation of each token based on the aggregated information from its context. A standard FFN consists of two linear transformations with a non-linear activation function—such as GELU or ReLU—sandwiched in between. The first transformation typically expands the dimensionality of the embedding space by a factor of four, allowing the model to perform complex non-linear feature transformations, while the second transformation maps it back to the original dimension. This depth provides the model with significant capacity to learn linguistic patterns and internal facts. By applying this to every token position equally, the Transformer maintains consistency across the sequence while the attention layer manages the cross-token interactions that define global context.
class FeedForward(nn.Module):
def __init__(self, d_model, dim_feedforward=2048):
super().__init__()
self.linear1 = nn.Linear(d_model, dim_feedforward)
self.dropout = nn.Dropout(0.1)
self.linear2 = nn.Linear(dim_feedforward, d_model)
self.act = nn.GELU()
def forward(self, x):
# Expand, activate, and project back to d_model
return self.linear2(self.dropout(self.act(self.linear1(x))))Assembling the Transformer Block
A single Transformer block encapsulates the two previous components within a structure that includes residual connections and layer normalization. Residual connections facilitate gradient flow through deep networks by allowing the model to learn identity mappings, effectively skipping layers that do not contribute to performance. Layer normalization is applied to stabilize the internal hidden states, ensuring that inputs to each sub-layer maintain a consistent distribution, which is crucial for efficient training. By stacking these blocks sequentially, the model learns increasingly abstract representations of language. In an autoregressive setting, we organize these layers to ensure the model only attends to past tokens during training. The modularity of this design allows researchers to scale the model size simply by adding more blocks or increasing the hidden dimensionality, making it a highly scalable approach for building large-scale generative models from scratch while maintaining robust performance metrics.
class TransformerBlock(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
self.attn = SelfAttention(d_model, num_heads)
self.ln1 = nn.LayerNorm(d_model)
self.ln2 = nn.LayerNorm(d_model)
self.ffn = FeedForward(d_model)
def forward(self, x):
# Add and Norm pattern
x = x + self.attn(self.ln1(x))
x = x + self.ffn(self.ln2(x))
return xTraining Loops and Causal Masking
To train a language model that generates coherent text, we must enforce causality so that a token at position 't' cannot 'see' future tokens at position 't+1'. We achieve this by applying a triangular mask to the attention scores before the softmax operation, setting the values corresponding to future tokens to negative infinity. During the training loop, we feed the model sequences of tokens, compute the loss using Cross-Entropy, and perform backpropagation to update weights. We use a sliding window approach where the input is shifted by one token to create the target labels. Monitoring the training loss and validation perplexity is critical to identifying underfitting or overfitting early. We must carefully choose an optimizer like AdamW to handle weight decay and maintain stable updates across the vast parameter space of a Transformer, ensuring the model effectively learns to predict the next token in the sequence.
def train_step(model, inputs, targets):
# Masking for autoregressive training
T = inputs.size(1)
mask = torch.tril(torch.ones(T, T)).view(1, 1, T, T)
logits = model(inputs)
# Calculate loss on flattened sequences
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
loss.backward()
return loss.item()Key points
- The self-attention mechanism relies on Query, Key, and Value projections to calculate token relevance.
- Positional encoding is required because Transformers do not have an inherent concept of order.
- Residual connections help mitigate the vanishing gradient problem in deep Transformer networks.
- Layer normalization ensures stable activation distributions across training steps.
- Feed-forward networks apply non-linear transformations to individual token representations independently.
- Causal masking prevents the model from peeking at future tokens during autoregressive training.
- The model architecture is modular, allowing for easy scaling by stacking identical transformer blocks.
- Cross-Entropy loss is the standard objective function for predicting the next token in a sequence.
Common mistakes
- Mistake: Forgetting to apply the mask in the self-attention mechanism. Why it's wrong: Without a causal mask, the model can look ahead at future tokens during training, leading to data leakage. Fix: Implement an upper-triangular matrix set to negative infinity to mask future tokens.
- Mistake: Neglecting to scale the dot product by the square root of the head dimension. Why it's wrong: Large values in the dot product push the Softmax function into regions with extremely small gradients, causing vanishing gradients. Fix: Divide the query-key dot product by sqrt(d_k).
- Mistake: Applying Layer Normalization after the residual connection rather than before. Why it's wrong: Pre-norm architectures are generally more stable and easier to train for deep transformer stacks compared to Post-norm. Fix: Place LayerNorm inside the residual block so the main path remains clean.
- Mistake: Using a fixed positional encoding that is not broadcastable. Why it's wrong: Positional embeddings must be added element-wise to the input embeddings; improper shapes cause broadcasting errors or incorrect alignment. Fix: Ensure positional embedding tensor shape matches input (batch, seq_len, embed_dim).
- Mistake: Improper initialization of linear layers. Why it's wrong: Poor initialization can lead to exploding or vanishing activations, especially in deep models. Fix: Use Xavier or Kaiming initialization tailored to the specific activation function used in the feed-forward network.
Interview questions
What is the fundamental role of the Scaled Dot-Product Attention mechanism in a Transformer model?
The Scaled Dot-Product Attention mechanism is the heart of the Transformer, enabling the model to weigh the importance of different words in a sequence relative to one another. We compute attention scores by taking the dot product of the Query and Key matrices, dividing by the square root of the dimension to stabilize gradients, and applying a softmax function. This creates a distribution that defines how much focus each token places on others. Mathematically, it is defined as Attention(Q, K, V) = softmax((QK^T) / sqrt(d_k))V. Without this, the model would treat every word in a sequence with equal importance, failing to capture the complex dependencies and contextual relationships necessary for generating coherent, human-like text in a custom LLM.
Why do we include residual connections and layer normalization in every sub-layer of the Transformer block?
Residual connections, or skip connections, are critical because they allow gradients to flow through the deep network during backpropagation without suffering from the vanishing gradient problem. By adding the input of a layer directly to its output, represented as x + Sublayer(x), we ensure that the model retains information from earlier layers. We pair this with Layer Normalization, which stabilizes the hidden state distributions across the batch. In a custom LLM, this architecture is essential because it allows us to stack many encoder or decoder blocks, enabling the model to learn increasingly abstract hierarchical features while maintaining training stability and faster convergence speeds across billions of parameters.
Explain the purpose of Multi-Head Attention versus using a single attention head.
Multi-Head Attention allows the model to jointly attend to information from different representation subspaces at different positions. If we only used one attention head, the averaging process would essentially collapse these diverse relationships into a single weighted representation, losing nuance. By using multiple heads, we can simultaneously focus on different linguistic structures, such as one head tracking syntactic dependencies, another tracking coreference resolution, and another tracking semantic associations. In code, this involves splitting the embedding dimension into smaller chunks, applying independent projections for each head, and concatenating the results. This parallel processing provides a richer, multi-faceted understanding of input sequences which is vital for high-performance generation.
How does the implementation of Positional Encodings work, and why is it necessary in a Transformer?
Unlike recurrent networks that process sequences sequentially, the Transformer processes all tokens in parallel, which means it has no inherent sense of word order or sequence position. Without Positional Encodings, the model would perceive the sentence 'the dog bit the man' exactly the same as 'the man bit the dog.' We solve this by adding a unique signal to the input embeddings. This is typically done using sine and cosine functions of varying frequencies, or via learned embedding vectors. By adding these values to the input, the model can distinguish between tokens based on their location, which is a fundamental requirement for maintaining the structure and grammar of generated sequences.
Compare using a causal mask in a decoder-only architecture versus a padding mask in an encoder-only architecture.
In a decoder-only architecture, such as a generative LLM, we must use a causal mask (also called a look-ahead mask) to prevent the model from 'cheating' by seeing future tokens during training. We implement this by applying a triangular matrix of negative infinity values to the attention scores before the softmax, ensuring the model only attends to current and previous positions. Conversely, a padding mask is used in encoders to ignore specific tokens, usually padding tokens added to make batch sequences uniform in length. While causal masks are about preserving the autoregressive nature of generation by enforcing temporal order, padding masks are about computational efficiency and data integrity, ensuring that null tokens do not contribute to the hidden state representation.
When implementing the Feed-Forward Network within each Transformer block, why is the hidden dimension usually four times larger than the embedding dimension?
The Feed-Forward Network typically consists of two linear transformations with a non-linear activation function, like GELU or ReLU, in between. Increasing the dimension to 4x the embedding size is a heuristic choice that provides sufficient 'workspace' for the model to perform non-linear transformations on the projected features. This expansion creates a higher-dimensional space where the model can separate complex patterns and perform sophisticated feature extraction before projecting back down to the model's standard embedding dimension. In a custom LLM, this projection is the primary site of knowledge storage and processing; if the hidden dimension is too small, the model becomes 'under-parameterized,' failing to capture the vast, nuanced internal logic required for high-quality language generation during inference.
Check yourself
1. Why do we divide the dot product of queries and keys by the square root of the head dimension (d_k)?
- A.To normalize the output vectors to unit length
- B.To prevent large dot product values from resulting in very small gradients during backpropagation
- C.To decrease the computational complexity of the attention matrix multiplication
- D.To ensure the query and key vectors have the same average magnitude as the values
Show answer
B. To prevent large dot product values from resulting in very small gradients during backpropagation
Scaling by sqrt(d_k) prevents the dot product from growing too large in magnitude, which would push the softmax function into its saturation region where gradients are near zero. Other options suggest normalization or complexity reduction, which are not the primary purpose of this scaling factor.
2. In a decoder-only transformer, why is the causal mask (look-ahead mask) essential during training?
- A.To force the model to learn bidirectional representations
- B.To improve the speed of the parallel computation during training
- C.To prevent the model from 'seeing' future tokens when predicting the current token
- D.To reduce the number of parameters in the attention mechanism
Show answer
C. To prevent the model from 'seeing' future tokens when predicting the current token
The causal mask ensures that the prediction for a position depends only on known past positions. Without it, the model would cheat by looking at future tokens. Birectional representation is used in encoders, not decoders, and masking does not improve speed or reduce parameters.
3. What is the primary function of the Feed-Forward Network (FFN) block in each transformer layer?
- A.To process each token position independently and introduce non-linearity
- B.To perform global communication between different token positions
- C.To reduce the dimensionality of the embedding space
- D.To store long-term memory about the training dataset
Show answer
A. To process each token position independently and introduce non-linearity
The FFN is applied position-wise, allowing the model to project features into a higher-dimensional space and apply non-linear activations. It does not perform communication between positions (that is the role of attention) and does not store the dataset.
4. When implementing multi-head attention, how does increasing the number of heads affect the model?
- A.It forces the model to ignore long-range dependencies
- B.It allows the model to attend to information from different representation subspaces simultaneously
- C.It increases the total number of parameters linearly with the number of heads
- D.It reduces the dimensionality of the hidden states in each head
Show answer
B. It allows the model to attend to information from different representation subspaces simultaneously
Multi-head attention enables the model to focus on different aspects of the input sequence in parallel. While the dimension per head decreases, the total parameter count remains roughly the same because the total hidden dimension is preserved. It does not ignore dependencies.
5. Why is the Residual Connection (Add & Norm) critical in a deep transformer?
- A.It acts as a substitute for the attention mechanism
- B.It allows gradients to propagate through the network more easily, preventing vanishing gradients
- C.It replaces the need for positional encodings
- D.It reduces the size of the attention matrix to save memory
Show answer
B. It allows gradients to propagate through the network more easily, preventing vanishing gradients
Residual connections provide a 'highway' for gradients to flow backward during training, which is essential for training very deep models. They are not substitutes for attention or embeddings, and they do not influence the size of the attention matrix.