9.2 Observational Methods¶

The Logit Lens (nostalgebraist, 2020) is one of the first tools developed to look inside transformers. It enables us to observe how a transformer refines its predictions layer by layer —allowing us to see not just the final output but the evolving "thought process" the model undergoes as it makes a prediction.

Transformers are trained to predict the next token in a sequence. They do this by transforming the inputs layer by layer, with each layer adding new information to improve the prediction. The Logit Lens enables us to “translate” each layer's internal representation back into tokens. By seeing what the model “predicts” at each layer, we can trace how its predictions evolve from a rough guess in the initial layers to a refined choice in the final one. However, it’s worth noting that while the Logit Lens lets us see the intermediate predictions at each layer, it doesn’t explain the mechanism of transformation. The Logit Lens is an inherently observational tool —it reveals what is the most probable token at the end of each layer but doesn’t allow us to understand why this precise token is predicted.

A transformer model is a powerful architecture for processing sequences of data, especially text. It is trained to predict the next word or subword in a sequence, referred to as tokens.

The set of all tokens (words or subwords) that the model can recognize is called the vocabulary. For example, GPT-2 has a vocabulary of 50,257 tokens. More precisely, a transformer takes a sequence of tokens as input and is trained to predict the next token in this sequence, outputting a probability distribution over the entire vocabulary.

Internally, a transformer consists of multiple stacked layers, each containing two sublayers: an MLP (multi-layer perceptron) and attention heads. For understanding the Logit Lens, we don’t need to go into the details of how these sublayers function.

The intermediate layers are connected through what’s known as the residual stream. The residual stream is a pathway that carries information from the input to the output, allowing it to flow through all layers of the model.

Figure: A minimalist illustration of a transformer with a single layer. Transformers cannot directly manipulate tokens, so they embed tokens into numerical vectors, an operation depicted in the bottom grey box. After embedding, the tokens are represented as vectors within the residual stream. The first sublayer (attention heads, represented by the two side-by-side boxes on the left) reads these embedded tokens, applies a transformation, and adds its output back into the residual stream. The second sublayer (MLP) does the same. Finally, to output a token, the embedded tokens must be converted back to the vocabulary space through an unembedding operation (depicted in the top box). The residual stream enables each layer to make incremental adjustments to the model’s predictions by accumulating and refining the information passed through each transformation. From (Elhage et al., 2021).

The essential idea behind the Logit Lens is that the unembedding operation, typically applied only after the last layer, can be applied at any point in the residual stream. After each intermediate layer, the residual stream can be unembedded—that is, converted back into the token space—allowing us to observe which token is currently the most probable.

In the previous section on Feature Visualization, we saw that CNNs build up their understanding of an image layer by layer, detecting simple patterns like edges in early layers and more complex shapes or objects in later ones. Feature visualization allows us to interpret these progressive transformations by displaying the visual patterns that different neurons respond to.

The Logit Lens can be seen as a parallel interpretability tool for transformers. Instead of visualizing image features, the Logit Lens lets us see the model’s step-by-step guesses for the next token.

Probing is a technique used to analyze neural networks and understand which representations or concepts they have learned and where (Belinkov, 2022). Probing techniques can be applied to a range of models, from transformers to CNNs, to investigate whether specific information or properties are encoded in a model’s intermediate representations.

What is a probe? A probe is a lightweight model, often a linear classifier, trained to detect whether a specific concept or feature is represented in the activations of a neural network. In a probing setup, researchers analyze activations (the intermediate outputs) from layers within a network to see if these activations encode information about a particular property or concept. For example, probes may be used to determine whether a language model’s activations contain information about grammar rules or whether a chess-playing model like AlphaZero encodes strategic knowledge about the game, where this knowledge is located within the network, and when it is acquired during training (Gurnee et al., 2023, McGrath et al., 2022).

