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:

ArchitectureCaptures context?Parallelizable?Complexity
Word2Vec❌ No (static)✅ Yes (training)O(1) per token
RNN/LSTM✅ Yes (sequential)❌ NoO(n) sequential
Transformer✅ Yes (self-attention)✅ YesO(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:

  1. Query (Q): "What am I looking for?"
  2. Key (K): "What do I offer?"
  3. 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:

AspectRNN/LSTMTransformer
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:

AspectWord2Vec/GloVeTransformer
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:

  1. Self-attention lets each token see all the other tokens
  2. An encoder stack of N layers captures context progressively
  3. Contextual embeddings > static embeddings (Word2Vec)

Additional resources

  1. Illustrated Transformer - Excellent visualization
  2. Attention Is All You Need - Original paper
  3. BERT Explained - Encoder-only architecture
  4. The Annotated Transformer - Annotated implementation
  5. 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