Transformer Architecture from Scratch: Attention, Encoders and Decoders

The original Transformer paper, "Attention Is All You Need," replaced a decade of recurrence and convolution with a single idea: let every token directly attend to every other token. When I first re-implemented the architecture from scratch for a production translation system in 2019, I underestimated how many subtle details matter beyond the famous attention formula. The difference between a model that converges and one that diverges to NaN often comes down to initialization of the query-key projection, the placement of LayerNorm, or how you implement causal masking. This article reconstructs the Transformer encoder-decoder piece by piece, with the practical decisions I have learned to get right after training and debugging these models at scale.

Why Self-Attention Replaced Recurrence for Sequence Modeling

Recurrent networks process tokens sequentially, which creates two fundamental bottlenecks. First, the path length between distant tokens grows linearly with sequence length, making long-range dependencies difficult to learn due to vanishing gradients. Second, sequential computation prevents parallelization across the time dimension, which throttles training throughput on modern accelerators.

Self-attention solves both by computing pairwise interactions in a single matrix multiplication. For a sequence of length n, each position can directly reference any other position with O(1) path length. The computational cost shifts to O(n²·d) where d is the model dimension, which is more expensive for very long sequences but highly parallelizable on GPUs. In my experience, this trade-off is overwhelmingly favorable for sequences up to 8k tokens, which covers most NLP and many vision tasks.

The Shift from Sequential State to Global Context

In an RNN, the hidden state h_t is a compressed summary of x_1...x_t. Self-attention discards this compression entirely. Instead, it keeps the full set of representations and learns to route information dynamically. This is why Transformers scale better with data and parameters: the model does not have to squeeze all prior context into a fixed-size vector. The attention weights themselves become interpretable routing decisions, and I have often used attention heatmaps during debugging to catch tokenizer or masking errors early.

If you are coming from convolutional architectures, the comparison is instructive as well. CNNs build receptive fields hierarchically through stacked local kernels, while attention builds a global receptive field in a single layer. That distinction explains why Vision Transformers eventually outperformed CNNs on large datasets, a trend analyzed in Modern CNNs: ResNet, EfficientNet and Vision Transformers Compared.

Dissecting Scaled Dot-Product Attention and Multi-Head Projections

At the core of every Transformer is scaled dot-product attention. Given queries Q, keys K, and values V, each of shape [batch, seq_len, d_k], the operation is:

Attention(Q, K, V) = softmax(QKᵀ / √d_k) V

The scaling factor √d_k is not optional. Without it, for large d_k the dot products grow in magnitude, pushing the softmax into regions with extremely small gradients. Vaswani et al. introduced this specifically to counteract that effect, and implementations that omit it consistently show slower convergence in my experiments. See the formal definition in the PyTorch Documentation for scaled_dot_product_attention, which also details memory-efficient kernel fusions.

