Positional Encodings Explained

Introduction

Welcome back to Deconstructing the Transformer Architecture! This is our third lesson together, and we're making excellent progress building the essential components that power modern language models. In our previous lessons, we explored multi-head attention mechanisms and the feed-forward networks with Add & Norm operations that process information within each Transformer block. Now, we turn to a fundamental challenge that every sequence model must solve: how do we help our model understand the order of elements in a sequence? Today, we'll tackle Positional Encodings, the ingenious solution that injects sequence order information into Transformer models. Unlike RNNs, which process sequences step by step and naturally encode position through their sequential processing, Transformers process all positions simultaneously through parallel attention operations. This creates both their greatest strength (parallelizability) and a critical challenge: without explicit positional information, attention mechanisms are completely permutation-invariant. We'll implement sinusoidal positional encodings from scratch, visualize their fascinating mathematical patterns, and understand why this elegant solution has become a cornerstone of modern NLP architectures. Let's go!

The Permutation Invariance Problem

\text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Implementing Sinusoidal Positional Encodings

Let's implement the complete PositionalEncoding class that creates these sinusoidal patterns:

Python

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        """
        Sinusoidal Positional Encoding
        
        Args:
            d_model: Model dimension
            max_len: Maximum sequence length
            dropout: Dropout rate
        """
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        
        # Create positional encoding matrix
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        
        # Compute the div_term for sinusoidal pattern
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
                           (-math.log(10000.0) / d_model))

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        """
        Sinusoidal Positional Encoding
        
        Args:
            d_model: Model dimension
            max_len: Maximum sequence length
            dropout: Dropout rate
        """
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        
        # Create positional encoding matrix
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        
        # Compute the div_term for sinusoidal pattern
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
                           (-math.log(10000.0) / d_model))

This initialization creates the foundation for our positional encodings. The position tensor contains values [0, 1, 2, ..., max_len-1] reshaped to enable broadcasting. The div_term computation is the key mathematical insight: it creates the frequency terms that make each dimension unique. By using torch.arange(0, d_model, 2), we generate frequencies for even indices, and the exponential computation creates the geometric progression of frequencies. The use of math.log(10000.0) transforms the original formulation into an exponential form that's computationally stable and efficient.

Now, let's complete the implementation by applying the sine and cosine functions:

Python

        # Apply sine to even indices (0, 2, 4, ...)
        pe[:, 0::2] = torch.sin(position * div_term)
        
        # Apply cosine to odd indices (1, 3, 5, ...)
        if d_model % 2 == 1:
            pe[:, 1::2] = torch.cos(position * div_term[:-1])
        else:
            pe[:, 1::2] = torch.cos(position * div_term)
        
        # Add batch dimension and register as buffer
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        """Add positional encoding to input embeddings"""
        seq_len = x.size(1)
        x = x + self.pe[:, :seq_len]
        return self.dropout(x)

        # Apply sine to even indices (0, 2, 4, ...)
        pe[:, 0::2] = torch.sin(position * div_term)
        
        # Apply cosine to odd indices (1, 3, 5, ...)
        if d_model % 2 == 1:
            pe[:, 1::2] = torch.cos(position * div_term[:-1])
        else:
            pe[:, 1::2] = torch.cos(position * div_term)
        
        # Add batch dimension and register as buffer
        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        """Add positional encoding to input embeddings"""
        seq_len = x.size(1)
        x = x + self.pe[:, :seq_len]
        return self.dropout(x)

The slicing operations pe[:, 0::2] and pe[:, 1::2] elegantly separate even and odd dimensions for sine and cosine application, respectively. The conditional logic handles the case where d_model is odd: since div_term has ceil(d_model/2) elements but odd indices have floor(d_model/2) positions, we use div_term[:-1] to match dimensions correctly. The register_buffer call is crucial: it ensures the positional encodings are part of the module's state but not trainable parameters, and they'll automatically move to the correct device with the model. The forward pass simply adds the appropriate slice of positional encodings to the input embeddings, enabling the model to distinguish between different positions.

Visualizing and Testing Our Implementation

Let's create the complete testing pipeline that generates the expected output and visualization:

Python

def visualize_positional_encoding():
    """Visualize positional encodings as heatmap"""
    print("Visualizing Positional Encodings...")
    
    d_model = 64
    max_len = 500
    
    # Create positional encoding
    pe_layer = PositionalEncoding(d_model, max_len, dropout=0.0)
    pe_matrix = pe_layer.pe.squeeze(0)  # (max_len, d_model)
    
    # Create visualization
    plt.figure(figsize=(12, 8))
    sns.heatmap(pe_matrix.numpy(), cmap='RdBu', center=0)
    plt.title('Sinusoidal Positional Encodings')
    
    print(f"Positional encoding shape: {pe_matrix.shape}")
    print(f"Value range: [{pe_matrix.min():.4f}, {pe_matrix.max():.4f}]")

