Chapter 3: Positional Encoding and Embeddings

Introduction

Unlike RNNs that process sequences token by token, Transformers process all positions simultaneously. This creates a fundamental challenge: how does the model know the order of tokens? This chapter explains positional encoding and the broader token embedding process.

The Position Problem

Why Position Matters

Consider these sentences:

• "The dog bit the man"
• "The man bit the dog"

Same tokens, completely different meanings. Position is crucial for understanding language.

Transformers and Permutation Invariance

Without positional information, self-attention is permutation invariant:

$$\text{SelfAttention}(\pi(X)) = \pi(\text{SelfAttention}(X))$$

Where $\pi$ is any permutation. The model treats "dog bit man" and "man bit dog" identically!

# Demonstration: attention scores are position-independent
import torch

def attention_scores(Q, K):
    return torch.softmax(Q @ K.T, dim=-1)

# Original sequence
x1 = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])

# Permuted sequence (swap positions 0 and 1)
x2 = torch.tensor([[0.0, 1.0], [1.0, 0.0], [1.0, 1.0]])

# Same attention pattern (just permuted)!
print(attention_scores(x1, x1))
print(attention_scores(x2, x2))

Token Embeddings

From Tokens to Vectors

Before adding positional information, we first convert tokens to dense vectors:

$$E = \text{Embedding}(\text{token\_ids})$$

Implementation

import torch
import torch.nn as nn

class TokenEmbedding(nn.Module):
    """
    Convert token IDs to dense vectors.
    """
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model
    
    def forward(self, token_ids):
        """
        token_ids: (batch, seq_len) of integer token IDs
        Returns: (batch, seq_len, d_model) embeddings
        """
        # Scale embeddings (from original Transformer paper)
        return self.embedding(token_ids) * (self.d_model ** 0.5)

# Example
vocab_size = 50000
d_model = 512
embedding_layer = TokenEmbedding(vocab_size, d_model)

tokens = torch.tensor([[101, 7592, 2088, 102]])  # [CLS] Hello World [SEP]
embeddings = embedding_layer(tokens)
print(f"Shape: {embeddings.shape}")  # (1, 4, 512)

Sinusoidal Positional Encoding

The Original Approach

The original Transformer paper used fixed sinusoidal functions:

$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$ $$PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$

Where:

• $pos$ = position in sequence
• $i$ = dimension index
• $d_{model}$ = model dimension

Intuition

Each dimension oscillates at a different frequency:

• Low dimensions: slow oscillation (capture long-range positions)
• High dimensions: fast oscillation (capture fine-grained positions)

Implementation

import numpy as np

def get_sinusoidal_encoding(max_len, d_model):
    """
    Generate sinusoidal positional encodings.
    
    Args:
        max_len: Maximum sequence length
        d_model: Model dimension
    
    Returns:
        Positional encoding matrix (max_len, d_model)
    """
    position = np.arange(max_len)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
    
    pe = np.zeros((max_len, d_model))
    pe[:, 0::2] = np.sin(position * div_term)  # Even indices
    pe[:, 1::2] = np.cos(position * div_term)  # Odd indices
    
    return pe

# Visualize
import matplotlib.pyplot as plt

pe = get_sinusoidal_encoding(100, 512)
plt.figure(figsize=(12, 4))
plt.imshow(pe.T, aspect='auto', cmap='RdBu')
plt.xlabel('Position')
plt.ylabel('Dimension')
plt.colorbar(label='Value')
plt.title('Sinusoidal Positional Encoding')
plt.show()

Properties of Sinusoidal Encoding

1. Unique: Each position has a unique encoding
2. Bounded: Values always in [-1, 1]
3. Relative positions: Linear relationship enables learning relative positions
4. Extrapolation: Can generalize to longer sequences than seen in training

PyTorch Implementation

class SinusoidalPositionalEncoding(nn.Module):
    """
    Sinusoidal positional encoding from 'Attention Is All You Need'.
    """
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        # Create positional encoding matrix
        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() * (-np.log(10000.0) / d_model)
        )
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)  # Add batch dimension
        
        # Register as buffer (not a parameter, but saved with model)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        """
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

Learned Positional Embeddings

An Alternative Approach

Many modern LLMs use learned position embeddings:

class LearnedPositionalEmbedding(nn.Module):
    """
    Learned positional embeddings (e.g., used in GPT-2).
    """
    def __init__(self, max_len, d_model):
        super().__init__()
        self.position_embedding = nn.Embedding(max_len, d_model)
    
    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        """
        positions = torch.arange(x.size(1), device=x.device)
        return x + self.position_embedding(positions)

Comparison

Aspect	Sinusoidal	Learned
Parameters	None	max_len × d_model
Extrapolation	Better	Limited to max_len
Performance	Comparable	Comparable
Flexibility	Fixed	Adapts to data

Rotary Position Embedding (RoPE)

Modern Innovation

RoPE, used in LLaMA, Mistral, and many modern models, encodes position through rotation:

$$f_q(x_m, m) = x_m e^{im\theta}$$

Key Benefits

1. Relative positions: Naturally captures relative position information
2. Extrapolation: Better handling of longer sequences
3. Efficiency: Applied directly to Q, K without additional parameters

Simplified Implementation

def rotary_embedding(x, seq_len, d_model):
    """
    Simplified RoPE implementation.
    x: (batch, seq_len, d_model)
    """
    # Generate rotation angles
    position = torch.arange(seq_len).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
    angles = position * div_term
    
    # Apply rotation to pairs of dimensions
    x_rotated = x.clone()
    x_rotated[..., 0::2] = x[..., 0::2] * torch.cos(angles) - x[..., 1::2] * torch.sin(angles)
    x_rotated[..., 1::2] = x[..., 0::2] * torch.sin(angles) + x[..., 1::2] * torch.cos(angles)
    
    return x_rotated

ALiBi: Attention with Linear Biases

Another Approach

ALiBi (Attention with Linear Biases) adds a linear penalty based on distance:

$$\text{softmax}(q_i \cdot k_j - m \cdot |i - j|)$$

Where $m$ is a head-specific slope.

Benefits

1. No additional parameters
2. Excellent length extrapolation
3. Simple implementation

Complete Embedding Layer

Putting It All Together

class TransformerEmbedding(nn.Module):
    """
    Complete embedding layer: token + position.
    """
    def __init__(self, vocab_size, d_model, max_len=5000, 
                 pos_encoding='sinusoidal', dropout=0.1):
        super().__init__()
        
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model
        self.dropout = nn.Dropout(dropout)
        
        if pos_encoding == 'sinusoidal':
            self.pos_encoding = SinusoidalPositionalEncoding(d_model, max_len)
        elif pos_encoding == 'learned':
            self.pos_encoding = LearnedPositionalEmbedding(max_len, d_model)
        else:
            raise ValueError(f"Unknown pos_encoding: {pos_encoding}")
    
    def forward(self, token_ids):
        """
        token_ids: (batch, seq_len)
        Returns: (batch, seq_len, d_model)
        """
        # Token embeddings (scaled)
        x = self.token_embedding(token_ids) * (self.d_model ** 0.5)
        
        # Add positional encoding
        x = self.pos_encoding(x)
        
        return self.dropout(x)

# Usage
embedding = TransformerEmbedding(
    vocab_size=50000,
    d_model=512,
    max_len=2048,
    pos_encoding='sinusoidal'
)

tokens = torch.randint(0, 50000, (32, 100))  # Batch of 32, seq_len 100
embedded = embedding(tokens)
print(f"Output shape: {embedded.shape}")  # (32, 100, 512)

The Complete Picture

Here's how tokens flow through embeddings to attention:

Token IDs: [101, 7592, 2088, 102]
     ↓
Token Embedding: Lookup in embedding table
     ↓
[0.23, -0.45, ...] [0.89, 0.12, ...] [0.45, -0.23, ...] [0.67, 0.34, ...]
     ↓
Add Positional Encoding
     ↓
[0.23+PE_0, -0.45+PE_0, ...] [0.89+PE_1, 0.12+PE_1, ...] ...
     ↓
Multi-Head Attention
     ↓
Output with position-aware context

Key Takeaways

Summary

1. Position information is essential—Transformers are otherwise permutation invariant
2. Token embeddings convert discrete tokens to continuous vectors
3. Sinusoidal encoding uses fixed frequency patterns (original Transformer)
4. Learned embeddings adapt to data but have fixed max length
5. RoPE encodes relative positions through rotation (modern LLMs)
6. ALiBi uses distance-based attention biases

Looking Ahead

Now that we understand the complete Transformer architecture—attention, multi-head attention, and positional encoding—we're ready to move on to Module 3, where we'll explore how these massive models are trained.

---

← Chapter 2: Self & Multi-Head Attention

Module 2 Quiz →