Transformer Positional Encoding — Flat Prediction Fix

In 2017, eight researchers at Google Brain published a 15-page paper that quietly made recurrent neural networks obsolete. 'Attention Is All You Need' introduced the Transformer architecture, and within three years it became the backbone of GPT, BERT, T5, DALL-E, Whisper, and virtually every state-of-the-art model in language, vision, audio, and protein folding. If you work in ML, this paper is not optional reading — it is the constitution of modern deep learning.

Before Transformers, sequence models like LSTMs and GRUs processed tokens one at a time, left to right. That sequential dependency meant you couldn't parallelise training across time steps, and long-range dependencies decayed badly across hundreds of tokens. The Transformer killed recurrence entirely. By replacing recurrence with self-attention, it achieved parallelism across the entire sequence and made long-range dependency a first-class citizen.

Here's the thing most tutorials skip: the paper isn't just theory you read once and file away. Every line — from the scaling factor to the learning rate schedule to the label smoothing — encodes a hard-won production lesson. The teams that deploy Transformers without internalising those details don't get elegant training curves. They get OOM crashes, flat predictions, and models that look great on validation but fail in the wild. This article makes sure you're not one of them.

Scaled Dot-Product Attention and Multi-Head Mechanics

The core operation is scaled dot-product attention. Given queries Q, keys K, values V (all matrices of shape [seq_len, d_k]), the attention output is softmax(Q·K^T / √d_k) · V. The division by √d_k prevents the dot products from growing too large, which would push the softmax into regions of extremely small gradients (saturation).

Multi-head attention: instead of one attention operation in d_model dimensions, project Q, K, V down to h lower-dimensional heads (each of dimension d_k = d_model / h), compute attention in parallel on each head, then concatenate and project back up. Each head learns different relationship types: some heads focus on local syntax (adjacent words), others on long-range dependencies, others on coreference (pronoun resolution).

In practice, h=8 for base model (d_model=512, d_k=64). The computational cost is the same as single-head attention because the total dimension is the same: h (d_k²) = d_model d_k. But multi-head adds a projection layer O(d_model²) after concatenation.

The paper found 8 heads performed best on translation; increasing to 16 gave marginal gains at higher compute cost. Don't chase more heads — the real wins come from better scaling, not more parallel subspaces.

One subtlety: the scaling factor. With d_k=64, √d_k=8, the dot products shrink by 8x. Without that, the variance of logits scales linearly with d_k. For d_k=512, logits have variance ~512, which pushes softmax nearly one-hot. Gradients become minuscule — your loss doesn't move. Production teams often forget this when increasing d_model and keeping n_heads constant (d_k grows).

A production reality: we once debugged a model where training loss flatlined at 4.3 for three days. The team tried different optimizers, learning rates, everything. The fix was one line: adding / math.sqrt(self.d_k) before softmax. The scaling factor was present in the paper's pseudocode but missing in the implementation. Three days of compute, gone. That's the kind of paper detail that separates working models from broken ones.

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        # Q, K, V: [batch, n_heads, seq_len, d_k]
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        return torch.matmul(attn, V)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        # 1. Linear projections
        Q = self.W_q(query)
        K = self.W_k(key)
        V = self.W_v(value)

        # 2. Split into heads: [batch, seq_len, n_heads, d_k] -> [batch, n_heads, seq_len, d_k]
        Q = Q.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

        # 3. Attention
        attn_output = self.scaled_dot_product_attention(Q, K, V, mask)

        # 4. Concatenate heads
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 5. Final projection
        return self.W_o(attn_output)

Positional Encoding — Giving Order to the Permutation-Invariant Transformer

The Transformer's attention mechanism is permutation-invariant: swapping two input tokens yields the same attention distribution over other tokens. This is a problem because language is fundamentally ordered — 'dog bites man' vs 'man bites dog' have opposite meanings. Positional encodings add information about each token's position in the sequence.

The original paper used sinusoidal encodings: PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) and PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model)). Each dimension of the positional encoding has a different wavelength, from 2π to 10000·2π cycles. This allows the model to attend to relative positions (because the encoding at position pos+k can be represented as a linear function of encoding at pos).

Sinusoidal encodings are not learned, so they can extrapolate to sequence lengths longer than those seen during training. Learned positional embeddings (trainable parameters) often perform better on fixed-length tasks but cannot generalise to longer sequences.

The encoding is added directly to the input embeddings: x = embedding + positional_encoding. Not concatenated — addition preserves the embedding dimension, concatenation would double it.

Here's the trap: if you forget positional encoding, the model still trains. Loss goes down. But you'll see flat predictions on any task requiring order. In the time-series incident above, the model literally predicted the mean of training values for every time step. The fix was adding positional encoding, not changing the model size or learning rate.

Another production pitfall: using learned embeddings and then hitting a sequence length longer than the pre-defined max_len during inference. The embedding matrix has fixed size; out-of-range positions throw IndexError. Always validate inference sequences against the embedding table size.

Modern practice has largely moved to Rotary Position Embedding (RoPE), used in Llama, Mistral, and GPT-NeoX. RoPE applies rotation matrices to Q and K based on position — it encodes relative position directly into the attention computation rather than adding a fixed vector. It extrapolates better than learned embeddings and has become the default for new LLM implementations.

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Create positional encoding matrix [max_len, d_model]
        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)  # [1, max_len, d_model]
        self.register_buffer('pe', pe)

    def forward(self, x):
        # x: [batch, seq_len, d_model]
        x = x + self.pe[:, :x.size(1), :]  # Add positional encoding to input embeddings
        return self.dropout(x)

