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.

How Attention Actually Remembers Position

The 'Attention Is All You Need' paper introduced the Transformer, a model that processes sequences in parallel rather than step-by-step. The core mechanic is the attention mechanism, which computes weighted sums over all input positions. But because attention is permutation-invariant — it treats the input as a set, not a sequence — the model has no inherent sense of word order. Positional encoding solves this by injecting a unique signal into each token's embedding, typically using sine and cosine functions of different frequencies. This lets the model distinguish "dog bites man" from "man bites dog" without sequential processing.

In practice, positional encodings are added directly to the input embeddings at the bottom of the encoder and decoder stacks. The sine/cosine functions produce values between -1 and 1, and their wavelengths range from 2π to 10,000·2π. This design gives two useful properties: the encoding for position pos+k can be expressed as a linear function of the encoding for pos, which helps the model learn relative positions; and the varying frequencies let the model attend to both nearby and distant tokens. The encoding dimension matches the model dimension (typically 512), so the addition doesn't change the tensor shape.

Use positional encoding in any Transformer-based architecture processing sequential data — NLP, time series, or even image patches. It's not optional: without it, a Transformer treats "hello world" and "world hello" identically. In production systems, the choice of encoding (learned vs. fixed sinusoidal) rarely matters for performance, but fixed encodings generalize better to sequence lengths unseen during training. This is critical when deploying models that must handle longer inputs than those seen in training.

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

The Modern Transformer Family — BERT, GPT, and T5

The original Transformer introduced an encoder-decoder architecture. But the family has diverged into three dominant lineages: encoder-only (BERT), decoder-only (GPT), and encoder-decoder (T5). Each has a different pretraining objective, inference pattern, and use case.

BERT (Bidirectional Encoder Representations from Transformers) uses only the encoder stack with bidirectional self-attention. Pretrained with masked language modeling (MLM) and next sentence prediction (NSP). BERT excels at understanding tasks: classification, NER, QA (extractive). It cannot generate text natively. Scaling: base (12 layers, 110M params) to large (24 layers, 340M). RoBERTa improved by removing NSP and dynamic masking.

GPT (Generative Pre-trained Transformer) uses only the decoder stack with causal masking. Pretrained with autoregressive language modeling (predict next token). GPT excels at generation: summarization, translation (prompt-based), chatbot, code generation. Scaling: GPT-1 (117M), GPT-2 (1.5B), GPT-3 (175B). Inference uses KV cache.

T5 (Text-to-Text Transfer Transformer) uses the full encoder-decoder stack with a text-to-text framework. Every task is cast as text input → text output. Pretrained with span corruption (mask spans of tokens and predict them). T5 is the Swiss Army knife: can do classification, generation, translation, QA (abstractive) in one model. Scaling: T5-small (60M) to T5-11B. The text-to-text format simplifies deployment but costs more compute per token than decoder-only.

Choose BERT when you need classification or extraction. Choose GPT when you need open-ended generation. Choose T5 when you need a single model for many tasks and can afford higher latency.

Attention Type Decision Matrix

When implementing a Transformer, you have several attention variants to choose from. The wrong choice leads to inefficient compute or incorrect behaviour. This matrix helps you decide based on task and constraints.

Decision steps: 1. If sequence length ≤ 1024 and you need full context → use full self-attention with Flash Attention for memory savings. 2. If you need autoregressive generation → use causal attention with KV cache. 3. If you're building a translation/seq2seq model → use cross-attention in decoder (Q=decoder, K,V=encoder). 4. If sequence length > 4096 and compute budget is tight → consider sparse or linear attention. 5. If you need long-context (100k+) → use Flash Attention (full but tiled) or linear variants.

In practice, the original 'Attention Is All You Need' attention serves nearly all modern models up to 4096 tokens. Beyond that, Flash Attention has become the default for training (PyTorch 2.0+), and KV cache for inference. Sparse and linear attention have niche but growing adoption.

Visual KV Cache Walkthrough for Production Serving

Autoregressive decoding is expensive: generating token by token, each step recomputes the entire sequence's attention. The KV cache eliminates this by storing the key and value vectors from previous steps. At step t, instead of computing Q,K,V for all t tokens, we only compute for the new token and retrieve cached K,V for positions 0..t-1.

