Transformers — Missing Positional Encoding Scrambles Order

Every time you use ChatGPT, Google Translate, GitHub Copilot, or a speech-to-text app, a Transformer is doing the heavy lifting. Since the landmark 2017 paper 'Attention Is All You Need,' Transformers have become the dominant architecture in NLP, vision (ViT), protein folding (AlphaFold2), audio (Whisper), and even reinforcement learning. Understanding how they work at the implementation level — not just the diagram level — is the difference between using these models and building or fine-tuning them confidently.

Before Transformers, sequence models like LSTMs and GRUs had to process tokens one at a time, left to right. That meant long-range dependencies got diluted — by the time the model reached word 200, the gradient signal from word 3 had nearly vanished. Attention was proposed as an add-on fix to encoder-decoder RNNs, but 'Attention Is All You Need' made the radical claim: throw away the recurrence entirely. Let attention do everything. The result was massively parallelisable, faster to train, and dramatically better at capturing long-range context.

By the end of this article you'll be able to implement scaled dot-product attention and multi-head attention from scratch in PyTorch, explain exactly why we scale by the square root of the key dimension, trace the full data flow through a Transformer encoder block, and spot the three most expensive production mistakes teams make when deploying attention-based models. Let's build this up piece by piece.

The Core Engine: Scaled Dot-Product Attention

At the heart of the Transformer is the Scaled Dot-Product Attention mechanism. It operates on three matrices: Queries (Q), Keys (K), and Values (V).

The mechanism calculates the attention score by taking the dot product of the Query with all Keys, scaling by the square root of the dimension $d_k$ to prevent gradients from vanishing during softmax, and finally applying a softmax to obtain weights that are multiplied by the Values. The formula is:

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

This allows the model to dynamically focus on different parts of the input sequence regardless of their distance. The scaling factor is not a hyperparameter choice — it's mathematically necessary. As $d_k$ grows, the variance of the dot product grows linearly. Without scaling, the softmax saturates and gradients vanish.

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# io.thecodeforge: Production-grade Scaled Dot-Product Attention
class ScaledDotProductAttention(nn.Module):
    def __init__(self, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key, value, mask=None):
        d_k = query.size(-1)
        
        # Compute dot product scores: (batch, heads, seq, seq)
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        # Softmax converts scores to probabilities
        p_attn = F.softmax(scores, dim=-1)
        p_attn = self.dropout(p_attn)
        
        return torch.matmul(p_attn, value), p_attn

# Usage in io.thecodeforge training pipelines
# q, k, v shapes: (batch, heads, seq_len, d_k)
attention = ScaledDotProductAttention()
context_vector, weights = attention(torch.randn(1, 8, 128, 64), 
                                    torch.randn(1, 8, 128, 64), 
                                    torch.randn(1, 8, 128, 64))

Multi-Head Attention: Attending to Multiple Contexts

A single attention head might focus only on the syntactic relationship between words. Multi-Head Attention allows the model to jointly attend to information from different representation subspaces at different positions.

Essentially, we project $Q, K, V$ into $h$ different subspaces, perform attention in parallel, concatenate the results, and project them back. This allows one head to focus on 'who' (the subject), another on 'what' (the action), and another on 'where' (the location). The number of heads $h$ must divide the model dimension $d_{\text{model}}$ evenly so each head gets $d_k = d_{\text{model}} / h$.

class MultiHeadAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        super().__init__()
        assert d_model % h == 0
        self.d_k = d_model // h
        self.h = h
        # Linear layers for Q, K, V projections
        self.linears = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(3)])
        self.output_linear = nn.Linear(d_model, d_model)
        self.attention = ScaledDotProductAttention(dropout)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # 1) Linear projections and split into h heads
        query, key, value = [l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
                             for l, x in zip(self.linears, (query, key, value))]
        
        # 2) Apply attention
        x, self.attn = self.attention(query, key, value, mask=mask)
        
        # 3) Concatenate and apply final linear layer
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
        return self.output_linear(x)

Positional Encoding: Giving Order to a Bag of Tokens

Since the Transformer processes all tokens simultaneously, it has no inherent notion of sequence order. Positional encodings solve this by injecting position information into the input embeddings. The original paper used sinusoidal functions of different frequencies:

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

These encodings are added directly to the token embeddings. The intuition: each position gets a unique signature, and the model can learn to attend based on relative positions because the encoding at position pos+k can be expressed as a linear function of the encoding at pos.