One well-known example of probing was applied to AlphaZero, the neural network trained to play chess that famously defeated top human players. Researchers used probes to investigate whether AlphaZero internally represents certain strategic concepts about chess, such as “Can the opponent capture my queen?” or “Is there a checkmate threat within one move?” (McGrath et al., 2022).

Probing classification is the process of training classifiers on the network’s intermediate activations to identify if specific properties are encoded. The steps of probing classification are as follows:

1. Choose a Property or Concept: Define a specific concept to explore, such as “Can the opponent capture my queen?”.

2. Generate or Select Input Data: Create a dataset with examples that vary in terms of the target property. For instance, in chess, we might create board configurations where “the player can capture the opponent’s queen” is either true or false.

3. Record Intermediate Activations: Feed this dataset through the model and record the activations of neurons at different layers.

4. Train a Classifier on the Activations: Use the recorded activations as input features to train a classifier (probe) that distinguishes between different classes based on the concept (e.g., true vs. false input prompts).

5. Evaluate the Probe’s Accuracy: If the probe achieves high accuracy, this suggests that the target concept is strongly encoded in the recorded layer’s activations.

Figure: Accuracy of two probes trained on AlphaZero’s intermediary activations, across layers and training steps The two two probes were trained on the two concepts: “Can the playing side capture their opponent’s queen?”, and “Could the opposing side checkmate the playing side in one move?”. The color represents the probe’s accuracy (the yellow color corresponds to a higher value, indicating that the feature is more represented). Here, the probes were trained at different stages of AlphaZero's training and from the activations of different layers of the network (a block corresponds to one layer of AlphaZero). It can be seen that the more AlphaZero is trained, the more it represents these two concepts. The concept "has_mate_threat" seems to be represented quite homogeneously across AlphaZero's layers, while "can_capture_queen_opponent" appears to be more represented in the earlier layers. "The ability to predict has_mate_threat from AlphaZero’s activations indicates that AlphaZero is not simply modeling its potential moves, but also its opponent’s potential moves and their consequences during position evaluation". From (MacGrath et al., 2022).

Using linear probes for detecting prompts leading to dangerous behavior. The alignment team from Anthropic is currently working on a research agenda called “sleeper agents” where they deliberately train deceptively aligned models (Hubinger et al., 2024) to test whether currently available techniques can detect misalignment. These "sleeper agents" are models with hidden backdoors introduced during training. One of the sleeper agents they trained writes secure code when the prompt indicates year “2023” but inserts vulnerabilities when the prompt indicates “2024”. Such dangerous outputs are hard to detect, and could easily be missed by a human overseer, this is why it would be extremely useful to develop an automatic detection method. Anthropic researchers have trained linear probes based on sleeper agents’ residual stream activations that classify inputs as safe or dangerous with very high accuracy (MacDiarmid et al., 2024). Whether or not a prompt will trigger defection appears to be linearly encoded across a wide range of middle residual streams. However, they remain wary of this result and acknowledge that the surprisingly high accuracy of the probes may be due to the backdoor insertion, and that “natural defection” may not be detected as easily.

It was also suggested to train probes on human-written scenarios of takeover attempts or catastrophic actions (Roger., 2023) to detect such goals or intentions in models before they happen. Similarly to sleeper agents, it is unclear whether probes trained in artificial settings will generalize in real-world scenarios.

Limitations of probing classifiers

• Correlation, not causation: Probing classifiers indicate that a concept (or a proxy of that concept) is encoded, but they don’t reveal whether this concept is actively used by the network during inference. The classifier’s high accuracy may reflect the ease with which it can detect patterns, not necessarily that the model relies on those patterns for decision-making. To discover causal effects we need to intervene in representations of the model instead of just observing them.

• Identifying proxies rather than “true concepts”: Probing may detect patterns that act as proxies for the concept of interest, rather than the concept itself. For example, a probe trained to detect “legal language” might actually be detecting correlated cues (such as certain formal words or sentence structures) rather than a genuine understanding of legal terminology. This makes it challenging to interpret probes as definitive indicators of the exact concepts represented within the model.