Step-by-step flow:

1. Initial step (t=0): The decoder receives the start token. Compute Q₀, K₀, V₀ for token 0 from the first decoder layer. Store K₀, V₀ in cache. Compute attention only over token 0 (no mask needed). Output token 1.
2. Step t=1: The decoder processes token 1. Compute Q₁, K₁, V₁ for token 1. Retrieve cached K₀, V₀. Concatenate: K_all = [K₀, K₁], V_all = [V₀, V₁]. Compute attention with Q₁ over K_all, V_all. Apply causal mask: [1, 1] for allowed, [0, 1] for future? Actually, at this step the valid positions are 0 and 1 (since we have two tokens). The mask is upper triangular -inf for future positions. But since we are at step 1, the only 'future' is position >1 which doesn't exist. The mask shape is [1, 2] (one query, two keys). No masking needed beyond that. Store K₁, V₁ in cache.
3. Step t=n: Repeat. Cache grows by 2 (K and V) per layer per step.

The diagram below shows the flow for a single decoder layer:

``mermaid graph LR subgraph Step 0 A0[Token 0] --> Q0[Compute Q₀] --> Att0[Attention: Q₀·K₀] A0 --> K0[Compute K₀, V₀] --> Cache0[(Cache: K₀, V₀)] Att0 --> Out0[Output token 1] end subgraph Step 1 A1[Token 1] --> Q1[Compute Q₁] Cache0 --> Concat[Concatenate K₀,V₀ with K₁,V₁] A1 --> K1[Compute K₁, V₁] --> Concat Q1 --> Att1[Attention: Q₁·[K₀,K₁]] Concat --> Att1 Att1 --> Out1[Output token 2] end ``

Memory cost: For a model with 24 layers, d_model=1024, half precision (2 bytes), each step adds 2 (K,V) 24 layers 1024 * 1 token = ~96 KB per step. For 2048 tokens, cache = ~192 MB. This is manageable. Without cache, each step would recompute all previous tokens, costing O(n²) compute: for 2048 tokens, that's ~4 million attention computations per step vs ~2000 with cache.

Implementation note: The cache is stored per layer, typically as two lists or tensors. At each step, we slice the latest Q from the decoder input, run through the decoder layer, and append K,V to the cache. The attention function must handle variable-length K,V.

See the PyTorch implementation below for how this works in code.

import torch
import torch.nn as nn

class DecoderLayerWithCache(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads)
        # ... cross-attn, ff, norms are similar
    
    def forward(self, x, encoder_output, cache=None, step=None):
        # x: [batch, 1, d_model] — only the new token
        # cache: dictionary with 'k' and 'v' for this layer
        # step: current generation step
        
        # Self-attention with KV cache
        Q = self.self_attn.W_q(x)
        K = self.self_attn.W_k(x)
        V = self.self_attn.W_v(x)
        
        if cache is not None:
            # Append new K,V to cache
            K = torch.cat([cache['k'], K], dim=1)  # [batch, t+1, d_model]
            V = torch.cat([cache['v'], V], dim=1)
        
        # Update cache
        new_cache = {'k': K, 'v': V}
        
        # Split into heads
        batch_size = Q.size(0)
        Q = Q.view(batch_size, -1, self.self_attn.n_heads, self.self_attn.d_k).transpose(1,2)  # [batch, n_heads, 1, d_k]
        K = K.view(batch_size, -1, self.self_attn.n_heads, self.self_attn.d_k).transpose(1,2)  # [batch, n_heads, t+1, d_k]
        V = V.view(batch_size, -1, self.self_attn.n_heads, self.self_attn.d_k).transpose(1,2)
        
        # Create causal mask: only the last query sees all previous tokens
        seq_len = K.size(2)
        mask = torch.triu(torch.ones(1, seq_len, dtype=torch.bool, device=x.device), diagonal=1)  # [1, seq_len]
        # Only the last row (current step) is valid
        mask = mask[step:step+1, :]  # [1, seq_len]
        
        attn_output = self.self_attn.scaled_dot_product_attention(Q, K, V, mask)
        attn_output = attn_output.transpose(1,2).contiguous().view(batch_size, -1, self.self_attn.d_model)
        # ... rest of forward pass
        return attn_output, new_cache

