Interview Prep
Explain the transformer architecture and its key components.
The transformer architecture is a deep learning framework designed to process sequential data using global attention mechanisms instead of recurrence. It enables massive parallelization and long-range dependency modeling, making it the fundamental building block of modern large language models. You reach for this architecture whenever you need to capture complex contextual relationships across long sequences of data efficiently.
The Core Philosophy of Self-Attention
The brilliance of the transformer lies in replacing sequential processing with the self-attention mechanism, which allows every token in a sequence to look at every other token simultaneously. In traditional models, information flow is bottlenecked by the need to pass a hidden state through time, leading to vanishing gradients and forgotten history. By using self-attention, the model assigns importance weights to different parts of the input sequence relative to the target token. This allows the network to learn structural relationships regardless of distance. Mathematically, this is done by computing dot products between query, key, and value vectors derived from the input. This interaction creates a dynamic mapping where the importance of information is determined by relevance rather than temporal proximity, solving the limitation of fixed-window attention or recurrent dependencies.
import torch
import torch.nn.functional as F
# Simulating self-attention weights
def simple_attention(q, k, v):
# q, k, v are vectors of dimension d_k
# Scaled dot-product attention computes relevance
scores = torch.matmul(q, k.transpose(-2, -1)) / (q.size(-1)**0.5)
weights = F.softmax(scores, dim=-1)
return torch.matmul(weights, v)
# Example vectors for demonstration
q = torch.randn(1, 4, 8) # batch, seq_len, head_dim
k, v = torch.randn(1, 4, 8), torch.randn(1, 4, 8)
print(simple_attention(q, k, v).shape)Multi-Head Attention: Capturing Multiple Perspectives
If self-attention is the ability to look at the whole sequence, multi-head attention is the ability to look at it from multiple perspectives at once. By splitting the embedding space into several independent 'heads', the model can dedicate different heads to different linguistic or semantic tasks—for instance, one head might track syntactic subject-verb agreement, while another identifies co-reference resolution between pronouns and nouns. These heads run in parallel and their outputs are concatenated and projected back to the original dimensionality. This structure allows the model to build a multifaceted understanding of the input without increasing the computational complexity significantly. The reasoning behind this is that a single attention map is too limiting to capture the diverse range of dependencies present in natural language, requiring a distributed representation of context to avoid information loss.
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, embed_dim, num_heads):
super().__init__()
self.mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
def forward(self, x):
# x: input tensor of shape (batch, seq, embed)
# Using standard multi-head implementation
out, _ = self.mha(x, x, x)
return out
# Initialize MHA with 512 dimensions and 8 heads
mha_layer = MultiHeadAttention(512, 8)Positional Encoding: Restoring Sequential Order
Since transformers process all input tokens simultaneously, they have no intrinsic notion of order or time. This is a potential weakness because, in language, the order of words is critical to meaning. To solve this, transformers inject positional encodings into the input embeddings. These encodings are often sinusoidal functions or learned vectors added directly to the input representations. By summing the token embedding with a positional signal, the model gains the ability to distinguish between a word at the beginning of a sentence and the same word at the end. The reasoning for using sine and cosine functions is that they provide a stable, bounded, and continuous signal that allows the model to potentially extrapolate to sequence lengths longer than those seen during training, maintaining a consistent relative positioning between tokens regardless of their absolute index.
import torch
import math
# Generate sinusoidal positional encoding
def get_pos_encoding(seq_len, d_model):
pe = torch.zeros(seq_len, d_model)
position = torch.arange(0, seq_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
return pe
print(get_pos_encoding(10, 512).shape)Feed-Forward Networks and Layer Normalization
After the attention mechanism extracts contextual information, the data passes through a position-wise Feed-Forward Network (FFN). This FFN consists of two linear transformations with a non-linear activation function in between, typically GELU or ReLU. The role of this component is to perform localized processing on each token's representation, allowing the model to project the attended information into a richer feature space. Layer normalization is applied either before or after these blocks to stabilize training by keeping the mean and variance of the activations within a manageable range. Without these normalization steps, deep stacks of layers would suffer from exploding or vanishing activation values, making convergence impossible. The FFN acts as the 'knowledge' component of the layer, where the network refines the abstracted context provided by the attention mechanism into more refined semantic features.
import torch.nn as nn
class FeedForward(nn.Module):
def __init__(self, d_model, d_ff):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Linear(d_ff, d_model)
)
def forward(self, x):
return self.layers(x)
# Standard FFN block with expansion ratio
ffn = FeedForward(512, 2048)Residual Connections: Enabling Depth
Deep networks are notoriously difficult to train because signals tend to degrade as they pass through many layers. Transformers utilize residual connections—also known as skip connections—to circumvent this issue. By adding the input of a layer to its output, the model creates a 'highway' for gradients to flow directly through the network during backpropagation. This simple architectural tweak is the primary reason why we can train transformers with hundreds of layers instead of just a handful. The residual connection ensures that each block in the transformer only needs to learn the 'residual' or the incremental improvement upon the previous state, rather than attempting to learn a completely new representation from scratch at every step. This makes the optimization landscape much smoother, which is essential for successfully training modern large-scale language models without encountering catastrophic training failures.
import torch.nn as nn
class ResidualBlock(nn.Module):
def __init__(self, d_model):
super().__init__()
self.norm = nn.LayerNorm(d_model)
self.linear = nn.Linear(d_model, d_model)
def forward(self, x):
# The core logic: x + Sublayer(x)
return x + self.linear(self.norm(x))
block = ResidualBlock(512)
# Apply block to input
output = block(torch.randn(1, 10, 512))Key points
- Self-attention allows the model to weight the importance of different tokens in a sequence globally.
- Multi-head attention enables the model to focus on diverse contextual relationships simultaneously.
- Positional encodings provide the necessary structural order that is missing in parallel processing.
- Residual connections facilitate the training of extremely deep stacks by preventing gradient degradation.
- Feed-forward networks apply localized, non-linear transformations to refine token representations.
- Layer normalization is critical for maintaining stable activation distributions across many layers.
- Transformers replace recurrence with attention, allowing for high-throughput parallel training.
- The combination of these components allows for the modeling of long-range dependencies in complex datasets.
Common mistakes
- Mistake: Confusing the Encoder with the Decoder in the original architecture. Why it's wrong: The Encoder processes input sequences while the Decoder generates them autoregressively. Fix: Remember the Encoder is for bidirectional context, the Decoder is for causal generation.
- Mistake: Thinking Attention mechanisms are a replacement for Feed-Forward Networks. Why it's wrong: Attention handles contextual relationships, but Feed-Forward layers are needed for projecting and processing information per token. Fix: View them as complementary components working in series.
- Mistake: Assuming Positional Encoding is learned for every position independently. Why it's wrong: Fixed sinusoidal encodings use periodic functions to allow the model to generalize to sequence lengths it hasn't seen. Fix: Distinguish between static sinusoidal embeddings and learned positional embeddings.
- Mistake: Believing Layer Normalization can be omitted in a Transformer. Why it's wrong: Transformers are highly sensitive to internal covariate shift due to deep stacks of residual connections; normalization is critical for stability. Fix: Always include LayerNorm inside the residual blocks.
- Mistake: Thinking Masked Multi-Head Attention is used in the Encoder. Why it's wrong: The Encoder must see the whole input at once, while the Decoder's masking is required to prevent looking at future tokens during training. Fix: Only apply causal masking to the Decoder block.
Interview questions
At a high level, what is the core purpose of the Transformer architecture when building our own LLM?
The Transformer architecture is designed to process sequential data in parallel rather than sequentially, which solves the bottleneck inherent in older architectures. By using self-attention mechanisms, the model can capture global dependencies between all words in a sequence simultaneously, regardless of their distance. This allows our LLM to understand context and relationships within long text sequences far more efficiently, drastically accelerating the training process on large datasets while improving the quality of the learned representations.
What is the role of positional encoding in a Transformer, and why can't we omit it?
Because the Transformer architecture processes all tokens in a sequence at the same time using parallel self-attention, it inherently lacks any concept of word order. Without positional encoding, the model would treat the sequence 'dog bites man' exactly the same as 'man bites dog'. We inject positional encodings—usually via sine and cosine functions or learned embeddings—into the input embeddings to provide the model with essential information about the relative or absolute position of tokens, allowing the model to distinguish structure.
Can you explain how the multi-head self-attention mechanism works?
Multi-head self-attention allows our LLM to jointly attend to information from different representation subspaces at different positions. Instead of computing one single attention function, we project the Queries, Keys, and Values multiple times with different, learned linear projections in parallel. These 'heads' are then concatenated and projected again. This is crucial because it enables the model to focus on various aspects of language simultaneously, such as one head focusing on syntactic relationships while another focuses on semantic or coreferential links.
Why do we use Layer Normalization and Residual Connections in the Transformer blocks?
Residual connections, defined as Output = Layer(Input) + Input, are vital because they allow gradients to flow directly through the network during backpropagation, mitigating the vanishing gradient problem in deep architectures. Layer Normalization is applied to stabilize the training process by normalizing the inputs across the features for each token independently. Together, these components ensure that we can stack many Transformer layers—enabling the learning of complex, high-level features—without the model becoming impossible to optimize effectively.
Compare the performance and utility of a standard Transformer decoder versus an encoder-only architecture for generating text.
An encoder-only architecture, like those used for bidirectional understanding, processes the entire input sequence at once, making it excellent for classification but poor for generative tasks. Conversely, a decoder-only architecture uses causal masking in its attention mechanism, ensuring that a position can only attend to previous positions. When building an LLM for text generation, the decoder-only approach is strictly superior because it respects the autoregressive property of language, ensuring that the model does not 'look ahead' at future tokens during training, which would be impossible during actual inference.
How do you implement the Feed-Forward Network component in a Transformer, and why does it have an expansion factor?
After the multi-head attention layer, each position passes through an identical position-wise feed-forward network. This usually consists of two linear transformations with a non-linear activation like GELU in between. The first layer expands the dimensionality (typically by a factor of 4), while the second projects it back to the original model dimension. The expansion factor is critical because it projects the data into a much higher-dimensional space, allowing the model to perform more complex non-linear feature transformations and store deeper knowledge about language patterns within the weights, acting as the primary knowledge repository for the LLM.
Check yourself
1. Why is the Scaled Dot-Product Attention mechanism divided by the square root of the dimension of the key vectors (sqrt(d_k))?
- A.To normalize the probability distribution for the Softmax function.
- B.To prevent the dot product from growing too large in magnitude, which pushes gradients into regions with tiny derivatives.
- C.To ensure the attention weights sum to exactly one.
- D.To allow the model to process longer sequences without exceeding memory limits.
Show answer
B. To prevent the dot product from growing too large in magnitude, which pushes gradients into regions with tiny derivatives.
Dividing by sqrt(d_k) prevents extremely large dot products that force the softmax into a 'saturated' state where gradients vanish. Option 0 is a side effect, not the primary reason; option 2 happens regardless of scaling; option 3 is false as scaling does not change the summation requirement.
2. What is the primary purpose of the residual connections around the attention and feed-forward sub-layers?
- A.To reduce the number of parameters in the model.
- B.To allow the model to learn identity mappings, which facilitates the training of very deep architectures.
- C.To force the model to prioritize the input tokens over learned representations.
- D.To eliminate the need for backpropagation through the entire network.
Show answer
B. To allow the model to learn identity mappings, which facilitates the training of very deep architectures.
Residual connections provide a 'shortcut' for gradients to flow through the network, preventing vanishing gradients in deep stacks. The other options are incorrect as they do not describe the fundamental function of residual blocks in deep learning.
3. In a decoder-only model, why is masking applied to the input during self-attention?
- A.To reduce the computational complexity of the attention matrix.
- B.To force the model to learn bidirectional context representation.
- C.To ensure that during training, a position cannot attend to tokens that appear after it in the sequence.
- D.To improve the diversity of the generated text.
Show answer
C. To ensure that during training, a position cannot attend to tokens that appear after it in the sequence.
Masking enforces the causal constraint required for autoregressive generation, ensuring the model only uses past info. Option 0 is irrelevant to masking; option 1 describes the encoder's task; option 3 is the purpose of masking.
4. How does Multi-Head Attention contribute to the transformer's performance compared to a single attention head?
- A.It allows the model to simultaneously attend to information from different representation subspaces at different positions.
- B.It significantly increases the inference speed by parallelizing head calculations.
- C.It replaces the need for positional encoding by providing spatial information.
- D.It reduces the total number of parameters by sharing weights across all heads.
Show answer
A. It allows the model to simultaneously attend to information from different representation subspaces at different positions.
Multiple heads allow the model to focus on different aspects of a sequence (e.g., syntax, semantics) concurrently. Option 1 is false because heads add overhead; option 2 is incorrect as positional encoding remains necessary; option 3 is false as heads have separate weights.
5. Why is the feed-forward network (FFN) applied to each position separately and identically in the transformer block?
- A.To allow for variable sequence lengths while maintaining spatial consistency across tokens.
- B.To facilitate the use of convolution-like properties without increasing computational cost.
- C.To allow for shared computation and parameter efficiency across all tokens in the sequence.
- D.To strictly limit the model's ability to cross-reference tokens.
Show answer
C. To allow for shared computation and parameter efficiency across all tokens in the sequence.
The position-wise FFN allows the model to process tokens independently, making it highly efficient for parallelization. Option 0 is not the primary reason; options 1 and 3 do not capture the efficiency/independence rationale.