10.4 Transformer Architecture and Position: Parallelism in Training, Not in Autoregressive Generation

The attention workbench now enables each position to directly access others, but a single attention mechanism is not enough to constitute a Transformer. The production line still requires position embeddings, residual connections, normalization layers, position-wise feed-forward networks, masking, and task-specific heads. The information flow between encoder and decoder also differs.

The value of the Transformer is not in eliminating all sequence dependencies. During training, known token positions can be computed in parallel. However, when the decoder generates the next token, it remains constrained by the causal, autoregressive generation order.

Learning Objectives

• Assemble a pre-norm/post-norm Transformer block;
• Distinguish between encoder-only, decoder-only, and encoder–decoder architectures;
• Understand position embeddings and causal, padding, and cross masks;
• Analyze the costs of training parallelism, KV cache generation, and long sequence processing;
• Write a PyTorch encoder block with correct shape and mask handling.

1. Token Representation = Content + Position

Discrete token IDs are first looked up in an embedding matrix:

$$ E\in\mathbb R^{|V|\times d_{model}}. $$

The sequence tensor [B,T,D] also requires a position signal, otherwise self-attention cannot distinguish token order. Input components may include:

$$ x_t = \text{tokenEmbed}(t) + \text{position}(t) + \text{segment/type}(t). $$

Adding these components is not the only approach; relative or rotary position embeddings encode positional information directly into the attention scores or the QK matrix relationships.

Even if the padding token has a zero embedding, the sum of position signals and bias may result in a non-zero value. Therefore, a padding mask must be used to explicitly exclude such tokens during loss computation and aggregation.

2. Trade-offs of the Position Method

Fixed sinusoidal

Different frequency sin/cos pairs produce deterministic absolute positions without adding parameters to the position table. The ability to compute longer positions does not mean the model has learned external extrapolation.

Learned absolute embedding

A vector is learned for each position. Simple, but typically constrained by a maximum training length or table boundary.

Relative position bias

Adjusts attention scores based on the relative distance between query and key, more directly encoding distance and direction.

Rotary position embedding

Applies position-dependent rotations to the components of Q and K, embedding relative positional information into the dot product. Multiple implementations exist for frequency scaling and long-context extension; RoPE should not be treated as a parameter-free solution.

The choice of position scheme affects extrapolation, caching, fine-tuning, and serving, and must be versioned alongside checkpoints and configurations.

3. A Transformer Block Has Two Types of Sublayers

Multi-head self-attention

Exchanges information across positions.

Position-wise FFN

Applies the same MLP independently to each position:

$$ FFN(x)=W_2\phi(W_1x+b_1)+b_2. $$

Attention mixes token dimensions, while FFN mixes feature/channel dimensions. Both employ residual connections and normalization; dropout, activation, and gating mechanisms vary depending on the architecture.

4. Post-norm vs Pre-norm

The original post-norm form is approximately:

$$ x' = LN(x + Attention(x)), $$

$$ y = LN(x' + FFN(x')). $$

A common pre-norm variant is:

$$ x' = x + Attention(LN(x)), $$

$$ y = x' + FFN(LN(x')). $$

Pre-norm often improves gradient flow in deep models, but the two forms represent and train differently; they cannot be arbitrarily switched during training or checkpoint loading. Dropout, scaling, and initialization settings for the residual branch are also part of the architectural definition.

5. PyTorch Pre-norm Encoder Block

import torch
from torch import nn

class EncoderBlock(nn.Module):
    def __init__(self, d_model, heads, d_ff, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attention = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=heads,
            dropout=dropout,
            batch_first=True,
        )
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
        )
        self.residual_dropout = nn.Dropout(dropout)

    def forward(self, x, padding_mask=None, causal_mask=None):
        normalized = self.norm1(x)
        attended, _ = self.attention(
            normalized,
            normalized,
            normalized,
            key_padding_mask=padding_mask,  # [B,T], True indicates ignore
            attn_mask=causal_mask,
            need_weights=False,
        )
        x = x + self.residual_dropout(attended)
        x = x + self.residual_dropout(self.ffn(self.norm2(x)))
        return x

block = EncoderBlock(d_model=128, heads=4, d_ff=512)
x = torch.randn(8, 20, 128)
padding = torch.zeros(8, 20, dtype=torch.bool)
assert block(x, padding_mask=padding).shape == x.shape

In the current MultiheadAttention Boolean padding mask, True indicates ignored tokens. The shape and dtype of the causal mask, as well as the behavior when passing two masks simultaneously, should be tested against a fixed version of PyTorch; semantic differences across SDPA API versions may exist.

This reference block serves as a foundational understanding and does not represent the full optimization suite of modern models (e.g., RMSNorm, SwiGLU, GQA, fused kernels, etc.).

6. Three Transformer Families

Encoder-only

Typically enables bidirectional self-attention, producing contextual representations for each token. Suitable for tasks such as classification, retrieval encoding, token labeling, and masked prediction.

Decoder-only

Uses causal self-attention, allowing only access to the current and previous tokens, and trains on next-token prediction. Ideal for autoregressive language modeling and generation.

Encoder–decoder

The encoder processes the full source sequence; the decoder employs causal self-attention and queries the encoder’s outputs via cross-attention. Best suited for translation, summarization, and conditional generation.

"BERT/GPT/T5" aren't just different in masking, they differ in objective functions, tokenizers, normalization methods, positional encoding, data distributions, and training strategies.