# Check that different positions have distinct encodings
pos_enc = PositionalEncoding(512, max_len=10)
print("Position 0 encoding (first 8 dims):", pos_enc.pe[0, 0, :8])
print("Position 1 encoding (first 8 dims):", pos_enc.pe[0, 1, :8])
assert not torch.allclose(pos_enc.pe[0, 0, :8], pos_enc.pe[0, 1, :8])

Encoder-Decoder Stack and Masking

The original Transformer has an encoder (processes input sequence) and a decoder (generates output sequence). Each encoder layer has multi-head self-attention (no masking) + feed-forward network (FFN). Each decoder layer has masked self-attention (prevents looking at future tokens) + cross-attention (attends to encoder output) + FFN.

The encoder sees the entire input sequence simultaneously. Self-attention is unmasked — every token can attend to every other token in the input.

The decoder is autoregressive: when generating token i, it can only attend to positions 0..i-1. This is enforced with a causal mask: an upper triangular matrix of -inf that zeros out attention to future tokens.

Cross-attention in the decoder uses the encoder output as K and V, and the decoder's previous layer output as Q. This allows the decoder to focus on different parts of the input sequence for each generated output token.

The feed-forward network (FFN) is a simple two-layer MLP with ReLU: FFN(x) = max(0, xW1 + b1)W2 + b2. It operates per token independently (no interaction across positions). This gives the model additional capacity to transform the attention output before the next layer.

One mistake I've seen in production code: applying the causal mask to the cross-attention. Cross-attention should have no mask — the decoder can attend to any encoder position, including those 'ahead' in the encoder sequence. The causal mask only applies to decoder self-attention. Mixing them up leads to artificially constrained generation.

Another production issue: when using KV cache for inference, the mask changes shape. During training, the mask is [seq_len, seq_len]. During inference with KV cache, the mask becomes [1, cached_len+1] — only the new token needs to mask out future tokens it shouldn't see. Getting this shape wrong causes either information leak or all tokens generating identical output.

import torch
import torch.nn as nn

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff=2048, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Pre-LN (stabler than original Post-LN)
        x = x + self.dropout(self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), mask))
        x = x + self.dropout(self.feed_forward(self.norm2(x)))
        return x

class TransformerDecoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff=2048, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), nn.Dropout(dropout)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, encoder_output, causal_mask=None, cross_mask=None):
        # Masked self-attention
        x = x + self.dropout(self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), causal_mask))
        # Cross-attention (encoder output as K, V)
        x = x + self.dropout(self.cross_attn(self.norm2(x), encoder_output, encoder_output, cross_mask))
        # Feed-forward
        x = x + self.dropout(self.feed_forward(self.norm3(x)))
        return x

Training Dynamics: Residual Connections, LayerNorm, and Dropout

The Transformer's depth (6 layers in base, 12 in big) requires careful architectural choices to enable gradient flow. The paper uses three key components: residual connections, layer normalization, and dropout.

Residual connections (skip connections): each sublayer output is added to its input: output = LayerNorm(x + Sublayer(x)). This lets gradients flow directly through the network, preventing vanishing gradients in deep stacks. Without residuals, a 6-layer Transformer would be nearly untrainable.

Layer Normalization: normalizes activations across the feature dimension (d_model). Unlike BatchNorm, LayerNorm is independent of batch size and works for variable-length sequences. The original paper placed LayerNorm after the residual addition (Post-LN), but modern practice places it before (Pre-LN).

Dropout: applied to the output of each sublayer (before addition) and to attention weights. Dropout rate 0.1 is standard. Insufficient dropout causes overfitting within 2-3 epochs on small datasets; too much dropout (>0.3) slows convergence.

Learning rate schedule: the paper uses a warm-up of 4000 steps with linear increase to 0.0005, then decays proportionally to inverse square root of step count. Pre-LN often makes warmup unnecessary.

Here's a practical insight: if you see training loss spike around step 4000 (the peak of warmup), the learning rate is too high. Reduce peak LR or extend warmup to 8000 steps. If loss plateaus and refuses to drop, you likely have one of two issues: Post-LN with insufficient warmup, or missing scaling factor in attention.

Also, label smoothing of 0.1 was used in the paper. It helps prevent the model from becoming overconfident, which degrades generation quality. Many modern LLMs skip label smoothing — but for translation it was critical.

For production training at scale, mixed precision (FP16/BF16) is standard. The original paper used FP32, but modern implementations use automatic mixed precision (AMP) for 2x throughput with negligible accuracy loss. The one gotcha: loss scaling in FP16 can overflow if the loss spikes during warmup. BF16 (if your hardware supports it) eliminates this problem entirely and is now the default for LLM training.

import torch
import torch.nn as nn

# Example: Pre-LN vs Post-LN comparison
class EncoderLayerPostLN(nn.Module):
    """Original Post-LN implementation"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads)
        self.ffn = nn.Sequential(nn.Linear(d_model, 2048), nn.ReLU(), nn.Linear(2048, d_model))
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(0.1)

    def forward(self, x):
        x = self.norm1(x + self.dropout(self.self_attn(x, x, x)))
        x = self.norm2(x + self.dropout(self.ffn(x)))
        return x

class EncoderLayerPreLN(nn.Module):
    """Modern Pre-LN implementation"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads)
        self.ffn = nn.Sequential(nn.Linear(d_model, 2048), nn.ReLU(), nn.Linear(2048, d_model))
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(0.1)

    def forward(self, x):
        x = x + self.dropout(self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x)))
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x