Token Embeddings & Positional Encoding

import torch
import torch.nn as nn
import math

class TokenEmbedding(nn.Module):
    """
    Convert token indices to dense vectors.
    
    This is simply a learned lookup table: each token ID maps to a d_model-dimensional vector.
    """
    
    def __init__(self, vocab_size: int, d_model: int):
        super().__init__()
        # The embedding matrix: [vocab_size, d_model]
        # PyTorch's nn.Embedding handles the lookup efficiently
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model
    
    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        """
        Args:
            token_ids: (batch_size, seq_len) - integer token IDs
        
        Returns:
            embeddings: (batch_size, seq_len, d_model) - dense vectors
        """
        return self.embedding(token_ids)

# Let's see it in action
vocab_size = 10000
d_model = 256

embed = TokenEmbedding(vocab_size, d_model)

# Simulate a batch of tokenized text
# batch_size=2, seq_len=5
token_ids = torch.tensor([
    [101, 2054, 2003, 2115, 2171],  # "What is your name"
    [101, 1045, 2066, 2653, 102],   # "I like pizza"
])

embeddings = embed(token_ids)

print(f"Input token IDs shape: {token_ids.shape}")
print(f"Output embeddings shape: {embeddings.shape}")
print(f"\nEach token is now a {d_model}-dimensional vector.")
print(f"\nFirst token (ID {token_ids[0, 0].item()}) embedding (first 8 dims):")
print(f"  {embeddings[0, 0, :8].detach().numpy().round(4)}")

