10 Exploring the transformer black box
This chapter covers
- Capturing internal activations with PyTorch hooks
- Visualizing activation heatmaps by token and by layer
- Observing the GLU mechanism in action on real data
- Trace how the model builds its response
Throughout the book, you've analyzed at different points the importance of Transformer blocks, individual layers, or even isolated neurons in the model's specific task. In this chapter, we're going to dive deeper into the study of activations to explore how information flows through the model and how it makes its decisions.
A Transformer block, as you've already seen, isn't an opaque unit; it has different modules like the MLP module, the Attention module, and normalization layers. The block can be explored as a whole, or we can select the specific components we want to examine.
Using hooks, which you've already mastered, we'll capture the model's different activations. From the information collected, we'll build three visualization tools that will let us analyze the LLM's behavior:
- Activation heatmaps: they show which neurons respond to each token and how that pattern changes across the model's depth.
- GLU mechanism view: the same structure we studied in Chapter 5, now visible over real activations.
- Logit lens: a technique that projects the internal state of each layer through the model's prediction head, letting us see how the final output takes shape layer by layer.