# Inference loop
cache = [None] * num_layers
output = start_token
for step in range(max_len):
    for layer_idx in range(num_layers):
        output, cache[layer_idx] = decoder_layers[layer_idx](output, encoder_output, cache[layer_idx], step)
    next_token = sample(output)
    output = embedding(next_token)

Keras/TensorFlow Implementation Snippets

While PyTorch dominates research, many production systems use TensorFlow/Keras for serving (via TF Serving, SageMaker, etc.). Here's how to implement the core Transformer components in Keras. The principles are identical to the PyTorch code above, but the API differs.

Multi-Head Attention in Keras:

```python import tensorflow as tf from tensorflow.keras import layers

class MultiHeadAttention(layers.Layer): def __init__(self, d_model, num_heads, dropout=0.1): super().__init__() self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads self.W_q = layers.Dense(d_model) self.W_k = layers.Dense(d_model) self.W_v = layers.Dense(d_model) self.W_o = layers.Dense(d_model) self.dropout = layers.Dropout(dropout) def call(self, query, key, value, mask=None, training=False): batch_size = tf.shape(query)[0] Q = self.W_q(query) # (batch, seq_len, d_model) K = self.W_k(key) V = self.W_v(value) # Reshape to (batch, seq_len, num_heads, d_k) and transpose to (batch, num_heads, seq_len, d_k) Q = tf.transpose(tf.reshape(Q, (batch_size, -1, self.num_heads, self.d_k)), perm=[0,2,1,3]) K = tf.transpose(tf.reshape(K, (batch_size, -1, self.num_heads, self.d_k)), perm=[0,2,1,3]) V = tf.transpose(tf.reshape(V, (batch_size, -1, self.num_heads, self.d_k)), perm=[0,2,1,3]) # Scaled dot-product attention scores = tf.matmul(Q, K, transpose_b=True) / tf.sqrt(tf.cast(self.d_k, tf.float32)) if mask is not None: scores += (mask * -1e9) attn_weights = tf.nn.softmax(scores, axis=-1) attn_weights = self.dropout(attn_weights, training=training) output = tf.matmul(attn_weights, V) # (batch, num_heads, seq_len, d_k) # Concatenate heads output = tf.transpose(output, perm=[0,2,1,3]) # (batch, seq_len, num_heads, d_k) output = tf.reshape(output, (batch_size, -1, self.d_model)) return self.W_o(output) ```

Positional Encoding in Keras:

``python class PositionalEncoding(layers.Layer): def __init__(self, max_len, d_model): super().__init__() self.pos_encoding = self._create_encoding(max_len, d_model) def _create_encoding(self, max_len, d_model): positions = tf.range(max_len, dtype=tf.float32)[:, tf.newaxis] div_terms = tf.exp(tf.range(0, d_model, 2, dtype=tf.float32) (-tf.math.log(10000.0) / d_model)) pe = tf.zeros((max_len, d_model)) pe[:, 0::2] = tf.sin(positions div_terms) pe[:, 1::2] = tf.cos(positions * div_terms) return pe[tf.newaxis, :, :] # (1, max_len, d_model) def call(self, x): return x + self.pos_encoding[:, :tf.shape(x)[1], :] ``

Transformer Encoder Layer (Pre-LN) in Keras:

``python class TransformerEncoderLayer(layers.Layer): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.attention = MultiHeadAttention(d_model, num_heads, dropout) self.ffn = tf.keras.Sequential([ layers.Dense(d_ff, activation='relu'), layers.Dense(d_model), layers.Dropout(dropout) ]) self.norm1 = layers.LayerNormalization(epsilon=1e-6) self.norm2 = layers.LayerNormalization(epsilon=1e-6) self.dropout = layers.Dropout(dropout) def call(self, x, mask=None, training=False): # Pre-LN attn_out = self.attention(self.norm1(x), self.norm1(x), self.norm1(x), mask, training=training) x = x + self.dropout(attn_out, training=training) ffn_out = self.ffn(self.norm2(x), training=training) x = x + ffn_out return x ``

