Identifying Linear Relational Concepts in Large Language Models (2024)

Abstract: Transformer language models (LMs) have been shown to represent concepts as directions in the latent space of hidden activations. However, for any human-interpretable concept, how can we find its direction in the latent space? We present a technique called linear relational concepts (LRC) for finding concept directions corresponding to human-interpretable concepts by first modeling the relation between subject and object as a linear relational embedding (LRE). We find that inverting the LRE and using earlier object layers results in a powerful technique for finding concept directions that outperforms standard black-box probing classifiers. We evaluate LRCs on their performance as concept classifiers as well as their ability to causally change model output.

## introduction how do large language models (llms) represent concepts, and how can we identify those concepts in hidden activations? if we can identify humaninterpretable concepts in model activations, we can analyze how concepts are created and changed during inference. identifying concept representations inside of models opens up the possibility of visualizing the computation process of a model as sentences are processed, and can help to understand incorrect or undesirable responses from the model. moreover, future work examining how concept directions arise in model weights and how models express relations between concepts may benefit from a robust method to find those concept directions as a first step. an intuitive approach when trying to locate a human-interpretable concept, like the concept of a city being in france, is to collect examples of sentences with cities that are in france and cities that are not, and train a probing classifier @xcite on hidden layers of the model, typically a simple linear classiwe then invert r(s) using a low-rank pseudo-inverse, resulting in r -1 (o). finally, we create an lrc v for each object o in the relation by applying r -1 (o) to the mean object activation e @xcite . above, we train an lre from the statement "san jose is in costa rica", then invert that lre and create linear relational concepts (lrcs) representing "located in england" and "located in china" from representations of objects "york" and "shanghai", respectively. our technique builds on work by @xcite , which models the relation between a subject s and object o as an affine linear transformation, called a linear relational embedding @xcite . while the lre work was mainly an investigation into how models represent relational knowledge, we find that inverting the lre can generate concept directions that achieve surprisingly strong performance as a classifier while also causally impacting model outputs, outperforming standard probing classifiers such as a linear support vector machine (svm). we refer to the concept direction that our method creates as a linear relational concept (lrc). an lrc represents a concept as a direction in a latent space, while also functioning as a linear classifier. figure 1 shows our method for generating an lrc. we first generate an lre for a relation, mapping subject activations to their corresponding object activations as a linear transformation. then, we perform a low-rank pseudo-inverse of the lre, mapping from object activations back to subject activations. applying this inverted lre to an object in the relation results in an lrc. lrcs beat traditional probing classifiers in both classification accuracy and causality, where causality is defined as being able to control the output of the model. for instance, we can force the model to output that "london is located in france" by subtracting the "located in england" lrc from the activation of "london" and adding the "located in france" lrc. in addition, since we use the lre only as an intermediate step to obtain lrcs, we can relax the requirement that lres must faithfully predict object output logits directly. this allows us to train the lre using object activations before the final model layer, and use all object token activations rather than only using the first object token. this improves classification accuracy for both singletoken and multi-token objects compared with the original lre work, where only the final object layer can be used and only the first object token can be modeled. in this paper, we investigate the problem of locating human-interpretable concepts within the hidden layer of auto-regressive llms such as gpt @xcite and llama @xcite . we evaluate our technique using the lre relations dataset @xcite in both multi-class classification accuracy, and causality (the ability of concepts to modify model output). our technique achieves high scores for both classification accuracy and causality across the four concept types in the dataset. our contributions include: (1) extending lres to handle multi-token objects, (2) using nonterminal model layers for the object activation, and (3) using inverted lres as an intermediate step to find concept directions (lrcs) in subject activations. our code is available on github foot_0 . ## background previous work on transformers has shown that features are stored as directions within the latent space of the model's hidden activations, known as the linear representation hypothesis @xcite . further work has shown that mid-level multilayer perceptron (mlp) layers in transformer llms act as key-value stores of information @xcite . these mlp layers enhance the final token of the subject of the sentence (e.g. the token "lin" in "berlin is located in the country of") with this information in factual relations @xcite . ## method our method finds an lrc, represented as a concept direction vector, v, for a given human-interpretable concept in the hidden activations of a transformer llm model at a layer l. since we are interested in concepts as directions, we do not add a bias term and focus on learning only a single unit-length vector to represent the lrc. formally, we consider an auto-regressive model g : @xmath0, . . . , x t ] ∈ x , x @xmath1. internally, g has a hidden state size h, and has l layers. the hidden activations of layer l of g at token i is represented by h (l) @xmath2. we follow the example of @xcite and @xcite , and consider statements of the form (s, r, o) consisting of a subject s, relation r, and object o. the statement "paris is located in the country of france" would have the subject "paris", object "france", and relation "located in country". our definition of a concept corresponds to a relation and object pair (r, o), which operates on the activations of the subject s. so in our case, we would learn an lrc for the concept "located in country: france", and would expect the lrc to have high similarity with the subject activations of "paris", but not "berlin" or "tokyo". we make the following changes to the original lre method by @xcite : (1) we use the mean of all object token activations rather than only the first object token activation to better handle multi-token objects. (2) we relax the requirement that only the final layer can be used for object activations, since we find that classification performance improves using earlier object layers. both (1) and ( 2 ) are possible because we do not directly evaluate the lre using faithfulness as in the original lre work, instead performing all evaluations on the lrc operating on the subject. (3) we restrict training samples for the lre to only contain examples where the model answers the prompt correctly. if the model does not answer a prompt correctly, we assume that the conceptual knowledge we hope to capture in the lrc is not present, and that the sample will likely be noise. for instance, if the model responds to the prompt "paris is located in the country of" with "japan", we would discard this prompt. for a relation r, we have a set of possible objects o, and each object o has a corresponding set of subjects s o . we first assemble prompts that elicit each object @xmath3. for example, for the relation "located in country", prompts follow the template "{} is located in the country of" where "{}" is replaced with the subject and the model is expected to predict the object. some examples of prompts and their corresponding objects are shown in when building an lrc for relation r and object o, we assume a set of prompts each containing their paris is located in the country of france suzhou is located in the country of china manaus is located in the country of brazil we first select n prompts for the relation r, balancing the prompts to have as even a distribution of prompts across o as possible. following hernandez et al. (2023b), we train a lre r(s) consisting of a weight matrix w and bias b using these prompts, however, in contrast with hernandez et al. (2023b), we calculate the weight matrix w using the jacobian of the mean of all object tokens relative to the subject, not only the first object token. this change means we model f (s, c) as this is identical to the original lre formulation if the object consists of a single token. we ignore the β scaling factor from the original lre definition. lrcs are normalized to have unit length, removing any scaling applied to the lre. our definition of an lre, denoted r(s), is thus simplified from equation 1 as follows we then invert equation 3 to map object activations to subject activations. following hernandez et al. (2023b) , we use a low-rank pseudo-inverse, denoted r † rather than the full matrix inverse r -1 : to calculate the lrc v o for o, we take the mean of all samples of r † (o) for each prompt (s, r, o) in our training set: finally, we normalize the lrc direction to have unit length: v @xmath4. ## results we evaluate our method using the relations dataset from @xcite . the dataset contains 47 relation types, and over 10,000 instances in total. the dataset divides relation types into four categories: factual knowledge, linguistic knowledge, commonsense knowledge, and implicit bias. a subset of data from a sample relation is shown in we evaluate against both llama2-7b @xcite and gpt-j @xcite . we focus on llama2-7b for our analysis as it is a more advanced model than gpt-j, but we include full results for gpt-j in appendix a. gpt-j is included as this model was used in the original lre paper. we evaluate our performance using classification accuracy and causality. for classification accuracy, we treat each relation as a multi-class classification problem, where the lrc with the largest dot product with the test subject activation a is considered to be the predicted object ŷ: to evaluate causality, we randomly pick a counterfactual object o c for each subject in a relation and edit the subject token activations to predict the new counterfactual object o ′ instead of the original object o. we subtract the original lrc from the final subject token activation at all layers, and add the new lrc. for instance, we may attempt to edit the prompt "paris is located in the country of" to predict "germany" instead of "france" by subtracting the "located in country: france" concept and adding the "located in country: germany" concept. lrcs are all normalized to unit length, so we scale by a hyperparameter β ∈ [0, 1] multiplied by the l2 norm of the subject activation before adding or subtracting them. the causal edit at a layer l is thus calculated as below: the causality intervention is successful if the probability of predicting the counterfactual object o ′ after the edit is higher than the probability of predicting the original object o. for multi-token predictions, we use the minimum probability across all predicted tokens to avoid penalizing objects that require more tokens to represent. experimentally, for each relation, we split the dataset into a 50%/50% train/test split by relation and object, ensuring at least one training example per object in the relation. we prepend four other examples from the same relation to each training prompt as fewshot examples. we train using few-shot prompts from the relations dataset, but evaluate using zeroshot prompts, following the procedure in the original lre paper. an example few-shot prompt is shown in figure 2 . we repeat five times with different random seeds for train/test splits, reporting mean and standard deviation. the shaded area in tokens llama2-7b gpt-j 1 2393 2108 2 451 39 3 371 2 4 107 6 5+ 4 0 the superlative form of bad is worst the superlative form of bright is brightest the superlative form of angry is the plots corresponds to this standard deviation. some relations contain a one-to-one mapping between subject and object, so it is impossible to create a test split with unseen subject/object pairs. for example in the relation "capital city of country", a country has one capital city, and a city is the capital of only one country. since our concepts require a unique r and o pair, we cannot evaluate these relations and exclude them from evaluation. we also exclude any samples the model answers incorrectly, and we exclude any relations with less than five test samples as few test samples make it hard to evaluate performance robustly. when training lrcs using our method, we use 20 training samples per lre for the main benchmark, and 5 training samples for sweep plots. we use rank 192 for pseudo-inverse. calculations are performed using a single nvidia a100 gpu with 16-bit quantization. we use subject layer 17 and object layer 21 for llama2-7b, and subject layer 14 and object layer 20 for gpt-j. ## related work previous work on understanding neural networks focuses on individual neurons @xcite . however, individual neurons have been found to activate in response to multiple concepts, making a clean understanding difficult @xcite . indeed, transformers can represent more concepts than they have neurons in their hidden layers @xcite . a source of inspiration of our work is knowledge editing in lms, specifically rome @xcite and remedi @xcite . in rome, factual knowledge is shown to reside in the mid-layer mlps of language models, and can be edited by updating a mid-layer mlp to insert any fact desired. in remedi, model outputs are edited by adding a vector to the subject of a sentence during forward inference. this is similar to our work in that this vector can be said to contain the concept that is desired to be elicited. however, the goal of remedi is to edit model outputs rather than identify concept directions and build a classifier as in our work. we also take inspiration from probing classifiers @xcite . probing classifiers are linear classifiers which operate on hidden activations inside of neural networks. tcav @xcite can be said to be a probing classifier for vision models, where a classifier is learned at multiple layers in the model. most similar to our work, @xcite build a probing classifier for textual games from lm hidden activations, and show that these hidden activations encode a basic world model. however, this work focuses on encoder-decoder models, and does not attempt to classify arbitrary human-interpretable concepts beyond the text game. closest to this paper is work on lres in llms @xcite , which is the source of our evaluation dataset and is the first step in our method. this work also attempts to estimate relations, and learns a linear mapping from the subject token activation to the first output token of the object. however, as lres only map to the first object token, they struggle with multi-token objects. for instance, an lre evaluated for faithfulness cannot distinguish between "bill gates" and "bill clinton" as they both begin with the same token. in addition, the original lre work is presented as an exploration of how llms encode relations rather than attempting to build a classifier or find concept directions. ## conclusion identifying and classifying a broad set of humaninterpretable concepts in language model activations is a vital step towards understanding how language models operate. in this work, we have shown a technique for identifying and classifying concepts in model hidden activations called linear relational concepts (lrcs). we show that lrcs outperform standard linear classifiers like svms on both classification accuracy and causality. while our technique performs well, there is a variance in performance depending on the training samples chosen. we expect further improvements to be achieved by optimizing the lre training samples chosen for each lrc. in addition, it may be possible to combine lrcs learned at multiple layers to achieve even better results to get around the causality / accuracy trade-off depending on the layer chosen to train the lrc. in the future, concept identification techniques such as lrc may make it possible to investigate the relations between concepts within model weights, and extract knowledge and even world models directly from pretrained language models.