Emergent Linear Representations in World Models of Self-Supervised Sequence Models (2023)

Abstract: How do sequence models represent their decision-making process? Prior work suggests that Othello-playing neural network learned nonlinear models of the board state (Li et al., 2023a). In this work, we provide evidence of a closely related linear representation of the board. In particular, we show that probing for “my colour” vs. “opponent’s colour” may be a simple yet powerful way to interpret the model’s internal state. This precise understanding of the internal representations allows us to control the model’s behaviour with simple vector arithmetic. Linear representations enable significant interpretability progress, which we demonstrate with further exploration of how the world model is computed.

## introduction how do sequence models represent their decisionmaking process? large language models are capable of unprecedented feats, yet largely remain inscrutable black boxes. yet evidence has accumulated that models extract features -articulable properties of the input 2 -and represent them in its internal activations @xcite @xcite @xcite @xcite . a key first step for interpreting them is understanding how these features are represented. @xcite introduce the linear representation hypothesis: that features are represented linearly as directions in activation space. this would be highly consequential if true, yet this remains controversial and without conclusive empirical justification. in this work, we present novel evidence of linear representations, and show that this hypothesis has real predictive power. we build on the work of @xcite , who demonstrate the emergence of a world model in sequence models. namely, the authors train oth-ellogpt, an autoregressive transformer model, to predict legal moves in a game of othello given a sequence of prior moves (section 2.2). they show that the model spontaneously learns to track the correct board state, recovered using non-linear probes, despite never being told that the board exists. they further show a causal relationship between the model's inner board state and its move predictions using model edits. namely, they show that the edited network plays moves that are legal in the edited board state even if illegal in the original board, and even if the edited board state is unreachable by legal play (i.e., out of distribution). critically, the original authors claim that othel-logpt uses non-linear representations to encode the board state, by achieving high accuracy with non-linear probes, but failing to do so using linear probes. in our work, we demonstrate that a closely related world model is actually linearly encoded. our key insight is that rather than encoding the colours of the board (black, white, empty), the sequence model encodes the board relative to the current player of each timestep @xcite . in other words, for odd timesteps, the model considers black tiles as mine and white tiles as yours, and vice versa for even timesteps (section 3). using this insight, we demonstrate that a linear projection can be learned with near perfect accuracy to derive the board state. we further demonstrate that we can steer the sequence model's predictions by simply conducting vectoral arithmetics using our linear vectors (section 4). put differently, by pushing the model's activations in the directions of mine, yours, or empty, we can alter the model's belief state of the board, and change its predictions accordingly. our intervention method is much simpler and interpretable than that of @xcite , which rely on gradients to update the model's activations (section 4.1). our results confirm that our interpretation of each probe direction is correct, but also demonstrates that a mechanistic understanding of model representations can lead to better control. our results do not contradict that of @xcite , but add to our understanding of emergent world models. we provide additional interpretations of the sequence model using linear operations. for example, we provide empirical evidence of how the model derives empty tiles of the board, and find additional linear representations, such as tiles being flipped at each timestep. finally, we provide a short discussion of our thoughts. how should we think of linear versus non-linear representations? perhaps most interestingly, why do linear representations emerge? ## preliminaries in this section we briefly describe othello, othel-logpt, and our notations. ## linearly encoded board states in this section we describe our experiments to find linear board state representations. ## intervening with linear directions in this section we demonstrate how we intervene on othellogpt's board state using linear probes. ## additional linear interpretations the linear representation hypothesis is of interest to the mechanistic interpretability community because it provides a foothold into understanding a system. the internal state of the transformer, the residual stream, is the sum of the outputs of all previous components (heads, layers, embeddings and neurons) @xcite . albeit the residual stream consisting of linear and non-linear transformations, linear functions of the residual stream allow us to identify where a computation of interest takes place, or trace how a representation of interest evolves over a forward pass. in this section we leverage our newfound linear representation of board state to provide additional interpretations of othellogpt, as proof of concept of how discovering linear representations unlocks downstream interpretability applications. ## discussions 6.1 on linear vs. non-linear interpretations one challenge with probing is knowing which features to look for. 13 for instance, classifying {black, white} versus {mine, yours} leads to different takeaways, which illustrates the danger of projecting our preconceptions. what might seem "sensible" to a human interpreter (black, white) may not be for a model. in hindsight, given the symmetric game-play of othello, encoding mine, yours is perfectly sensible for the model (for more examples of non-obvious, sensible features, see @xcite ). more broadly, what is sensible, or alternatively, how we choose to interpret linear or non-linear encodings, can be relative to how we see the world. suppose we had a perfect world model of our physical world. further suppose that if and when it 13 for a longer discussion on probing, see appendix. computes a gravitational force between two objects (newton's law), we discover a neuron whose square root was the distance between two objects. is this a non-linear representation of distance? or, given the form of netwon's law, is the square of the distance a more natural way for the model to represent the feature, and thus considered a linear representation? as this example shows, what constitutes a natural feature may be in the eye of the beholder. ## conclusion in this work we demonstrated that the emergent world model in othello-playing sequence models is full of linear representations. previously unbeknownst, we demonstrated that the board state in othellogpt is linearly represented by encoding the colour of each tile relative to the player at each timestep (mine, yours, empty) as opposed to absolute colour (black, white, empty). we showed that we can accurately control the model's behaviour with simple vector arithmetic on the internal world model. lastly, we mechanistically interpreted multiple facets of the sequence model, analysing how empty tiles are detected, and linear representations of which pieces are flipped. we find hints that multiple circuits might exist for predicting legal moves in the end game, as well as further evidence that residual networks iteratively refine their features across layers.