2 An end-to-end rearchitecting project
The chapter presents a practical, end-to-end example of rearchitecting a language model to make it smaller and faster while preserving as much capability as possible. It frames model performance as a balance between two goals: reasoning quality and inference efficiency. To make the process concrete, it walks through a compact Gemma model, first establishing a clear baseline with a sanity-check prompt, benchmark scores, and timing measurements so that every later change can be judged against the original behavior.
The structural optimization phase then applies depth pruning by removing the final two transformer layers from the model. The chapter explains the mechanics of copying the model, trimming its layer list, and updating the configuration so the modified architecture remains valid. This pruning yields a modest parameter reduction but a meaningful speedup, and the evaluation shows the trade-off clearly: the pruned model runs faster, yet it loses noticeable capability on several benchmarks and begins to produce less reliable, more hallucinatory outputs.
To address that loss, the chapter introduces knowledge recovery through distillation, using the original model as teacher and the pruned model as student. A subset of a general text dataset is tokenized, batched, and used in a training loop where the student learns to match the teacher’s output distributions by minimizing KL divergence. After this recovery step, the smaller model regains most of its lost performance while keeping the speed advantage, demonstrating that careful pruning followed by distillation can produce a more efficient model with only limited capability loss.
The rearchitecting cycle followed in this chapter. We start with an original model, make it smaller through pruning, and finally recover its capabilities through knowledge distillation.
The original gemma-3-270m model is reduced from 18 to 16 layers. The capabilities results shown are achieved via the knowledge recovery phase.
In the knowledge distillation cycle, the Teacher model generates its reasoning (logits) for a given text. The Pruned model (our student) does the same. The KL Divergence loss function measures the difference between both "reasonings," and this error is used to adjust the student's weights, making it think more like the Teacher over time.
Summary
- Establishing a quantitative baseline is an indispensable first step in any optimization project because it provides an objective metric against which the success or failure of interventions is measured.
- Depth pruning removes entire Transformer layers to reduce model size and speed up inference, typically targeting the final layers that hold refinement details rather than core knowledge.
- Knowledge distillation lets smaller “student” models recover lost capabilities by learning to mimic the reasoning patterns of larger “teacher” models.
- Combining structural pruning and knowledge distillation lets you build more efficient models that retain most of the capabilities of their larger, original versions.
FAQ
What is the main goal of the end-to-end rearchitecting project in this chapter?
The goal is to take a working base model, make it smaller and faster, and preserve as much of its original capability as possible. The chapter does this by combining structural optimization with knowledge recovery.
What two phases of the rearchitecting pipeline are used in this chapter?
The chapter uses two phases: Structural Optimization, where the model architecture is modified to be lighter and faster, and Knowledge Recovery, where knowledge is transferred back into the smaller model using distillation.
Why is gemma-3-270m chosen as the model for this project?
It is compact enough to run well on a T4 GPU in Google Colab, but still advanced enough to demonstrate real optimization techniques. It also reflects current architectural trends.
What does the chapter mean by model performance?
Model performance is defined as two equally important dimensions: capabilities, which measure reasoning and correctness on benchmarks, and inference efficiency, which measures speed and resource usage such as tokens per second and latency.
Why is it important to establish a baseline before pruning a model?
A baseline gives a reference point for comparison. Without it, you cannot tell whether pruning improved or degraded the model, or by how much.
Which benchmarks are used to evaluate the baseline model?
The chapter uses ARC-Easy, Winogrande, HellaSwag, and LAMBADA. Together, they measure reasoning, commonsense understanding, and word prediction behavior.
What is depth pruning in this chapter?
Depth pruning is the process of removing complete transformer layers from the model. In this chapter, the last two layers are removed to reduce depth and improve inference speed.
Why does the pruned model need its configuration updated after layers are removed?
After changing the architecture, the model’s configuration must match the new structure. Updating fields like num_hidden_layers prevents mismatches and runtime errors.
How does knowledge distillation help recover lost capabilities?
Knowledge distillation trains the smaller student model to imitate the teacher model’s output distribution, not just its final answers. By minimizing KL divergence, the student learns to behave more like the original model.
What was the overall result of pruning and then distilling the model?
The final student model kept the smaller size and speed gain while recovering most of the original capabilities. The chapter reports that the recovered model is about 10% faster, 4.16% smaller, and retains roughly 87% to 98% of the original benchmark performance.
Rearchitecting LLMs ebook for free