Multi-head attention runs this operation h times in parallel with different learned projections. Instead of a single attention function on d_model = 512, we project to h = 8 heads of d_k = d_v = 64, apply attention independently, then concatenate and project back. This allows each head to specialize: in trained models, I regularly see distinct heads focusing on syntactic relations, positional adjacency, and long-range coreference.

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    # Q, K, V: [batch, heads, seq_len, d_k]
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    
    if mask is not None:
        # mask shape: [batch, 1, seq_len, seq_len] or [batch, 1, 1, seq_len]
        # masked positions are set to -inf before softmax
        scores = scores.masked_fill(mask == 0, float('-inf'))
    
    attn_weights = F.softmax(scores, dim=-1)
    # In my experience, dropout on attention weights (0.1) is critical regularization
    attn_weights = F.dropout(attn_weights, p=0.1, training=True)
    return torch.matmul(attn_weights, V), attn_weights

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model=512, n_heads=8):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_k = d_model // n_heads
        self.n_heads = n_heads
        self.q_proj = torch.nn.Linear(d_model, d_model, bias=False)
        self.k_proj = torch.nn.Linear(d_model, d_model, bias=False)
        self.v_proj = torch.nn.Linear(d_model, d_model, bias=False)
        self.o_proj = torch.nn.Linear(d_model, d_model, bias=False)

    def forward(self, query, key, value, mask=None):
        B, L, _ = query.shape
        Q = self.q_proj(query).view(B, L, self.n_heads, self.d_k).transpose(1, 2)
        K = self.k_proj(key).view(B, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = self.v_proj(value).view(B, -1, self.n_heads, self.d_k).transpose(1, 2)
        attn_out, _ = scaled_dot_product_attention(Q, K, V, mask)
        attn_out = attn_out.transpose(1, 2).contiguous().view(B, L, -1)
        return self.o_proj(attn_out)

What I Initialize Differently for Stable Attention

I've found that Xavier uniform initialization for the Q and K projections often leads to overly large initial attention logits. Using a smaller gain or initializing the final output projection to zero (or near-zero) helps the residual branch start as an identity, which stabilizes early training. Many open-source implementations in the Hugging Face Documentation now use this trick. I also tie the scaling of the residual addition by multiplying the sub-layer output by 1/√(2·N) where N is the number of layers, which helps when scaling beyond 24 layers.

Constructing the Transformer Encoder Block: Residuals, LayerNorm, and Feed-Forward Networks

The encoder is a stack of identical blocks, each containing two sub-layers: multi-head self-attention and a position-wise feed-forward network (FFN), both wrapped with residual connections and normalization.

The original paper used Post-LN (LayerNorm after the residual addition): x + Sublayer(LayerNorm(x)) vs. x + Sublayer(x) then LayerNorm. In practice, Post-LN is more sensitive to learning rate and often diverges without warmup. Most production systems I build now use Pre-LN: LayerNorm before the sub-layer, with a final LayerNorm after the stack. Pre-LN gives more stable gradients at initialization, allowing larger learning rates without warmup, though it can slightly underperform Post-LN at convergence if tuned carefully. For a deeper look at why normalization placement and learning rate schedules interact, see Training Deep Networks: Learning Rate Schedules, Batch Normalization and Gradient Clipping.

The FFN is deceptively simple but accounts for two-thirds of the parameters: FFN(x) = max(0, xW₁ + b₁)W₂ + b₂, where W₁ expands from d_model to d_ff = 2048 (4×) and W₂ projects back. I treat this as a key-value memory where the first layer detects patterns and the second writes them. Replacing ReLU with GELU or SwiGLU consistently improves performance by 0.5-1% in my language modeling runs, at minimal cost.

class EncoderLayer(torch.nn.Module):
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads)
        self.norm1 = torch.nn.LayerNorm(d_model)
        self.norm2 = torch.nn.LayerNorm(d_model)
        self.ffn = torch.nn.Sequential(
            torch.nn.Linear(d_model, d_ff),
            torch.nn.GELU(),
            torch.nn.Dropout(dropout),
            torch.nn.Linear(d_ff, d_model)
        )
        self.dropout = torch.nn.Dropout(dropout)

    def forward(self, x, src_mask=None):
        # Pre-LN architecture
        residual = x
        x = self.norm1(x)
        x = self.self_attn(x, x, x, mask=src_mask)
        x = self.dropout(x)
        x = residual + x

        residual = x
        x = self.norm2(x)
        x = self.ffn(x)
        x = self.dropout(x)
        x = residual + x
        return x

Residuals Are Not Optional

Without residual connections, a 6-layer Transformer loses gradient flow within the first few thousand steps. With them, I have trained 48-layer encoders from scratch without auxiliary losses. The critical detail is to apply dropout to the sub-layer output *before* the addition, not to the sum itself, otherwise you randomly zero the residual path.

The Decoder Stack: Masked Self-Attention and Cross-Attention Mechanics

The decoder adds two complications: it must be autoregressive, and it must condition on the encoder output. Each decoder layer therefore has three sub-layers: masked self-attention, cross-attention, and the FFN.

Masked self-attention prevents position i from attending to positions j > i. This is implemented by adding a causal mask with -inf in the upper triangle before softmax. A common bug I see is creating the causal mask on CPU with the wrong dtype, which then fails to broadcast on GPU with mixed precision. Always create it as a boolean tensor and expand it to [batch, 1, seq_len, seq_len] to match the attention scores.

Cross-attention is where the decoder queries the encoder. Here, Q comes from the previous decoder layer, while K and V come from the final encoder output. The encoder output is not masked (except for padding), so each decoder position can attend to the entire source sequence. This is the mechanism that enables sequence-to-sequence tasks like translation and summarization.

Teacher Forcing and Exposure Bias in Practice

