Appendix A. A Gemma-3 case study in SLM architectural refinements
This case study shows how modern small language models keep the Transformer foundation while refining it for efficiency and stability. The text begins by emphasizing that, although GQA and GLU became popular design patterns, newer model families still introduce their own adjustments, especially as model sizes shrink from massive LLMs to compact SLMs. Using Gemma-3-270M as the example, the discussion frames SLM design as a set of architectural trade-offs aimed at preserving capability while prioritizing compactness and faster inference.
The main attention-side refinement is a move from GQA to MQA. By inspecting Gemma’s layer shapes and configuration, the text shows that the model uses four query heads but only one key/value head, meaning all query heads share the same K and V vectors. This is presented as the most aggressive form of KV-cache compression, and the practical effect is dramatic: memory use drops sharply while the number of processable tokens increases substantially. The comparison with MHA and GQA makes clear that this sharing strategy is a major reason small models can handle longer contexts efficiently on limited hardware.
The chapter also explains why Gemma adds q_norm and k_norm inside attention: when multiple queries share a single KV set, normalization helps keep vector magnitudes balanced so no single head dominates the softmax calculation. Without this, the attention mechanism could become unstable or behave too much like a single-head model, reducing expressiveness. The text closes by noting that these shared-component designs complicate head pruning, since removing one head can disrupt the shared projections, and then transitions toward the model’s MLP refinements as the next part of the analysis.
Figure A.1 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
A.2 Refinements in the MLP module: expansion and stability
Throughout this book you've worked with GLU structures across multiple chapters, so its function as an expansion-contraction mechanism should already be familiar, and Gemma doesn't change these fundamental premises. Its MLP module has the same foundations, but introduces some interesting modifications, designed specifically to optimize performance in a model of such ”small” dimensions.
To start, let's recall the structure of Llama-3.2-1B's MLP Module:
We compare it with the relevant and equivalent layers of Gemma-3.2-270M:
As you can see, we've included the normalization layers in the structure. Although they are defined outside the MLP Module, they work together with it inside the Transformer block.
The first key difference we'll analyze is the expansion ratio, a parameter that defines the "width" of the processing in the MLP. Let's calculate it for both models using the structures we've just seen:
Llama-3.2-1B: 8192 (intermediate dimension) / 2048 (embedding dimension) = 4.0x
Gemma-3-270M: 2048 (intermediate dimension) / 640 (embedding dimension) = 3.2x
Gemma, with its 18 blocks and a more conservative expansion ratio, exemplifies a deeper and narrower model. It performs more sequential transformations, but each one is computationally less intensive.
In contrast, Llama-3.2-1B, with 16 layers but a much larger embedding and expansion ratio, represents a less deep and wider model. It executes fewer transformation steps, but each one has much greater processing capacity.
These characteristics are already giving us a hint about what type of pruning can be more effective in each model.
Along with the expansion ratio, the other variation we find is the activation function. While Llama uses SiLU(), Gemma opts for PytorchGELUTanh(), a variant of the GELU function. The GELU-Tanh function offers smoother gradients and more stable convergence, which is especially valuable in smaller models where each parameter has a greater impact.
Finally, we find the pre_feedforward_layernorm and post_feedforward_layernorm normalization layers. Although in the structure they appear listed at the Gemma3DecoderLayer level, their function is intimately linked to the MLP. These layers act like a "corset" that wraps around the MLP module, stabilizing the data that enters and exits it. In a "thin and deep" architecture like Gemma's, with a small embedding of 640 (you can see the embedding size in listing 3.1), but 18 layers, maintaining numerical stability throughout the entire stack of blocks is a challenge, and these additional layers are a safeguard to prevent problems during training and inference.
The most important aspect for us is that none of these optimizations alter the fundamental principle of the MLP module. The GLU structure that expands and contracts the information, and which is the target of our Width Pruning, remains intact, ready to be optimized with the techniques we discuss in Chapter 5.
Rearchitecting LLMs ebook for free