To make sense of data, classify it, or make decisions, neural networks need to learn features—representations that capture meaningful patterns in the data. However, networks have a limited number of neurons to store a vast amount of information. Instead of assigning a single neuron to each feature, neural networks often "share" neurons across multiple features. This shared, overlapping storage is known as superposition.

Superposition occurs because it allows the network to handle more features using fewer neurons, making it more memory-efficient. However, this efficiency comes at a cost: polysemanticity. Polysemanticity means that a single neuron or component in the network represents multiple, often unrelated, features. For example, one neuron might activate in response to both "cats" and "cars," even though these concepts are entirely different.

The vocabulary surrounding superposition and related concepts is often messy. The terms superposition and polysemanticity are often used interchangeably, even though they refer to slightly different aspects of the same phenomenon in some contexts: superposition describes the overlapping storage of features, while polysemanticity highlights the behavior of neurons that respond to multiple distinct features.

Also keep in mind that superposition should not be confused with the superposition hypothesis, which is a specific theory proposed to explain why and how superposition occurs in neural networks.

The Toy Models of Superposition paper introduces simplified models that help researchers study superposition in a controlled environment (Elhage et al., 2022). These models provide evidence supporting the superposition hypothesis—a theory about how neural networks efficiently store and organize information.

The superposition hypothesis suggests that neural networks can represent more features or concepts than they have individual neurons by encoding information as linear combinations across multiple neurons. This means:

1. Efficient compression: Neural networks can compress information by representing more features than there are available dimensions, optimizing memory use.

2. Distributed representation: Features are not exclusively tied to single neurons; instead, they are distributed across multiple neurons.

3. Non-orthogonal directions: Features are stored in directions that are not perfectly orthogonal in the network's activation space, which leads to overlaps and potential interference between concepts.

This paper is a cornerstone of mechanistic interpretability because it reveals fascinating phenomena about how neural networks organize and store information:

• Demonstrating superposition: The experiments show that superposition occurs and identify the conditions under which it arises.

• Explaining mono- and polysemantic neurons: The paper clarifies why some neurons specialize in a single feature (monosemantic) while others represent multiple features (polysemantic).

• Phase transitions in training: It highlights a phase change¹ during training that determines whether features are stored in superposition.

• Geometric feature organization: Features in superposition are arranged into geometric structures such as digons, triangles, pentagons, and tetrahedrons, providing insight into the network’s internal organization.

The following figure is a great illustration of superposition in a toy model. The toy model has 2 dimensions, represented by the x and y axis in the figure below, but it needs to learn 5 features. This means that the model needs to find a way to fit more information (5 features) into fewer dimensions (2-dimensional space). Each feature is given an importance, represented with colors. An important feature is a feature that has a significant impact on the model’s accuracy or loss function (if removing or weakening the representation of this feature causes a large drop in performance, it would be considered important). The key challenge in superposition is how to efficiently encode important features while minimizing overlap and interference between them.

Each feature also has a sparsity. This refers to how often it is active in the dataset. A dense feature (low sparsity) is activated frequently, making it harder for other features to coexist in the same dimensions without interference.

Figure: How does a 2-dimension model encode 5 features as their sparsity increases? 0% sparsity means that features are very dense, or frequent. The model only encodes the two most important features in orthogonal dimensions. As sparsity increases and the important features are less frequently useful, the toy model encodes additional features in non-orthogonal directions. The intuition is that the less frequent the features are, the less likely two overlapping features are to be activated at the same time and cause interference. So the cost of interference between features is outweighed by the advantage of learning more features. At 90% sparsity the 5 features are represented as a pentagon. From (Elhage et al., 2022).

When encoding features, there is a tradeoff to make between the usefulness of having as many features as possible and low interference between them. Models embed their features into very complex geometric structures to reach optimal encoding. The figure below shows the different geometric structures that models use to encode features.