These snippets mirror the PyTorch versions exactly. The key differences: Keras uses layers.Dense instead of nn.Linear, LayerNormalization instead of nn.LayerNorm, and tf.matmul with transpose option instead of torch.matmul. The call method accepts a training argument for dropout behaviour.

Note: TensorFlow's tf.keras.layers.MultiHeadAttention (built-in) is an efficient implementation with relative attention bias. Use that for production. The custom code above is for learning the internals.

import tensorflow as tf
from tensorflow.keras import layers

class TransformerEncoderLayer(layers.Layer):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attention = layers.MultiHeadAttention(num_heads, d_model, dropout=dropout)
        self.ffn = tf.keras.Sequential([
            layers.Dense(d_ff, activation='relu'),
            layers.Dense(d_model),
            layers.Dropout(dropout)
        ])
        self.norm1 = layers.LayerNormalization(epsilon=1e-6)
        self.norm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout = layers.Dropout(dropout)
    
    def call(self, x, mask=None, training=False):
        # Pre-LN
        attn_out = self.attention(query=self.norm1(x), value=self.norm1(x), key=self.norm1(x),
                                 attention_mask=mask, training=training)
        x = x + self.dropout(attn_out, training=training)
        ffn_out = self.ffn(self.norm2(x), training=training)
        x = x + ffn_out
        return x

# Usage
encoder = TransformerEncoderLayer(d_model=512, num_heads=8, d_ff=2048)
x = tf.random.normal((4, 32, 512))
output = encoder(x, mask=None, training=True)
print(output.shape)  # (4, 32, 512)

Why Hard Attention is a Production Nightmare (and Soft Attention Saves You)

Most blog posts treat attention types like a buffet — pick what looks good. In production, you pick soft attention or you pick a fire drill. Hard attention selects discrete input positions via sampling. That's non-differentiable. You can't backprop through it without REINFORCE or some other high-variance gradient estimator. Training becomes a coin flip. Soft attention uses a weighted sum over all inputs, with weights from a softmax. Differentiable end-to-end. Stable gradients. Predictable loss curves.

Hard attention sounds appealing for efficiency — only look at 10% of the input. But the variance in training kills any throughput gain. Every team I've seen attempt hard attention for sequence tasks reverts to soft within two sprints. The Transformer paper uses scaled dot-product attention exclusively, which is a soft variant. Follow that lead.

The practical choice: soft attention for training, and if you must prune for inference, use a separate sparsity technique like top-k after training. Don't bake non-differentiability into your architecture unless you enjoy debugging NaN gradients at 3 AM.

// io.thecodeforge — ml-ai tutorial

import torch
import torch.nn.functional as F

def soft_attention(query, keys, values):
    # Differentiable. Use this.
    scores = torch.matmul(query, keys.T) / (keys.size(-1) ** 0.5)
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, values)

def hard_attention(query, keys, values, num_samples=2):
    # Non-differentiable. Avoid in training.
    scores = torch.matmul(query, keys.T)  # no scaling — who cares, it's broken anyway
    probs = F.softmax(scores, dim=-1)
    indices = torch.multinomial(probs, num_samples, replacement=False)
    selected_values = values[indices]
    return selected_values.mean(dim=0)  # crude approximation

q = torch.randn(4, 8)
k = torch.randn(6, 8)
v = torch.randn(6, 32)

soft_out = soft_attention(q, k, v)
print(f"Soft output shape: {soft_out.shape}")  # torch.Size([4, 32])

hard_out = hard_attention(q, k, v)
print(f"Hard output shape: {hard_out.shape}")  # torch.Size([32])

The Encoder-Decoder Handshake — Where Attention Actually Connects

Most diagrams show encoder and decoder as two blocks with an arrow labeled "attention." That arrow hides the critical interface. In the Transformer, the encoder outputs a sequence of key-value pairs. The decoder generates queries from its own hidden states. The cross-attention layer in the decoder computes attention between those decoder queries and the encoder's keys and values. This is not self-attention. It's encoder-decoder attention, and it's the bridge that lets the decoder "look at" the input.

