LLM Fundamentals
Transformer Architecture
The Transformer architecture is a deep learning model based on the self-attention mechanism that replaces recurrence with parallelizable global dependencies. It serves as the foundational backbone for modern large language models, enabling them to process vast contexts with unprecedented efficiency and representational power. You reach for this architecture whenever you need to model long-range relationships in sequential data where training scalability and context retention are the primary requirements.
The Input Embedding Layer
The journey of a token through a Transformer begins with the embedding layer, which maps discrete vocabulary indices into a continuous vector space. Why does this matter? Neural networks operate on mathematical operations rather than raw text, and these embeddings capture semantic similarities; words that appear in similar contexts reside closer together in the vector space. We append positional encodings to these vectors because the Transformer processes tokens simultaneously rather than sequentially. Without positional info, the model would be permutation invariant, treating 'dog bites man' identical to 'man bites dog.' By adding sinusoidal functions or learned offsets to the embeddings, we inject a signal that tells the model the relative or absolute order of the tokens, allowing the model to distinguish structure and syntactic roles within the input sequence.
import torch
import torch.nn as nn
# Create a simple embedding layer with positional encoding placeholder
vocab_size = 1000
d_model = 64
embedding = nn.Embedding(vocab_size, d_model) # Converts indices to vectors
# Simulate input indices [batch_size, seq_len]
input_indices = torch.randint(0, vocab_size, (1, 10))
embedded = embedding(input_indices)
# Note: In practice, we add positional encoding here to the 'embedded' tensorScaled Dot-Product Attention
Self-attention is the mechanism that allows the model to weigh the importance of other words in the sequence when processing a specific token. We calculate three vectors for each token: Query (what I am looking for), Key (what I offer), and Value (what I contain). By taking the dot product of the Query of one token with the Keys of all others, we generate attention scores. We scale these scores by the square root of the dimension of the keys to prevent the gradients from becoming too small during backpropagation, a phenomenon common in high-dimensional dot products. This mechanism allows the model to dynamically construct context; for example, in 'the bank of the river,' the model attends heavily to 'river' when processing 'bank,' resolving ambiguity by leveraging global information rather than just the immediately preceding word.
import torch.nn.functional as F
# Scaled Dot-Product Attention function
def scaled_dot_product_attention(q, k, v):
d_k = q.size(-1)
# Compute scores and scale by sqrt of head dimension
scores = torch.matmul(q, k.transpose(-2, -1)) / (d_k**0.5)
# Apply softmax to get attention weights
attn = F.softmax(scores, dim=-1)
return torch.matmul(attn, v)
# Example usage with dummy tensors: (batch, heads, seq, head_dim)
q = k = v = torch.randn(1, 8, 10, 64)
output = scaled_dot_product_attention(q, k, v)Multi-Head Attention
Rather than relying on a single attention mechanism, Multi-Head Attention projects the input into multiple lower-dimensional subspaces. Why? A single head might focus on syntactic relationships, while another might focus on semantic associations or specific grammatical nuances. By running multiple heads in parallel, the model can attend to different facets of the information simultaneously, enriching the internal representation of each token. These heads operate independently and their outputs are concatenated and linearly projected back to the original dimension. This design increases the capacity of the model to capture complex, multi-layered dependencies. By distributing the computational burden across several heads, the architecture ensures that the model is not forced to compromise between different types of relational information during the representation learning process at each layer.
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
self.num_heads = num_heads
self.head_dim = d_model // num_heads
# Single linear layer to project input into Q, K, V for all heads
self.qkv_proj = nn.Linear(d_model, 3 * d_model)
self.out_proj = nn.Linear(d_model, d_model)
def forward(self, x):
# Logic: split input into heads, apply attention, concatenate, project
return self.out_proj(x) # Placeholder for concated attention outputThe Feed-Forward Network and Residual Connections
After the attention mechanism, each position passes through an identical position-wise feed-forward network. This network acts as a memory or knowledge store for the patterns identified by the attention layers. Crucially, Transformers use residual connections (or skip connections) around each sub-layer followed by layer normalization. Residual connections address the vanishing gradient problem, enabling information to flow unimpeded through very deep stacks of layers. Without these, the gradients would degrade, preventing effective training of deep architectures. By adding the input of a layer to its output, we force the network to learn only the 'residual' or the incremental transformation required, making the optimization process significantly more stable and efficient. This combination of attention for routing and feed-forward networks for processing allows for sophisticated transformations across depth.
class FeedForward(nn.Module):
def __init__(self, d_model):
super().__init__()
# Standard two-layer MLP with expansion factor
self.net = nn.Sequential(
nn.Linear(d_model, 4 * d_model),
nn.GELU(),
nn.Linear(4 * d_model, d_model)
)
def forward(self, x):
# Add residual connection: x + sublayer(x)
return x + self.net(x)The Transformer Decoder Stack
The decoder stack mirrors the encoder but includes masked attention and cross-attention. Masked self-attention ensures that each token can only attend to preceding tokens, preserving the autoregressive property necessary for text generation. Without this mask, the model would 'see' the answers during training. Cross-attention then allows the decoder to attend to the processed encoder output, acting as a bridge between the input representation and the target generation sequence. The final output passes through a linear layer and a softmax to produce a probability distribution over the vocabulary. This systematic approach—combining input encoding, masked generation, and cross-token attention—is what allows modern language models to translate, summarize, and generate coherent text by predicting the next logical token in a sequence with high probability.
class DecoderLayer(nn.Module):
def __init__(self, d_model):
super().__init__()
self.masked_attn = nn.MultiheadAttention(d_model, 8)
self.norm = nn.LayerNorm(d_model)
def forward(self, x, memory):
# Masked self-attention for autoregressive generation
out = self.masked_attn(x, x, x, attn_mask=None)
return self.norm(x + out[0]) # Simplified skip connectionKey points
- Embeddings map discrete tokens into continuous vectors to allow for mathematical processing.
- Positional encodings are necessary because Transformers process all tokens in a sequence simultaneously.
- The self-attention mechanism computes weighted relationships between all pairs of tokens in a sequence.
- Scaled dot-product attention utilizes a scaling factor to maintain gradient stability during training.
- Multi-head attention enables the model to simultaneously attend to different semantic and syntactic information subspaces.
- Residual connections are essential for preventing gradient decay in deep neural network architectures.
- Masked attention in the decoder prevents the model from seeing future tokens during autoregressive training.
- The Transformer's ability to parallelize computations across sequence length makes it superior to recurrent architectures for long-range data.
Common mistakes
- Mistake: Thinking Attention is the only mechanism. Why it's wrong: Attention is a component, but Feed-Forward Networks and Positional Encodings are essential for processing. Fix: View Transformers as an integration of Attention, Add & Norm, and Position-wise networks.
- Mistake: Misinterpreting 'Self-Attention' as global memory. Why it's wrong: It is a dynamic weighting mechanism for token relationships within a context window, not a stored database. Fix: Remember it computes scores based on current input relationships, not learned past knowledge.
- Mistake: Confusing Encoder and Decoder roles. Why it's wrong: The Encoder produces rich contextual representations, while the Decoder uses those to generate sequence output with causal masking. Fix: Understand that causal masking prevents the model from seeing the 'future' during generation.
- Mistake: Believing Positional Encoding is optional. Why it's wrong: Transformers are permutation-invariant; without ordering information, the model sees a 'bag of words'. Fix: Always include sine/cosine or learned embeddings to inject sequence order.
- Mistake: Overestimating parameter efficiency. Why it's wrong: Large models require immense compute due to quadratic scaling of self-attention memory. Fix: Account for the memory complexity scaling as sequence length increases.
Interview questions
What is the primary innovation of the Transformer architecture compared to previous sequence modeling approaches?
The primary innovation of the Transformer is the self-attention mechanism, which allows the model to process all tokens in a sequence simultaneously rather than sequentially. Unlike recurrent architectures that process data step-by-step, creating bottlenecks and long-range dependency issues, the Transformer uses parallelization. By calculating global dependencies between every pair of words in a sentence instantly, it significantly speeds up training and enables the model to understand context across very long inputs.
How does the Multi-Head Attention mechanism function within a Transformer model?
Multi-Head Attention works by running multiple self-attention layers in parallel, which we call heads. Each head independently learns different types of relationships between tokens—for example, one head might focus on syntactic structure while another focuses on semantic associations. Specifically, the model projects input embeddings into multiple sets of Query, Key, and Value vectors. After computing the attention scores for each, these outputs are concatenated and linearly projected back into a single representation, allowing the model to capture diverse perspectives of the input sequence simultaneously.
Explain the role and importance of Positional Encodings in Transformer architectures.
Because the Transformer processes all input tokens in parallel rather than sequentially, it lacks an inherent sense of the order or position of words in a sequence. Without positional encoding, the model would treat the phrase 'the cat chased the dog' the same as 'the dog chased the cat.' We inject positional information by adding sine and cosine waves of different frequencies to the input embeddings. This allows the model to retain critical information about the relative and absolute positions of tokens, which is fundamental for understanding grammatical structure and context in generative tasks.
Compare and contrast the Encoder-only architecture with the Decoder-only architecture in the context of generative tasks.
An Encoder-only model is designed to build a deep, bidirectional understanding of input text, making it excellent for classification or feature extraction. Conversely, Decoder-only architectures are autoregressive; they are trained to predict the next token based solely on the preceding ones. While Encoders look at the entire context simultaneously to understand the full meaning, Decoders enforce a causal mask to ensure the model only looks at past tokens during generation, which is essential for creating coherent, sequential text in generative applications.
What is the purpose of the Feed-Forward Network and Layer Normalization within each Transformer block?
Within each Transformer block, the Feed-Forward Network applies two linear transformations with a non-linear activation like ReLU or GeLU to the output of the attention layer. The FFN acts as a memory or knowledge base where the model processes the features extracted by the attention mechanism. Layer Normalization is applied around these sub-layers with residual connections to stabilize training. By normalizing the activations, we prevent gradient exploding or vanishing, which ensures that deep networks remain trainable even as we add more layers for increased generative complexity.
Explain the mechanics and necessity of the Masked Multi-Head Attention in decoder-based models.
Masked Multi-Head Attention is the mechanism that prevents the model from cheating during training. In a generative task, we want the model to predict the next word without knowing what follows. We apply a causal mask—typically a lower triangular matrix of negative infinity values—to the attention scores before the softmax operation. This ensures that the dot product for any token i only considers positions j less than or equal to i. This forces the model to learn conditional probabilities which is the backbone of all autoregressive text generation.
Check yourself
1. What is the primary purpose of the multi-head attention mechanism?
- A.To compress the input data into a fixed-length vector
- B.To allow the model to attend to information from different representation subspaces
- C.To remove the need for backpropagation during training
- D.To eliminate the need for activation functions in the network
Show answer
B. To allow the model to attend to information from different representation subspaces
Multi-head attention allows the model to jointly attend to information at different positions from different representation subspaces. Option 0 describes pooling, Option 2 is incorrect as training requires gradients, and Option 3 is false because non-linearity is essential.
2. Why is a causal mask applied in the Decoder portion of a Transformer?
- A.To reduce the number of calculations in the Feed-Forward layer
- B.To prevent the model from attending to future tokens during training
- C.To force the model to focus only on the first token
- D.To simplify the calculation of gradients in the Encoder
Show answer
B. To prevent the model from attending to future tokens during training
Causal masking ensures that when predicting a token at position i, the model only has access to tokens at positions less than i. This prevents 'cheating' by looking at upcoming tokens. The other options do not describe the function of causal masks.
3. If you double the input sequence length in a standard Transformer, how does the self-attention memory complexity change?
- A.It remains constant
- B.It doubles linearly
- C.It increases quadratically
- D.It decreases by half
Show answer
C. It increases quadratically
Standard self-attention computes an N x N attention matrix for N tokens, resulting in O(N^2) complexity. It does not stay constant, scale linearly, or decrease.
4. What role do Residual Connections play in deep Transformer architectures?
- A.They replace the need for normalization layers
- B.They facilitate gradient flow during training by preventing vanishing gradients
- C.They automatically determine the optimal sequence length
- D.They increase the memory usage for each layer
Show answer
B. They facilitate gradient flow during training by preventing vanishing gradients
Residual (skip) connections allow gradients to bypass transformations, which is critical for training deep networks. They do not replace normalization, determine sequence length, or primarily exist to increase memory usage.
5. How does a Transformer handle the lack of inherent recurrence or convolution?
- A.By using Positional Encodings to inject token order information
- B.By processing tokens one by one in a loop
- C.By using a global average pooling layer on the entire input
- D.By discarding the input ordering completely
Show answer
A. By using Positional Encodings to inject token order information
Because Transformers process the entire sequence in parallel, they need Positional Encodings to know the relative order of tokens. Option 1 describes RNNs, Option 2 is ineffective, and Option 3 would lose the semantic structure.