Link to original article

Welcome to The Nonlinear Library, where we use Text-to-Speech software to convert the best writing from the Rationalist and EA communities into audio. This is: Extracting and Evaluating Causal Direction in LLMs' Activations, published by Fabien Roger on December 14, 2022 on The AI Alignment Forum.This post was written by Fabien at SaferAI. Simeon has prompted Fabien in relevant directions and has provided valuable feedback.Thanks to Jean-Stanislas Denain, Alexandre Variengien, Charbel-Raphael Segerie, and Nicole Nohemi for providing helpful feedback on early experiments and drafts of this post.In this postI describe a method to evaluate that a specific direction encodes information relative to a specific concept,I use it to evaluate how good directions found through probing techniques are,I present a way of using causal intervention to find directions that maximally encode information relative to a specific concept, which I call Causal Direction Extraction (CDE),I present some preliminary results about the directions found through Causal Direction ExtractionOne of the main goals of this post is to get suggestions on further experiments to run. I think the current data I gathered is not strong enough to prove or disprove claims like “CDE finds a direction, which mostly captures the concept of gender in GPT-J”, but further experiments could.The code for all experiments is available here.How to Evaluate How Much a Direction Encodes a ConceptWhat It Means for a Direction to Encode a ConceptLet’s say you have a neural network N. Let N<=L be the network up to a given layer L, and N>L be the network after that, such that for all input x, N>L(N<=L(x))=N(x).The claim that a (normalized) direction d in the space of activations A encodes a target concept between layers L and L+1 can be understood as follows:Activations after layer L can be cleanly cut into two pieces:c(x), the projection of N<=L(x) along d, which only encodes information relative to the target conceptrest(x), the projection of N<=L(x) orthogonal to d, which only encodes all information not related to the target conceptAnother way to put it is that there exists a function c:XRd which only depends on properties of x related to the concept and a function rest:X(Rd)⊥ which depends on every other characteristic of the input, such that for all input x, N(x)=N>L(c(x)+rest(x)).Using Activation Patching to Evaluate How Much a Direction Encodes a ConceptWith this definition of what it means for a direction to encode a concept, you can use activation patching (introduced by Redwood’s Indirect Object Identification paper, Wang 2022) to quantify how well this hypothesis works. The setup here is simpler than in the original paper, since we patch the activations directly in the residual stream (between two layers instead of at the output of one attention head).Let’s take the concept of gender, and let’s run the network on the two following sentences: xA = “but her favorite color is”, xB = “but his favorite color is”, which only differ by the concept of gender (we’ll discuss this claim later). If the hypothesis is correct, then rest(xA)=rest(xB), which means we can patch activations after layer L:N(xA)=N>L(c(xA)+rest(xB))Let’s define Nd(xA,xB)=N>L(c(xA)+rest(xB)). The first argument is the one we should determine the output of the function if the hypothesis is correct, the second argument is a “distraction”.How much this holds can be quantified in different ways:Measure the KL-Divergence between the relevant output N(xA) and the mixed output Nd(xA,xB)Measure the probability of a specific token in the mixed output Nd(xA,xB) (“pink” for example) and see how much closer it is to the relevant output N(xA) than the irrelevant N(xB). More precisely, we can define the success rate as Nd(xA,xB)pink−N(xB)pinkN(xA)pink−N(xB)pink. A success rate of 1 means that the probability of “ pink” in the mixed output is the probability given to pink by the model on the relevant i...