3 A blueprint to modern transformers

In this chapter, the book shifts from a previous re-architecture success story to a deeper look at how modern Transformers are built and why their internal parts matter for optimization. Using DistilGPT2 as a compact example of the classic design, it shows how a model is organized into embeddings, repeated Transformer blocks, and an output head, and how execution actually flows through nested forward passes rather than simply through the order in which layers are listed. The key idea is that understanding the static structure of a model is the starting point for understanding its runtime behavior and for choosing the right re-architecture technique.

The chapter then explains the two main functional pillars of a Transformer block: attention and the MLP. Attention gives tokens context by letting each token compare its Query with other tokens’ Keys and mix in their Values, but this creates a growing KV cache that can become a major memory bottleneck during generation. To address that, modern models like Llama replace classic multi-head attention with Grouped-Query Attention, reducing KV storage by sharing keys and values across groups of query heads. The MLP side is also modernized: instead of a simple expand-activate-contract pipeline, newer models use GLU variants such as SwiGLU, where a learned gate modulates which information is allowed through, improving flexibility and shaping how knowledge is processed.

Finally, the chapter connects these structural choices to practical optimization strategies. It distinguishes depth from width, showing that depth pruning removes whole Transformer blocks, while width pruning targets the expansion layers inside the MLP; with GLU-based models, pruning must consider paired gate and content projections together. It also links internal activations to behavior, emphasizing that model weights are fixed learned parameters while activations reveal how the model responds to specific inputs, which makes it possible to diagnose inefficiency or bias. The result is a clear map from architecture component to optimization method, preparing the reader to apply the right technique to attention, MLPs, blocks, or activations in later chapters.

Embedding creation process in the input block. The wte layer converts each text token into a Token Embedding based on its meaning. The wpe layer generates a Positional Embedding based on the token's position in the sequence. The combination of both produces a vector that contains both semantic and positional information.

Anatomy of DistilGPT2. On the left, the general architecture is shown: data flows from the Embeddings layer, passes through the stack of six Transformer Blocks, and ends at the output layer. On the right, a "zoom" into one of those blocks reveals a LayerNorm → GPT2Attention → LayerNorm → GPT2MLP flow, where each module processes the information before passing it to the next.

Multi-Head Attention architecture. The input embedding (768 dimensions) is divided into 12 specialized attention heads, each processing 64 dimensions in parallel. Each head focuses on different contextual relationships, before concatenating their outputs into a context-aware embedding that incorporates information from the entire sequence. Each head focuses on different contextual relationships (labels are illustrative; real specialization emerges from training and is not directly interpretable).

Attention mechanism for the token "spring." The Query (Q) from "spring" compares with the Keys (K) of all tokens through dot product operations. The resulting scores, after softmax normalization, determine what fraction of each Value (V) contributes to create the enriched embedding.

Attention mechanisms and memory optimization strategies. Multi-Head Attention (MHA) maintains independent K,V pairs for each Q head — one pair per processed token, forming a growing cache. Grouped-Query Attention (GQA) groups Q heads to share K,V pairs (here 6 Q heads in 3 groups, each sharing one K,V pair per token). Multi-Query Attention (MQA) compresses to the extreme: all Q heads share a single K,V pair per token.

KV Cache memory growth during sequential token generation. The graph compares cache growth for three attention mechanisms using the example from the text: a 6-layer model with 12 attention heads. The solid line shows MHA storing 12 KV pairs per token, requiring 1,152 vectors to cache 8 tokens. GQA 4:1 with 3 KV pairs per token can process 32 tokens with the same 1,152-vector budget (+300%), while MQA with a single KV pair can process 96 tokens (+1100%).

Comparison of the "shapes" of different models. The narrow column represents the embedding dimension, the main box shows the model's depth (its height, proportional to the number of blocks) and its processing width (its width, proportional to the size of the MLP expansion).

Processing flow in MLP Modules: classic vs. GLU. While the classic MLP follows a simple expansion-activation-contraction sequence, GLU uses two parallel projections: one up_proj that transforms the content and another gate_proj that creates a dynamic filter.

How activations flow between layers. Each neuron applies its static weights (learned during training) to the input values and passes the result through an activation function (SiLU). The output values, the activations, vary in intensity: high values propagate strong signals, while low values suppress information. These activations become the inputs for the next layer, creating the continuous flow of information through the model.

• Transformers consist of repeating blocks, each containing two core modules: an attention module for contextualizing information and an MLP module for processing knowledge.
• The two fundamental dimensions for rearchitecting a model are its depth (the number of stacked Transformer blocks) and its width (the size of the expansion within the MLP modules).
• Modern architectures like Llama evolved from the classical design by incorporating key innovations such as Grouped-Query Attention (GQA) to reduce memory consumption and Gated Linear Units (GLU) to create more sophisticated MLP processing.
• A model's static structure defines its components, while its dynamic behavior is revealed through its activations during inference.
• By understanding the specific function of each component, you can strategically map techniques like Depth Pruning, Width Pruning, and Attention Bypass to the parts of the model where they will have the most impact.

FAQ