Module 2: How Do Embeddings Work?
Transformer Encoders: The Architecture of Modern Embeddings
Capsule overview
Transformers revolutionized natural language processing in 2017 and are the foundation of every modern embedding model (BERT, GPT, OpenAI embeddings, Sentence-BERT).
In this capsule you'll learn the encoder-only architecture (BERT-style) that generates embeddings, how self-attention works to capture context, and why Transformers surpassed RNNs/LSTMs. You won't implement a Transformer from scratch (that's ML Engineering), but you'll understand each component conceptually so you can make informed decisions as an AI Engineer.
You'll also see concrete comparisons of pre-Transformer embeddings (Word2Vec) vs post-Transformer (BERT) and understand why the shift was so disruptive.
The Transformer revolution (2017)
Before Transformers (pre-2017):
Word2Vec, GloVe (2013-2015):
# Static word embeddings:
embed("bank") = [0.5, 0.3, 0.8, ...] # ALWAYS the same vector
# Problem: does NOT capture context
"I'm going to the river bank" → "bank" = [0.5, 0.3, 0.8]
"I'm going to the bank to withdraw money" → "bank" = [0.5, 0.3, 0.8] # ❌ Same vector
RNNs/LSTMs (2015-2017):
# They did capture context, but:
# ❌ Slow (sequential processing, not parallelizable)
# ❌ Problems with long sequences (vanishing gradients)
# ❌ Hard to train
After Transformers (2017+):
BERT, GPT, OpenAI embeddings (2018+):
# Contextual embeddings:
embed("I'm going to the river bank") → "bank" = [0.2, 0.5, 0.1, ...] # Riverbank
embed("I'm going to the bank to withdraw money") → "bank" = [0.8, 0.3, 0.6, ...] # Financial
# ✅ Different context → different embedding
# ✅ Fast (parallel processing)
# ✅ Scalable (billions of parameters)
Fundamental shift: From static embeddings to contextual embeddings.
Transformer architecture: Encoder-Decoder
Original paper (2017): "Attention Is All You Need"
Full architecture:
┌─────────────────────────────────┐
│ INPUT TOKENS │
└─────────────────────────────────┘
↓
┌─────────────────────────────────┐
│ ENCODER (N layers) │ ← For embeddings we use ONLY this
│ - Self-Attention │
│ - Feed-Forward │
└─────────────────────────────────┘
↓
┌─────────────────────────────────┐
│ DECODER (N layers) │ ← For text generation (GPT)
│ - Self-Attention │
│ - Cross-Attention │
│ - Feed-Forward │
└─────────────────────────────────┘
↓
┌─────────────────────────────────┐
│ OUTPUT TOKENS │
└─────────────────────────────────┘
For embeddings: We only use the ENCODER (encoder-only models).
Encoder-Only Models (BERT-style)
Encoder-only architecture:
Input: "Python is a programming language"
↓
┌─────────────────────────────────────┐
│ Tokenization + Initial Embedding │
└─────────────────────────────────────┘
↓
┌─────────────────────────────────────┐
│ Encoder Layer 1 │
│ - Multi-Head Self-Attention │
│ - Add & Norm │
│ - Feed-Forward Network │
│ - Add & Norm │
└─────────────────────────────────────┘
↓
┌─────────────────────────────────────┐
│ Encoder Layer 2 │
│ (same structure) │
└─────────────────────────────────────┘
↓
... (N layers)
↓
┌─────────────────────────────────────┐
│ Hidden States (output) │
│ One vector per token │
└─────────────────────────────────────┘
↓
┌─────────────────────────────────────┐
│ Pooling │
│ (reduce to 1 vector) │
└─────────────────────────────────────┘
↓
Final Embedding
Each encoder layer processes all tokens in parallel (not sequentially like an RNN).
Self-Attention: The heart of the Transformer
What self-attention is:
A mechanism that lets each token "pay attention" to all the other tokens in the sequence.
Visual example:
Input: "The river bank"
Self-attention for "bank":
- "The" → 0.1 (little attention)
- "river" → 0.8 (LOTS of attention) ✅
- "bank" → 1.0 (attention to itself, always high)
Result: "bank" understands it's near "river"
→ The embedding reflects "riverbank" (not "financial")
Attention weights visualization:
# Conceptual (not executable without a loaded model)
text = "The river bank is very full"
# Attention matrix (simplified):
# The river bank is very full
# The 1.0 0.1 0.2 0.1 0.0 0.0
# river 0.1 1.0 0.8 0.1 0.0 0.3 ← "river" attends to "bank"
# bank 0.2 0.8 1.0 0.1 0.0 0.0 ← "bank" attends to "river"
# is 0.1 0.1 0.1 1.0 0.2 0.5
# very 0.0 0.0 0.0 0.2 1.0 0.6
# full 0.0 0.3 0.0 0.5 0.6 1.0
High values indicate a semantic relationship between tokens.
Comparison with previous architectures:
| Architecture | Captures context? | Parallelizable? | Complexity |
|---|---|---|---|
| Word2Vec | ❌ No (static) | ✅ Yes (training) | O(1) per token |
| RNN/LSTM | ✅ Yes (sequential) | ❌ No | O(n) sequential |
| Transformer | ✅ Yes (self-attention) | ✅ Yes | O(n²) but parallel |
Trade-off: Transformers are O(n²) in memory but parallelizable (GPUs).
Components of an Encoder Layer
1. Multi-Head Self-Attention
What it does: Computes attention weights between all pairs of tokens.
"Multi-Head": Multiple attention mechanisms in parallel (e.g., 12 heads in BERT-base).
┌─────────────────────────────────────┐
│ Multi-Head Self-Attention │
│ │
│ Head 1: Attention to syntactic │
│ relationships │
│ Head 2: Attention to semantic │
│ relationships │
│ Head 3: Attention to local │
│ context │
│ ... │
│ Head 12: Attention to global │
│ context │
└─────────────────────────────────────┘
Each head learns different patterns automatically during training.
2. Add & Normalize (Residual Connection)
What it does: Adds the original input + the attention output, then normalizes.
# Pseudocode
output_attention = self_attention(input)
output = layer_norm(input + output_attention) # Residual connection
Why: Avoids vanishing gradients in deep networks (12-24 layers).
3. Feed-Forward Network
What it does: A non-linear transformation applied to each token independently.
# Pseudocode
def feed_forward(x):
# Expansion
hidden = linear_1(x) # [hidden_size] → [4 * hidden_size]
hidden = relu(hidden) # Activation
# Contraction
output = linear_2(hidden) # [4 * hidden_size] → [hidden_size]
return output
# Apply to each token
for token_vector in hidden_states:
token_vector = feed_forward(token_vector)
Intuition: Projects to a larger space (4x), applies non-linearity, projects back.
4. Add & Normalize (second time)
# Pseudocode
output_ffn = feed_forward(input)
output = layer_norm(input + output_ffn) # Residual connection
Result: The output of a complete encoder layer.
Encoder Stack (N layers)
BERT-base: 12 encoder layers
Input tokens
↓
┌────────────┐
│ Layer 1 │ → Captures local patterns (bigrams, trigrams)
├────────────┤
│ Layer 2 │
├────────────┤
│ Layer 3 │ → Captures syntactic relationships
├────────────┤
│ Layer 4-6 │ → Captures basic semantics
├────────────┤
│ Layer 7-9 │ → Captures deep semantics
├────────────┤
│ Layer 10-12│ → Captures global context
└────────────┘
↓
Hidden states (contextualized)
More layers = more contextual depth.
Self-Attention: The math (simplified)
Key concepts:
For each token, we compute:
- Query (Q): "What am I looking for?"
- Key (K): "What do I offer?"
- Value (V): "My content"
Intuition: Token A makes a query, compares it with the keys of all tokens, obtains attention weights, and computes a weighted sum of the values.
Attention formula (simplified):
Attention(Q, K, V) = softmax(Q · K^T / √d_k) · V
Where:
- Q, K, V: Matrices of queries, keys, values
- Q · K^T: Dot product (similarity)
- √d_k: Scaling factor (numerical stability)
- softmax: Normalizes to probabilities (sums to 1)
- · V: Weighted sum of the values
You don't need to memorize this. What matters: self-attention computes similarity between tokens.
Conceptual example:
# Token "bank" in "The river bank"
# Query of "bank": [0.5, 0.3, 0.8]
# Keys of all tokens:
# "The": [0.1, 0.2, 0.1]
# "river": [0.6, 0.4, 0.9]
# "bank": [0.5, 0.3, 0.8] (itself)
# Dot products (similarity):
# "bank" · "The" = 0.19 (low)
# "bank" · "river" = 1.02 (high - related) ✅
# "bank" · "bank" = 0.98 (high - itself)
# Softmax (normalize):
# "The": 0.1
# "river": 0.5 ← Most attention
# "bank": 0.4
# Weighted sum of values → New vector of "bank" (contextualized)
Result: "bank" updated with information from "river".
Positional Encoding
Problem:
Self-attention has no notion of order (permutation-invariant).
# Without positional encoding:
embed("The river bank") ≈ embed("bank river The") # ❌ Order doesn't matter
Solution: Add position information.
Positional Encoding (sinusoidal):
# Pseudocode (OpenAI/BERT use learned embeddings)
def positional_encoding(position, d_model):
"""
A unique encoding for each position
Args:
position: The token's position (0, 1, 2, ...)
d_model: The embedding dimensions
"""
encoding = []
for i in range(d_model):
if i % 2 == 0:
encoding.append(sin(position / 10000^(i/d_model)))
else:
encoding.append(cos(position / 10000^(i/d_model)))
return encoding
# Add to the token embeddings
token_embedding = [0.5, 0.3, 0.8, ...]
position_encoding = positional_encoding(position=2, d_model=768)
final_embedding = token_embedding + position_encoding
Now the model knows that "bank" is in position 2.
Comparison: Word2Vec vs Transformer Embeddings
Word2Vec (2013):
# Static embeddings - NOT contextual
vocab = {
"bank": [0.5, 0.3, 0.8], # ALWAYS this vector
"river": [0.2, 0.6, 0.1],
"money": [0.9, 0.1, 0.7]
}
# Problem:
text_1 = "river bank"
text_2 = "money bank"
# "bank" has the SAME embedding in both contexts ❌
BERT / OpenAI Embeddings (2018+):
# Contextual embeddings
def get_contextual_embedding(text, word):
# Tokenize the full text
tokens = tokenize(text)
# Process with the Transformer
hidden_states = transformer_encoder(tokens)
# The embedding of "word" depends on context
word_idx = tokens.index(word)
return hidden_states[word_idx]
# Example:
emb_1 = get_contextual_embedding("river bank", "bank")
# → [0.2, 0.5, 0.1, ...] (riverbank)
emb_2 = get_contextual_embedding("money bank", "bank")
# → [0.8, 0.3, 0.6, ...] (financial)
# DIFFERENT embeddings depending on context ✅
Popular encoder-only models
1. BERT (2018) - Google
Architecture: 12 layers, 768 hidden, 12 attention heads
Parameters: 110M (base), 340M (large)
Embeddings: 768 dims (base), 1024 dims (large)
Training: Masked Language Modeling (MLM) + Next Sentence Prediction
Use: A base for fine-tuning on specific tasks.
2. Sentence-BERT (2019)
Architecture: BERT fine-tuned for sentence embeddings
Parameters: 22M (MiniLM-L6), 110M (base)
Embeddings: 384 dims (MiniLM), 768 dims (base)
Training: Siamese network with similar pairs
Pooling: Mean pooling (default)
Use: Sentence embeddings out-of-the-box (better than vanilla BERT).
3. OpenAI text-embedding-3-small (2024)
Architecture: Transformer encoder (details not public)
Parameters: Not public
Embeddings: 1536 dims
Training: Contrastive learning at massive scale
Pooling: Not documented (probably mean)
Use: API, not local. Excellent out-of-the-box quality.
4. BGE (BAAI General Embedding) - 2023
Architecture: BERT-large fine-tuned
Parameters: 340M (large)
Embeddings: 1024 dims
Training: Contrastive learning + hard negatives
Pooling: CLS token
Use: State-of-the-art open-source, local.
Advantages of Transformers over previous architectures
vs RNN/LSTM:
| Aspect | RNN/LSTM | Transformer |
|---|---|---|
| Parallelization | ❌ No (sequential) | ✅ Yes (all tokens in parallel) |
| Long sequences | ❌ Vanishing gradients | ✅ Self-attention captures long range |
| Speed | ❌ Slow (GPU underused) | ✅ Fast (GPU fully utilized) |
| Scalability | ❌ Hard to scale | ✅ Scales to billions of parameters |
vs Word2Vec/GloVe:
| Aspect | Word2Vec/GloVe | Transformer |
|---|---|---|
| Contextual | ❌ No (static) | ✅ Yes (dynamic context) |
| Polysemy | ❌ No (1 vector per word) | ✅ Yes (different embeddings by context) |
| Quality | ⚠️ Basic | ✅ State-of-the-art |
Limitations of Transformers
1. Quadratic complexity (O(n²)):
# Attention computes similarity between ALL pairs of tokens
n_tokens = 1000
attention_computations = n_tokens ** 2 # 1,000,000 operations
# For long texts (10K tokens):
n_tokens = 10000
attention_computations = n_tokens ** 2 # 100,000,000 operations 😱
Limit: OpenAI embeddings = 8191 tokens max.
2. Memory intensive:
# Attention matrix: [batch_size, n_heads, seq_len, seq_len]
batch_size = 32
n_heads = 12
seq_len = 512
attention_matrix_size = batch_size * n_heads * seq_len * seq_len * 4 # bytes (float32)
# = 32 * 12 * 512 * 512 * 4 ≈ 402 MB (attention alone!)
Requires GPUs with a lot of VRAM.
3. Doesn't understand math/logic perfectly:
# Transformers are good at natural language, not exact logic:
query = "How much is 123456 × 789012?"
embedding = get_embedding(query)
# The embedding captures "a multiplication question" but NOT the exact result
For calculations: use tools (tool use), not embeddings.
Exercises
Exercise 1: Identify the architecture
Which architecture is best for each case?
Case A: Generate embeddings of sentences Case B: Generate text (chatbot) Case C: Machine translation
See solution
-
Case A: Encoder-only (BERT, Sentence-BERT, OpenAI embeddings)
- You only need a vector representation, not generation
-
Case B: Decoder-only (GPT, Claude)
- Text generation token by token
-
Case C: Encoder-Decoder (T5, BART)
- The encoder processes the source language, the decoder generates the target language
Embeddings = Encoder-only.
Exercise 2: Predict attention
Given the text "The cat chases the mouse", which token should receive the MOST attention from "chases"?
Options: A) "The" B) "cat" C) "the" D) "mouse"
See solution
Answer: B) "cat" and D) "mouse"
Reason:
- "chases" is a verb
- The verb attends to the subject ("cat") and the object ("mouse")
- "The" and "the" are articles (less semantically relevant)
Self-attention learns syntactic relationships automatically.
Exercise 3: Compare Word2Vec vs Transformer
Which approach better captures the difference in "bank"?
text_1 = "I'm going to the river bank to fish"
text_2 = "I'm going to the bank to deposit money"
Options: A) Word2Vec (static embeddings) B) Transformer (contextual embeddings)
See solution
Answer: B) Transformer
Word2Vec:
- "bank" = [0.5, 0.3, 0.8] in BOTH contexts ❌
- No difference between "riverbank" and "financial"
Transformer:
- Text 1: "bank" with the context of "river" + "fish" → Embedding A (riverbank)
- Text 2: "bank" with the context of "deposit" + "money" → Embedding B (financial)
- Embeddings A and B are DIFFERENT ✅
Transformers solve polysemy (multiple meanings).
Exercise 4: Calculate complexity
How many attention operations for 512 tokens?
n_tokens = 512
# How many operations?
See solution
n_tokens = 512
attention_operations = n_tokens ** 2
print(f"Operations: {attention_operations}")
# → 262,144 operations
# With 12 attention heads (BERT-base):
total_operations = attention_operations * 12
print(f"Total: {total_operations}")
# → 3,145,728 operations per layer
# × 12 layers = ~37M operations total
That's why Transformers require powerful GPUs.
Common troubleshooting
Problem 1: "Token limit exceeded"
# Error: Text > 8191 tokens (OpenAI)
long_text = "..." * 10000
embedding = get_embedding(long_text) # ❌ Error
Solution: Chunk the text (Module 4).
Problem 2: Identical embeddings for different contexts
# You expected different, but they're almost identical:
emb_1 = embed("river bank")
emb_2 = embed("money bank")
similarity = cosine_similarity(emb_1, emb_2) # 0.95 (very similar)
Cause: The model isn't contextual enough (e.g., legacy Word2Vec).
Solution: Use a modern model (OpenAI, SBERT, BGE).
Summary
What you learned:
- ✅ Transformers: Encoder-only architecture for embeddings
- ✅ Self-attention: Captures context across all tokens
- ✅ Encoder layers: A stack of 12-24 layers processes in parallel
- ✅ Contextual embeddings: "bank" has different embeddings depending on context
- ✅ Advantages: Parallel, scalable, captures long range
- ✅ Limitations: O(n²) complexity, token limit
Key concepts:
- Self-attention lets each token see all the other tokens
- An encoder stack of N layers captures context progressively
- Contextual embeddings > static embeddings (Word2Vec)
Additional resources
- Illustrated Transformer - Excellent visualization
- Attention Is All You Need - Original paper
- BERT Explained - Encoder-only architecture
- The Annotated Transformer - Annotated implementation
- Transformers from Scratch - Technical deep dive
In the next capsule
Capsule 03: Tokenization
You'll learn:
- The BPE (Byte Pair Encoding) algorithm
- WordPiece tokenization (BERT)
- tiktoken code (OpenAI)
- Why sub-words > full words
- Vocabulary and out-of-vocabulary handling
From architecture to the first step of the pipeline.
Module 2 - Embeddings Deep Dive Guide Understanding the architecture that revolutionized NLP