Representational Analysis of Binding in Language Models (2024)

Abstract: Entity tracking is essential for complex reasoning. To perform in-context entity tracking, language models (LMs) must bind an entity to its attribute (e.g., bind a container to its content) to recall attribute for a given entity. For example, given a context mentioning “The coffee is in Box Z, the stone is in Box M, the map is in Box H”, to infer “Box Z contains the coffee” later, LMs must bind “Box Z” to “coffee”. To explain the binding behaviour of LMs, existing research introduces a Binding ID mechanism and states that LMs use a abstract concept called Binding ID (BI) to internally mark entity-attribute pairs. However, they have not directly captured the BI information from entity activations. In this work, we provide a novel view of the Binding ID mechanism by localizing the BI information. Specifically, we discover that there exists a low-rank subspace in the hidden state (or activation) of LMs, that primarily encodes BIs. To identify this subspace, we take principle component analysis as our first attempt and it is empirically proven to be effective. Moreover, we also discover that when editing representations along directions in the subspace, LMs tend to bind a given entity to other attributes accordingly. For example, by patching activations along the BI encoding direction we can make the LM to infer “Box Z contains the stone” and “Box Z contains the map”.

## introduction the ability of a model to track and maintain information associated with an entity in a context is essential for complex reasoning @xcite @xcite @xcite . to recall attribute information for a given entity in a context, the model must bind entities to their at-tributes @xcite . for example, given sample 1 and 2, a model must bind the entities (e.g., "box z", "box m", "box h", "alex", "john" and "carl") to their corresponding attributes (e.g., "coffee", "stone", "map", "bean", "pie" and "fruit") so as to recall (or answer) such as what is in "box z" or what is sold by "alex" without confusion. binding has also been studied as a fundamental problem in psychology @xcite . to uncover how language models (lms) realize binding in term of internal representation, @xcite introduce a binding id mechanism and state that lms apply a abstract concept called binding id (bi) to bind and mark entity-attribute (ea) pairs (e.g., "box z" and "coffee" in sample 1, where bi is denoted as a numbered square). they also claim that the bi is represented as a vector to be added on the representation (or activation) of a ea pair so that the common vector is used as a key clue to search attribute for a given entity. however, they have not captured the ordering id (oi) information from the entity (or attribute) activations that causally affects binding behaviour and thus bi information as well. here, oi is defined as the input order (or ordering index) of entities and attributes, no matter they are bound by a relation (e.g., "is_in" in sample 1) or not, such as the indexing number in sample 1 and sample 3. we can observe that in a 1e-to-1a bound context, such as in sample 1, bi and oi are interchangeable. oi information into a subspace of lms' activations that primarily encodes the order index of entities and attributes in a given context. what is more, the discovered oi subspace plays a crucial role in the in-context binding computation. in addition, we find that such oi subspace that determines binding is prevalent across multiple lm families such as llama2 @xcite (and llama3 (ai@meta, 2024)), qwen1.5 @xcite and pythia @xcite , and the code fine-tuned lm float-7b @xcite . ## finding oi subspace in this section we describe our principle component analysis (pca) based method to localize the oi subspace in activations of lms. as shown in figure 1 , we firstly extract entity activation from lms. given a lm (e.g., llama2), and a collection of texts which describes a set of ea pairs related by a relation such as "is_in" in sample 1, we extract the activation of entity token (e.g., "z") in query (denoted as x i ) from certain layer foot_0 and construct a activation matrix m r p r nˆd for a relation r, where n denotes the number of entities and d denotes the dimension of the activation. the row i of m r is the activation of an entity token (i.e., x i ). pca has been applied for identifying various subspace (or direction) such as the subspace encoding language bias @xcite , truth value of assertions @xcite and sentiment @xcite . inspired by these studies, we choose pca as our first attempt to localize oi subspace. in addition, we also apply other di- figure 2 : layer-wise oi subspace visualization on llama2-7b, where "bi" primarily denotes oi. we hypothesize that in a activation subspace, entities with the same oi tend to cluster together (w.r.t the ones with different ois), even though these entities have different semantic meaning, and the ois are encoded as directions (or a pc) in the subspace. for convenience, we number ois in leftto-right order, and the leftmost oi" 0. since pca is applied to identify the principle directions from a multidimensional space, we leverage pca to capture oi subspace (or direction) from activations of lms. specifically, the pca of a activation matrix is m r " u r σ r v t r , where the columns of v r p r dˆd are principle directions of m r . we takes first c columns of v r as the oi direction, denoted as b r p r dˆc . ## oi subspace visualization we adopt a subset of the entity tracking dataset @xcite , which contains n " 1000 samples, to create layer (l) wise activation matrix m l r . we then uses the m l r to extract the layer-wise oi subspace pro-jection matrix b l r p r dˆ2 to visualize the activations. figure 2 shows the embedding visualization on llama2-7b, where each point represents the activation of an entity projected via the b l r , and the colors represent ois. from which, we can observe that middle layers, such as layer 8, have a clearly visible direction along which oi increases, while the others have tangled distribution. we also observe similar pattern of distribution on llama3-8b, float-7b ( §a.5) and other lm families such as qwen1.5 and pythia ( §a.6). this indicates that lms use the middle layers to encode oi information, and the finding is prevalent across multiple lm families. this finding is also consistent with the "stages of inference hypothesis" @xcite stating that the function of early layers is to perform detokenization, middle layers do feature engineering, and late layers map the representations from the middle layers into the output embedding space for next-token prediction. according to the hypothesis, we would expect to find the ordering feature most prominently represented in middle layers, which is exactly what the visualization shows. we call this dimension that represents oi as oi principle component (oi-pc). in the following section, we apply causal intervention on the oi-pc to analyze how oi-pc affect the model output. ## causal interventions on oi subspace by projecting the activation matrix m r into the oi subspace, we have found a correlative evidence for the existence of the direction (i.e., oi-pc) that encodes oi information. however, it is possible that the oi information is encoded in the oi subspace but has no effect on lms' binding behaviour. in order to test if ois are not only encoded in the oi subspace, but that these representations can be steered so as to swap the binding and change lm's output, in this section, we perform interventions to analyze the causality. that is, we want to find out if making interventions along oi-pc leads to a change in lm's binding computation. ## related work linear representation recent research found that sequence models trained only on next token prediction linearly represent various semantic concepts including othello board positions @xcite , the truth value of assertions @xcite , sentiment @xcite , and numeric values such as elevation, population, birth year, and death year @xcite . continuing this line of research, in this work, we discover that multiple lms such as llama2 can also linearly encode oi along a oi increasing direction in the activations. pc1 pc2 pc3 0 0.2 0.4 0.6 0.8 1 related yes no correlation figure 10: oi-pc based correlation between attributes and their corresponding entities, where "pci" denotes the i-th pc of the oi subspace, "yes" and "no" represent the related (i.e., original) and non-related dataset respectively. see appendix ( §a.1) for additional related work. ## conclusion and future work in this work, we study the in-context binding, a fundamental skill underlying many complex reasoning and natural language understanding tasks. we provide a novel view of the binding id mechanism introduced by @xcite that there exists a low-rank subspace in the hidden state (or activation) of lms that primarily encodes the ordering information and which is used as the prototype of bis to causally determine binding. our future work includes: 1. the analysis of oi subspace in a more realistic setting; 2. the study of interaction between in-context binding and factual knowledge learned from pretraining; 3. oi subspace based mechanistic analysis.