def test_embeddings_integration():
    """Test token embeddings with positional encodings"""
    print("Testing Token + Positional Embeddings...")
    
    vocab_size, d_model, seq_len, batch_size = 1000, 64, 12, 2
    torch.manual_seed(42)
    token_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
    
    token_embedding = TokenEmbedding(vocab_size, d_model)
    positional_encoding = PositionalEncoding(d_model, dropout=0.0)
    
    token_embeds = token_embedding(token_ids)
    final_embeds = positional_encoding(token_embeds)
    
    print(f"Token IDs shape: {token_ids.shape}")
    print(f"Token embeddings shape: {token_embeds.shape}")
    print(f"Final embeddings shape: {final_embeds.shape}")
    print(f"Token embedding mean: {token_embeds.mean():.4f}")
    print(f"Final embedding mean: {final_embeds.mean():.4f}")

def visualize_positional_encoding():
    """Visualize positional encodings as heatmap"""
    print("Visualizing Positional Encodings...")
    
    d_model = 64
    max_len = 500
    
    # Create positional encoding
    pe_layer = PositionalEncoding(d_model, max_len, dropout=0.0)
    pe_matrix = pe_layer.pe.squeeze(0)  # (max_len, d_model)
    
    # Create visualization
    plt.figure(figsize=(12, 8))
    sns.heatmap(pe_matrix.numpy(), cmap='RdBu', center=0)
    plt.title('Sinusoidal Positional Encodings')
    
    print(f"Positional encoding shape: {pe_matrix.shape}")
    print(f"Value range: [{pe_matrix.min():.4f}, {pe_matrix.max():.4f}]")

def test_embeddings_integration():
    """Test token embeddings with positional encodings"""
    print("Testing Token + Positional Embeddings...")
    
    vocab_size, d_model, seq_len, batch_size = 1000, 64, 12, 2
    torch.manual_seed(42)
    token_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
    
    token_embedding = TokenEmbedding(vocab_size, d_model)
    positional_encoding = PositionalEncoding(d_model, dropout=0.0)
    
    token_embeds = token_embedding(token_ids)
    final_embeds = positional_encoding(token_embeds)
    
    print(f"Token IDs shape: {token_ids.shape}")
    print(f"Token embeddings shape: {token_embeds.shape}")
    print(f"Final embeddings shape: {final_embeds.shape}")
    print(f"Token embedding mean: {token_embeds.mean():.4f}")
    print(f"Final embedding mean: {final_embeds.mean():.4f}")

This testing pipeline demonstrates the complete workflow from token indices to position-aware embeddings. When we run our complete implementation, we get the following output, which confirms our positional encodings are working correctly:

text

Visualizing Positional Encodings...
Positional encoding shape: torch.Size([500, 64])
Value range: [-1.0000, 1.0000]
Testing Token + Positional Embeddings...
Token IDs shape: torch.Size([2, 12])
Token embeddings shape: torch.Size([2, 12, 64])
Final embeddings shape: torch.Size([2, 12, 64])
Token embedding mean: -0.0141
Final embedding mean: 0.4499

Visualizing Positional Encodings...
Positional encoding shape: torch.Size([500, 64])
Value range: [-1.0000, 1.0000]
Testing Token + Positional Embeddings...
Token IDs shape: torch.Size([2, 12])
Token embeddings shape: torch.Size([2, 12, 64])
Final embeddings shape: torch.Size([2, 12, 64])
Token embedding mean: -0.0141
Final embedding mean: 0.4499

Heatmap of sinusoidal positional encodings (500 tokens × 64 dims).

The heatmap reveals the elegant mathematical structure: rows represent token positions, and columns represent model dimensions. Leftmost channels oscillate rapidly with high-frequency patterns, while rightmost channels show nearly constant values, creating a spectrum from fine-grained to coarse-grained positional information. These results demonstrate several key properties: the value range [-1.0000, 1.0000] confirms that our sine and cosine functions are properly bounded, preventing positional encodings from overwhelming token embeddings. The shape preservation across all transformations shows that our implementation maintains dimensional consistency throughout the pipeline. Most importantly, the change in embedding statistics (from mean -0.0141 to 0.4499) indicates that positional information has been successfully integrated with token representations.

Introduction

The Permutation Invariance Problem

Join the 1M+ learners on CodeSignal

Understanding Sinusoidal Positional Encodings

Implementing Sinusoidal Positional Encodings

Creating Token Embeddings and Integration

Visualizing and Testing Our Implementation

Conclusion and Next Steps