Why this matters for debugging: If your translation model outputs garbage, check the cross-attention weights first. A common failure mode is the decoder attending to the wrong input tokens — often the start token or padding. Visualize the attention matrix. If it's a flat, uniform distribution, the model isn't learning alignment. If it peaks on the wrong positions, your positional encodings might be misaligned or your dataset has alignment errors.

The encoder doesn't attend to the decoder. The decoder attends to the encoder. That asymmetry is intentional. The encoder builds a rich representation of the input. The decoder reconstructs the output by selectively focusing on that representation. Flip the direction and you break causality.

// io.thecodeforge — ml-ai tutorial

import torch
import torch.nn as nn

class CrossAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.query_proj = nn.Linear(d_model, d_model)
        self.key_proj = nn.Linear(d_model, d_model)
        self.value_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, decoder_hidden, encoder_output):
        Q = self.query_proj(decoder_hidden)  # [batch, tgt_len, d_model]
        K = self.key_proj(encoder_output)    # [batch, src_len, d_model]
        V = self.value_proj(encoder_output)  # [batch, src_len, d_model]

        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / (Q.size(-1) ** 0.5)
        attn_weights = torch.softmax(attn_scores, dim=-1)
        return self.out_proj(torch.matmul(attn_weights, V))

batch, tgt_len, src_len, d_model = 2, 5, 7, 512
cross_attn = CrossAttention(d_model, 8)
decoder_hidden = torch.randn(batch, tgt_len, d_model)
encoder_output = torch.randn(batch, src_len, d_model)
output = cross_attn(decoder_hidden, encoder_output)
print(f"Output shape: {output.shape}")  # torch.Size([2, 5, 512])

Where Attention Fails — Four Production-Ready Workarounds

Attention is not magic. It has known failure modes that crash production systems. Here are four, and how to beat each one.

1. Quadratic complexity. Self-attention is O(n²) in sequence length. At 512 tokens, fine. At 8192, your GPU OOMs. Fix: Use sparse attention patterns (like Longformer's sliding window) or linear attention variants (like Linformer or Performer). Production rule: never let sequence length grow without a scaling plan.
2. Positional confusion. The vanilla Transformer is permutation-invariant. Without positional encoding, "cat sat" and "sat cat" produce identical representations. Absolute positional encodings fix this but don't generalize to unseen lengths. Fix: Use relative positional encodings (like T5's bias) or rotary embeddings (RoPE). RoPE is my go-to — it's clean and extrapolates well.
3. Attention collapse. In deep transformers, attention heads often converge to nearly identical patterns, reducing effective capacity. This is "attention redundancy." Fix: Use regularization like attention dropout or dedicated loss terms that encourage head diversity. I've also seen success with initializing heads with different temperature parameters.
4. Over-attention to padding. Models learn to attend to padding tokens because they're frequent. This dilutes signal. Fix: Explicitly mask padding tokens by setting their attention scores to -inf before softmax. Standard practice, but I still see new codebases miss this and wonder why validation loss plateaus.

Each of these has bitten my teams in production. Don't let them bite yours.

// io.thecodeforge — ml-ai tutorial

import torch
import torch.nn.functional as F

def masked_self_attention(query, key, value, mask=None):
    # mask: [batch, seq_len] with 1 for valid, 0 for padding
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / (d_k ** 0.5)
    if mask is not None:
        # Expand mask to [batch, 1, seq_len, seq_len]
        mask = mask.unsqueeze(1).unsqueeze(2)
        scores = scores.masked_fill(mask == 0, float('-inf'))
    weights = F.softmax(scores, dim=-1)
    return torch.matmul(weights, value)

batch, seq_len, d_model = 4, 128, 64
query = torch.randn(batch, seq_len, d_model)
key = torch.randn(batch, seq_len, d_model)
value = torch.randn(batch, seq_len, d_model)
# Mask: first 10 tokens are padding
mask = torch.ones(batch, seq_len)
mask[:, :10] = 0

output = masked_self_attention(query, key, value, mask)
print(f"Output shape: {output.shape}")  # torch.Size([4, 128, 64])

# Check: no nan or inf
print(f"Has NaN: {torch.isnan(output).any().item()}")  # False