7. Decoder Training Can Be Parallel, Generation Still Serial

During training, the full target sequence is known, so input and label sequences can be shifted and masked with a causal mask, enabling a single matrix computation to compute next-token losses across all positions. This is a key parallelization advantage of Transformers over recurrent training approaches.

In autoregressive generation, the token at position $t+1$ does not exist and must wait until the token at position $t$ is sampled or selected before proceeding. While parallelism across batch, heads, and layers remains, token steps are inherently dependent.

"The training is hundreds of times faster" is not a guaranteed architectural benefit; performance depends on sequence length, batch size, hardware, kernel efficiency, RNN baseline performance, communication overhead, and memory usage.

8. KV Cache Avoids Redundant Projection of Historical Context

In decoder generation, recalculating the entire prefix's K/V values at each step is highly inefficient. Instead, KV cache stores the historical keys and values per layer. When a new token is processed, only the new Q/K/V values are computed, and the query attends to the cached K/V data.

Benefits and trade-offs:

• Avoids redundant computation of historical K/V projections;
• Each token's attention still scales with the growing context length;
• Cache memory grows approximately proportionally to layers × sequence length × KV heads × head dimension × data type;
• Changes in beam or batch size, GQA/MQA configurations, or offloading/quantization affect memory usage;
• Position indices, masking, and cache eviction must be implemented correctly.

The cache represents the serving state of the model, and improper handling (such as tenant isolation failures, TTL mismatches, or model version incompatibilities) can lead to serious errors.

9. Long-Sequence Cost

The computational cost of dense self-attention typically scales as $O(T^2)$, while projection and feed-forward network (FFN) operations scale at $O(TD^2)$. This bottleneck varies with sequence length $T$, model dimension $D$, and hardware capabilities.

Strategies for long-context scenarios include:

• Flash/memory-efficient exact kernels: reduce memory I/O and materialization overhead;
• local/sliding-window/block-sparse attention: limit the number of connections;
• low-rank, kernel, or linear approximations: provide approximate attention computation;
• recurrence, memory compression, or retrieval mechanisms: alter the source of information;
• chunking: introduces challenges due to cross-chunk dependencies.

Claiming support for 128k context only indicates that the interface or training configuration allows such lengths, it does not prove that the model can reliably retrieve or reason at every position. Validation is required through real-world evaluations of needle queries, multi-hop reasoning, positional awareness, and the "lost-in-the-middle" scenario.

10. Mask is a Security Boundary for Information

At minimum, the following must be tested:

• causal: Changing future tokens should not affect past logits;
• padding: Modifying padded token IDs should not influence valid outputs;
• cross-attention: Source padding should not be read or processed;
• loss: Loss computation should only consider valid target positions;
• packed/batched: Samples of different lengths must not interfere with one another;
• cache: Incremental logits generated from caching must be within tolerance of full-prefix logits.

If the mask is off by even one position, the model might train faster and achieve lower loss, because it effectively sees the answer.

11. Training Stability

The Transformer still relies on the principles established in Chapter 9:

• residual/norm layout and initialization;
• optimizer and weight decay parameter groups;
• warmup and learning schedule;
• gradient clipping, mixed precision, and loss scaling;
• dropout and its mode;
• validation, checkpointing, seed setting, and data ordering.

LayerNorm does not depend on batch statistics, but it cannot automatically prevent activation outliers, attention logit overflow, or deep residual accumulation.

12. The Model Still Requires Task Bias

The flexibility of Transformer connections does not mean there is no prior knowledge: tokenization, masking, position encoding, context window size, weight sharing, and the training objective all introduce bias.

Image Transformers require patching, positional information, and augmentation; time series models need causal cutoffs, scaling, and calendar-based features; collection data may not require absolute positional information. The architecture should align with both the data generation process and the deployment workflow.

Common Misconceptions

• Transformers eliminate order: They still require position embeddings; the decoder must maintain causal ordering during generation.
• All tokens can be generated in parallel: Parallelization applies only to teacher-forced training positions, not to autoregressive generation steps.
• KV cache makes long-context cost constant: Attention and cache memory usage still scale with increasing context length.
• Long context means effective use of long context: The model must be validated with task-level positional or retrieval mechanisms.
• LayerNorm solves all training stability issues: Problems remain in initialization, precision, optimization, and residual connections.

Exercise

1. Draw information flow diagrams and mask configurations for encoder-only, decoder-only, and encoder–decoder architectures.
2. Implement sinusoidal positional encoding and verify behavior for odd/even dimensions and length extrapolation.
3. Modify the future token input in a pre-norm block and validate causal invariance.
4. Compare full-prefix attention with incremental logits using KV cache.
5. Estimate the memory footprint of an attention matrix, FFN activations, and KV cache for a given model.

Summary

The Transformer block combines attention, position encoding, feed-forward networks, residual connections, and normalization into a stackable architecture. Encoders, decoders, and encoder-decoder structures are distinguished by the flow of information and the training objective. Training can parallelize over known positions, but autoregressive generation still advances one token at a time; long-context capabilities remain constrained by attention mechanisms, caching, and actual hardware utilization.

The next chapter moves into pretraining: the tokenizer determines the discrete units the model observes, while the objective dictates what the model learns from vast corpora. Fine-tuning and prompting represent different interfaces for altering model behavior.