10 Exploring the transformer black box

In this chapter, the model is treated less like a black box and more like a system whose internal signals can be observed and compared. The core idea is to use hooks to capture activations from selected transformer submodules, especially the MLP path, and then turn those tensors into readable views of how information moves through the network. Because transformer blocks are hierarchical and model families name their parts differently, the chapter also emphasizes how to find the right hook points in architectures such as Llama, Gemma, and Qwen, while accounting for details like forward hooks versus forward pre-hooks when the goal is to inspect an expanded representation before it is compressed again.

Once the activations are available, the chapter builds a practical visualization toolkit to make them interpretable. Heatmaps show neuron activity either for one layer across tokens or across all layers after averaging tokens, revealing how different words trigger different response patterns and how those patterns evolve with depth. The GLU mechanism is then inspected in action by comparing the activation function output, the up projection, and the input to the down projection side by side, making clear that the elementwise product acts as a filter rather than a simple amplifier. These visualizations are designed not just to describe one prompt, but to provide a reusable method for later investigations of behavior and bias.

The final perspective comes from the logit lens, which projects hidden states from each layer through the model’s final normalization and output head to show how a prediction gradually forms. Early layers produce noisy or incoherent guesses, intermediate layers begin to settle on semantically related candidates, and the last layers crystallize a coherent answer. The chapter closes by connecting these lightweight interpretability tools to broader mechanistic interpretability work, contrasting the accessible logit lens with more resource-intensive approaches like sparse autoencoders, and positioning the whole toolkit as a practical foundation for future model analysis and intervention.

Token × neuron activation heatmap for act_fn_layer_12 on the prompt "The cat sleeps". Each row is a token, each column is a bin of 256 neurons averaged together. Activation intensity varies visibly across tokens: "cat" shows consistently stronger responses across most neuron bins, while "The" and "sleeps" remain more uniform. The BOS token has been excluded from the visualization.

Layer × neuron heatmap for act_fn on the prompt "The cat sleeps", averaged over tokens. Each row is a transformer layer, each column a bin of 256 neurons. The BOS token has been excluded before averaging.

GLU three-panel visualization for layer 14 on the prompt "The cat sleeps". From left to right: act_fn (SiLU(gate)), up_proj, and down_proj_input (SiLU(gate) × up). Each panel shares the same token axis but has its own independent color scale. The BOS token has been excluded.

Logit lens top-three tokens per layer for the prompt "The most significant building in Barcelona is the". Each cell shows the top-three predicted tokens and their probabilities for one layer, projected through the model's final normalization and lm_head. Background intensity encodes the top-1 probability, with darker cells indicating higher confidence.

• PyTorch hooks let you capture a model's internal activations by attaching to the input or output of any submodule via forward hooks and forward pre-hooks.
• The internal structure of a Transformer block is hierarchical and can be explored separately. Although submodule names vary between model families (self_attn versus linear_attn, for example), the MLP module keeps the same structure in Llama, Gemma, and Qwen, which lets you apply the same capture techniques consistently.
• Activation heatmaps and grouping neurons into bins turn tensors of thousands of values into a manageable visual representation.
• The GLU mechanism can be observed in action on real activations. Visualizing its three components side by side shows that the multiplication between act_fn and up_proj does not amplify the signal, it filters and selects it.
• The logit lens technique extracts hidden states from intermediate layers and projects them through the final normalization and the prediction head. It lets you visualize how the model builds its output layer by layer.
• Direct activation interpretability offers a viable alternative for resource-constrained environments compared to computationally expensive methods like sparse autoencoders.

FAQ