import torch
import math

class PositionalEncoding(torch.nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = torch.nn.Dropout(p=dropout)
        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() * 
                             (-math.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)  # shape: (1, max_len, d_model)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x shape: (batch, seq_len, d_model)
        seq_len = x.size(1)
        x = x + self.pe[:, :seq_len, :]
        return self.dropout(x)

The Feed-Forward Network: Adding Non-Linearity and Depth

After the multi-head attention sub-layer, each token passes through a feed-forward network (FFN) that consists of two linear transformations with a ReLU activation in between:

$$\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2$$

The FFN is applied identically to every position — same weights, different activations per token. The inner dimension is typically 4x the model dimension (e.g., d_model=512, d_ff=2048). This expansion-contraction pattern lets the model learn complex transformations while keeping the parameter count manageable.

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        return self.linear2(self.dropout(F.relu(self.linear1(x))))

Layer Normalization & Residual Connections: Stabilizing Deep Networks

Each sub-layer (attention and FFN) is wrapped with a residual connection and followed by layer normalization. The original Transformer uses post-norm (norm after addition), but modern implementations often use pre-norm (norm before each sub-layer) because it stabilizes training.

Residual connection: $x = x + \text{Sublayer}(x)$ — this helps gradients flow through deep stacks.

Layer normalization: Normalizes across the feature dimension (d_model) to keep activations in a consistent range across layers.

class TransformerEncoderBlock(nn.Module):
    def __init__(self, d_model, h, d_ff, dropout=0.1):
        super().__init__()
        self.attention = MultiHeadAttention(h, d_model, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Pre-norm: norm before each sub-layer
        attn_out, _ = self.attention(self.norm1(x), self.norm1(x), self.norm1(x), mask)
        x = x + self.dropout(attn_out)
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x

Production Gotchas: Memory, Inference & Deployment

Deploying Transformers in production brings three major pain points: memory explosion from quadratic attention, inference latency from autoregressive decoding, and position extrapolation for sequences longer than training.

Memory: For a batch size of 1 and sequence length 4096 with d_model=512 and 12 heads, the attention logits alone take 4KB per head * 4096^2 = ~64MB per layer. Stack 12 layers and you exceed 1GB for just the attention scores.

Inference: Autoregressive decoding (common in GPT-style models) processes one token at a time, recomputing attention for all previous tokens each step. This is O(N^2) per step, making long generation expensive. Caching keys and values (KV cache) reduces complexity to O(N) per step.

Position extrapolation: If you trained on 512 tokens and try to generate 1024, learned positional embeddings will fail. Use Rotary Position Embedding (RoPE) which naturally allows extrapolation.

# io.thecodeforge: KV cache for autoregressive inference
class AttentionWithCache(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)
        self.wo = nn.Linear(d_model, d_model)

    def forward(self, x, past_kv=None):
        batch, seq_len, _ = x.shape
        q = self.wq(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        k = self.wk(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        v = self.wv(x).view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)
        past_kv = (k, v)

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        # causal mask here
        p_attn = F.softmax(scores, dim=-1)
        out = torch.matmul(p_attn, v)
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)
        return self.wo(out), past_kv

Training Transformers: Practical Tips for Stability and Speed

Training a Transformer from scratch is expensive and prone to instability. Here are the most impactful levers: - Learning rate schedule: Use a warmup phase (linear increase over first ~10k steps) followed by cosine decay. Without warmup, the attention weights can destabilise. - AdamW optimizer: Use weight decay separately from the learning rate (decoupled weight decay). The original Adam with L2 regularization can interact badly with LayerNorm. - Gradient clipping: Clip global norm to 1.0. The attention softmax can produce large gradients when logits are extreme. - Precision: Use mixed precision (fp16/bf16) to cut memory and speed up training. But ensure loss scaling works with attention softmax. - Initialization: Use small initial weights (e.g., xavier_uniform with gain 1.0 for FFN, and for attention projections scale by 1/sqrt(2 * num_layers) as in T5).

import torch
from torch.optim import AdamW

optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1e-4, total_steps=10000)

scaler = torch.cuda.amp.GradScaler()
for batch in dataloader:
    with torch.cuda.amp.autocast():
        loss = model(batch)
    scaler.scale(loss).backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()
    scheduler.step()