Figure: As features get sparser (less frequently activated), more of them can be encoded optimally and in more complex geometric structures. The first figure on the top left shows that when features are very dense, only the most important features are represented and they are organized in tetrahedrons. This model learns 28 features and encodes them in 7 tetrahedrons. On the second figure features are slightly sparser, more of them can be optimally encoded, and tetrahedrons are replaced by triangles and digons. This model learns 46 features encoded in triangles and digons. Superposition exhibits complex geometric structure. From (Elhage et al., 2022).

A major challenge in mechanistic interpretability is polysemanticity, where a single neuron or feature represents multiple, unrelated concepts. Polysemanticity arises naturally in LLMs’ MLPs and residual streams, and makes it more complicated to identify which specific features influence a model’s outcome. Sparse AutoEncoders (SAES) are a promising approach for disentangling features within a network (Bricken et al., 2023, Gao et al., 2024, Templeton et al., 2024).

SAEs are gaining popularity because they have shown promising results in separating out features. Features extracted using SAEs can then be used to:

• Steer language models behavior away from undesirable outcomes (see section Activation Steering). (Templeton et al., 2024) found a bunch of safety-relevant features in Claude 3 Sonnet, including features for unsafe code, bias, sycophancy, deception and power seeking, and dangerous or criminal information. These features activate on text involving these topics and causally influence the model’s outputs when intervened upon.

• Find more interpretable circuits directly made of features instead of model components. In particular, finding and understanding the circuits in which safety-relevant features are involved could be valuable (see section on Automating and Scaling Interpretability).

SAEs are also great because they are trained in an unsupervised manner, which enables us to discover abstractions or associations formed by the model that we might not have anticipated beforehand.

An autoencoder is a neural network designed to learn compressed representations of input data by encoding it into a lower-dimensional latent space and then reconstructing the original input from that representation. The latent space is the layer where the data is represented in a compressed or abstracted form, containing the key features necessary for reconstructing the input. This learned representation often captures the essential patterns or features in the data.

F*igure: An autoencoder learns two transformations, represented by encoder weights and decoder weights, to compress input data into a lower-dimensional latent space and then reconstruct the original input from this representation. *

Why are SAEs “sparse”? Autoencoders can vary in structure and purpose. SAEs are autoencoders that introduce sparsity constraints on the activations in the latent space, encouraging the model to use a limited number of latent neurons for each input. This “sparsity” forces the model to learn distinct and specific features, making the representations more interpretable.

How do SAEs help disentangling features in transformers? The process of disentangling model activations into interpretable features typically involves training a sparse autoencoder to reconstruct activations from specific parts of a model, such as the MLP of a particular layer or a residual stream, with a latent space larger than the input², so that each neuron in the latent space will hopefully be monosemantic - representing a single feature - and interpretable.

Figure: Models activations can be decomposed into features using a sparse autoencoder. This figure illustrates an SAE trained to disentangle features in an MLP. From (Bricken et al., 2023).

Although promising, SAEs are still early work with limitations:

• Incomplete understanding of model usage: Identifying model features doesn’t reveal how they are used during inference, we still have to find the circuits involving them.

• Difficulty in feature interpretation: Not all features discovered by SAEs are easily interpretable; some features remain challenging to understand.

• Lack of validation methods: Currently, there are limited methods to test the validity of feature interpretations.

• SAEs have poor reconstruction quality: Sparse autoencoders don’t reconstruct model activations very well, which means that they don’t completely capture the behavior of our models. For instance, passing GPT-4’s activations through an SAE results in performance equivalent to a model trained with 10x less compute (Gao et al., 2024).

SAEs present intriguing research questions for AI safety and interpretability:

• What features activate during jailbreaks?

• What features need to activate or to remain inactive for a model to give advice on producing cyberattacks, bioweapons, etc. ?

• Can we use the feature basis to detect when fine-tuning a model increases the likelihood of undesirable behaviors?

• etc.