class LearnedPositionalEmbedding(nn.Module):
    """
    Learned positional embeddings (GPT-2 / BERT style).
    
    Each position 0, 1, 2, ... gets its own learned d_model-dimensional vector.
    These are added to the token embeddings to inject position information.
    """
    
    def __init__(self, d_model: int, max_seq_len: int = 512):
        super().__init__()
        # Position embedding matrix: [max_seq_len, d_model]
        self.pos_embedding = nn.Embedding(max_seq_len, d_model)
        self.max_seq_len = max_seq_len
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (batch_size, seq_len, d_model) - token embeddings
        
        Returns:
            (batch_size, seq_len, d_model) - embeddings with position info added
        """
        batch_size, seq_len, d_model = x.shape
        
        # Create position indices: [0, 1, 2, ..., seq_len-1]
        positions = torch.arange(seq_len, device=x.device)
        
        # Look up position embeddings: (seq_len, d_model)
        pos_emb = self.pos_embedding(positions)
        
        # Add to input (broadcasting handles batch dimension)
        return x + pos_emb

# Example
pos_embed = LearnedPositionalEmbedding(d_model=256, max_seq_len=512)

# Token embeddings (from before)
x = torch.randn(2, 10, 256)  # batch=2, seq_len=10, d_model=256

# Add position information
x_with_pos = pos_embed(x)

print(f"Input shape: {x.shape}")
print(f"Output shape: {x_with_pos.shape}")
print(f"\nPosition embedding for position 0 (first 8 dims):")
print(f"  {pos_embed.pos_embedding.weight[0, :8].detach().numpy().round(4)}")
print(f"\nPosition embedding for position 1 (first 8 dims):")
print(f"  {pos_embed.pos_embedding.weight[1, :8].detach().numpy().round(4)}")

# Count parameters
n_params = sum(p.numel() for p in pos_embed.parameters())
print(f"\nPosition embedding parameters: {n_params:,} ({n_params/1e6:.2f}M)")

class ALiBiPositionalBias(nn.Module):
    """
    ALiBi: Attention with Linear Biases.
    
    Instead of adding position info to embeddings, we add a distance-based
    penalty directly to attention scores. Different heads use different
    slopes, letting some focus locally and others globally.
    
    Reference: https://arxiv.org/abs/2108.12409
    """
    
    def __init__(self, num_heads: int):
        super().__init__()
        
        # Slopes form a geometric sequence: 2^(-8/n), 2^(-16/n), ...
        # Head 0 has the steepest slope (strongest local focus)
        # Last head has the gentlest slope (most global)
        slopes = torch.tensor([
            2 ** (-8 * (i + 1) / num_heads) 
            for i in range(num_heads)
        ])
        # Register as buffer (not a parameter - we don't learn it)
        self.register_buffer('slopes', slopes)
        self.num_heads = num_heads
    
    def forward(self, seq_len: int) -> torch.Tensor:
        """
        Compute ALiBi bias matrix.
        
        Args:
            seq_len: length of the sequence
        
        Returns:
            bias: (num_heads, seq_len, seq_len) - add this to attention scores
        """
        # Position indices
        positions = torch.arange(seq_len, device=self.slopes.device)
        
        # Compute pairwise distances: |i - j|
        # positions.unsqueeze(0) is [1, seq_len]
        # positions.unsqueeze(1) is [seq_len, 1]
        # Broadcasting gives [seq_len, seq_len]
        distances = torch.abs(positions.unsqueeze(0) - positions.unsqueeze(1))
        
        # Apply slopes: -slope * distance
        # slopes is [num_heads], distances is [seq_len, seq_len]
        # Result is [num_heads, seq_len, seq_len]
        biases = -self.slopes.view(-1, 1, 1) * distances.unsqueeze(0).float()
        
        return biases

# Let's see the ALiBi biases
alibi = ALiBiPositionalBias(num_heads=4)
biases = alibi(seq_len=6)

print(f"ALiBi bias shape: {biases.shape}")
print(f"Slopes per head: {alibi.slopes.numpy().round(4)}")
print()

# Head 0: steepest slope (strongest local preference)
print("Head 0 biases (steep slope - prefers nearby tokens):")
print(biases[0].numpy().round(3))
print()

# Head 3: gentlest slope (most global)
print("Head 3 biases (gentle slope - more global attention):")
print(biases[3].numpy().round(3))

class RotaryPositionalEmbedding(nn.Module):
    """
    RoPE: Rotary Position Embeddings.
    
    Rotates query and key vectors by position-dependent angles.
    The dot product Q·K then naturally encodes relative position.
    
    Reference: https://arxiv.org/abs/2104.09864
    """
    
    def __init__(self, d_model: int, max_seq_len: int = 5000, base: int = 10000):
        super().__init__()
        
        # Compute inverse frequencies for each dimension pair
        # These determine the rotation speed for each pair
        # Lower dimensions rotate faster (higher frequency)
        inv_freq = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model))
        self.register_buffer('inv_freq', inv_freq)
        
        # Precompute cos/sin for all positions
        self._build_cache(max_seq_len)
    
    def _build_cache(self, seq_len: int):
        """Precompute rotation matrices for efficiency."""
        positions = torch.arange(seq_len, device=self.inv_freq.device)
        
        # Outer product: positions × frequencies
        # Shape: [seq_len, d_model/2]
        freqs = torch.outer(positions.float(), self.inv_freq)
        
        # Cache cos and sin
        self.register_buffer('cos_cached', freqs.cos(), persistent=False)
        self.register_buffer('sin_cached', freqs.sin(), persistent=False)
    
    def forward(self, seq_len: int):
        """
        Get rotation matrices for the given sequence length.
        
        Returns:
            cos: (seq_len, d_model/2)
            sin: (seq_len, d_model/2)
        """
        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]


def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """
    Apply rotary embeddings to input tensor.
    
    Args:
        x: (batch, seq_len, d_model) or (batch, n_heads, seq_len, d_k)
        cos: (seq_len, d_model/2)
        sin: (seq_len, d_model/2)
    
    Returns:
        Rotated tensor with same shape as x
    """
    # Split into even and odd dimensions
    x1 = x[..., 0::2]  # Even indices
    x2 = x[..., 1::2]  # Odd indices
    
    # Apply rotation
    # [x1, x2] → [x1*cos - x2*sin, x1*sin + x2*cos]
    x_rotated = torch.stack([
        x1 * cos - x2 * sin,
        x1 * sin + x2 * cos
    ], dim=-1).flatten(-2)
    
    return x_rotated

# Example
d_model = 64
rope = RotaryPositionalEmbedding(d_model=d_model, max_seq_len=512)

cos, sin = rope(seq_len=10)

print(f"Cos cache shape: {cos.shape}")
print(f"Sin cache shape: {sin.shape}")
print()
print("Frequencies (inverse) for first 4 dimension pairs:")
print(f"  {rope.inv_freq[:4].numpy().round(4)}")
print()
print("Notice: lower dimensions have higher frequencies (rotate faster).")
print("This lets the model capture both local (fast) and global (slow) patterns.")

# Visualize the rotation frequencies
import matplotlib.pyplot as plt
import numpy as np

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# Left: show cos values across positions for different frequency pairs
ax = axes[0]
positions = np.arange(50)
cos_np = rope.cos_cached[:50].numpy()

for i in [0, 4, 8, 12]:  # Different frequency indices
    ax.plot(positions, cos_np[:, i], label=f'Dim pair {i}')

ax.set_xlabel('Position')
ax.set_ylabel('Cos value')
ax.set_title('RoPE: Rotation patterns by dimension pair')
ax.legend()
ax.grid(True, alpha=0.3)

# Right: show the full cos cache as a heatmap
ax = axes[1]
im = ax.imshow(cos_np.T, aspect='auto', cmap='RdBu', vmin=-1, vmax=1)
ax.set_xlabel('Position')
ax.set_ylabel('Dimension pair')
ax.set_title('RoPE: Cos values (heatmap)')
plt.colorbar(im, ax=ax)

plt.tight_layout()
plt.show()

print("Low dimension pairs (top of heatmap) oscillate rapidly.")
print("High dimension pairs (bottom) change slowly across positions.")

class TransformerEmbedding(nn.Module):
    """
    Complete embedding module: token embeddings + positional encoding.
    
    Converts token IDs to dense vectors with position information.
    """
    
    def __init__(self, vocab_size: int, d_model: int, max_seq_len: int, dropout: float = 0.1):
        super().__init__()
        
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.pos_embedding = nn.Embedding(max_seq_len, d_model)
        self.dropout = nn.Dropout(dropout)
        self.d_model = d_model
    
    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        """
        Args:
            token_ids: (batch_size, seq_len) - integer token IDs
        
        Returns:
            (batch_size, seq_len, d_model) - embedded tokens with positions
        """
        batch_size, seq_len = token_ids.shape
        
        # Token embeddings: (batch, seq_len, d_model)
        tok_emb = self.token_embedding(token_ids)
        
        # Position embeddings: (seq_len, d_model)
        positions = torch.arange(seq_len, device=token_ids.device)
        pos_emb = self.pos_embedding(positions)
        
        # Combine and apply dropout
        x = tok_emb + pos_emb
        x = self.dropout(x)
        
        return x

# Create embedding module with our config
embedding = TransformerEmbedding(
    vocab_size=10000,
    d_model=256,
    max_seq_len=512,
    dropout=0.1
)

# Test it
token_ids = torch.randint(0, 10000, (4, 32))  # batch=4, seq_len=32
embeddings = embedding(token_ids)

print(f"Input tokens: {token_ids.shape}")
print(f"Output embeddings: {embeddings.shape}")

# Count parameters
n_params = sum(p.numel() for p in embedding.parameters())
print(f"\nEmbedding parameters: {n_params:,}")
print(f"  Token embeddings: {10000 * 256:,}")
print(f"  Position embeddings: {512 * 256:,}")