During training, decoder inputs are the target sequence shifted right by one, with ground-truth tokens fed in (teacher forcing). At inference, the model feeds its own predictions autoregressively. This mismatch causes exposure bias. In my experience, scheduled sampling rarely helps for large Transformers; instead, I rely on label smoothing (ε=0.1) and beam search with coverage penalties to mitigate repetition and over-confidence. For pure language modeling, the encoder is removed entirely and the model becomes decoder-only, as in GPT.

Positional Encodings and Input Embeddings That Actually Generalize

Because self-attention is permutation-equivariant, it has no notion of order without positional information. Input representations are therefore token embeddings plus positional encodings.

The original sinusoidal encoding PE(pos,2i) = sin(pos/10000^(2i/d_model)), PE(pos,2i+1) = cos(pos/10000^(2i/d_model)) was chosen because it allows the model to extrapolate to longer sequences and to learn relative positions via linear transformations. Learned positional embeddings perform similarly for fixed-length tasks but fail to generalize beyond the trained length. For sequences longer than 512, I have shifted entirely to Rotary Positional Embeddings (RoPE) or ALiBi, which extrapolate far better and require no additional parameters.

Embeddings themselves need careful scaling. I scale token embeddings by √d_model before adding positional encodings, which balances their variance and prevents the positional signal from drowning out early in training. Weight tying between the input embedding matrix and the final output projection (before softmax) reduces parameters and consistently improves BLEU/perplexity in my runs.

Handling Padding and Variable Lengths Correctly

Padding masks must be applied in every attention call. I create a key padding mask of shape [batch, seq_len] where 1 indicates real tokens and 0 indicates padding, then unsqueeze it to broadcast over queries. Failure to mask padding in cross-attention will allow the decoder to attend to meaningless encoder positions, which often manifests as spuriously high attention weights on the first or last padded token during inspection.

Training Instabilities and Optimization Tactics for Transformers at Scale

Transformers are notoriously sensitive to optimization choices. The original schedule — linear warmup for 4,000 steps to a peak learning rate of ~7e-4 with inverse square-root decay — remains a strong baseline. Without warmup, Adam's second moment estimates are inaccurate in early steps, leading to large updates that can permanently damage initialization.

I now use AdamW with β1=0.9, β2=0.98, ε=1e-9, weight decay 0.01, and gradient clipping at 1.0. Clipping by global norm is essential; attention logits can spike when a batch contains unusually long sequences, producing gradient norms an order of magnitude larger than average. I also enable mixed-precision (FP16/BF16) with loss scaling — but keep the softmax and LayerNorm in FP32, as I have seen overflow in FP16 attention cause NaN losses that are difficult to trace. Insights from arXiv Machine Learning papers on large-scale training repeatedly confirm that numerical stability, not model capacity, is the primary barrier past 1B parameters.

Batch size matters more than most practitioners expect. Transformers benefit from large effective batch sizes (e.g., 25k-400k tokens per batch) with gradient accumulation. Smaller batches increase variance in the attention statistics and often require lower learning rates to remain stable.

def transformer_lr_schedule(step, d_model=512, warmup_steps=4000):
    # Original "Attention Is All You Need" schedule
    return (d_model ** -0.5) * min(step ** -0.5, step * (warmup_steps ** -1.5))

# Modern alternative I use with AdamW
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, betas=(0.9, 0.98), 
                              eps=1e-9, weight_decay=0.01)

# Framework choice influences distributed training throughput
# See detailed benchmarks in PyTorch vs TensorFlow comparison
# Mixed precision + gradient clipping is non-negotiable at scale
scaler = torch.cuda.amp.GradScaler()
for batch in dataloader:
    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
        logits = model(batch.src, batch.tgt_input)
        loss = F.cross_entropy(logits.view(-1, vocab_size), 
                               batch.tgt_output.view(-1), 
                               ignore_index=pad_id, label_smoothing=0.1)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

When to Use Which Framework for Implementation

In my experience, the choice between frameworks rarely affects model quality but significantly affects iteration speed. For research and custom attention variants, the flexibility documented in PyTorch vs TensorFlow in 2026: Which Framework for Which Use Case aligns with what I see in production: PyTorch dominates for Transformer research due to its dynamic graph and torch.compile support for fused attention kernels.

Frequently Asked Questions