Learning Transformer Programs with Dan Friedman - #667

0:09All right, everyone, welcome to another episode of the Twimble. AI podcast. I am, of course, your host, Sam Charrington. And today I'm joined by Dan Friedman. Dan is a PhD student in the Princeton NLP group. Before we get into today's conversation, be sure to take a moment to head over to Apple Podcasts, Spotify, YouTube, or your listening of Platform of Choice. And if you don't already subscribe to the show, be sure to hit that subscribe button. Dan, welcome to the podcast. Thank you for having me. I'm really looking forward to digging into our conversation. We'll be talking about your research that will be presented at this year's NeurIPS.

0:45on the topic of interpretability for transformers, and you take a unique approach to that challenge. Before we get into that, though, I'd love to have you share a little bit about your background and how you came to work in the field of machine learning. Sure. So I'm currently a fourth-year PhD student in computer science at Princeton. I took maybe a roundabout path to machine learning. So my undergraduate degree was in English literature. Oh, wow. And yeah, I kind of got into, I just took an introductory programming class and a junior or senior year. And I liked it a lot and I kind of gradually pivoted.

1:22That's awesome. Yeah. So I did a little bit of research as an undergraduate, kind of in NLP and sort of computational linguistics a little bit. And I worked for a few years as just a software engineer. And then I decided to go back to graduate school to kind of see what was going on in the research world. Awesome. Tell us a little bit about your research interests. What have you been studying over the past four years? I'd say what I'm mainly interested in now is the kind of general topic of interpretability. So I'm trying to better understand how these machine learning models that we have now can process natural language text.

1:59Awesome. And are you specifically focused on transformers exclusively in your research or are you kind of exploring the topic of interpretability broadly? That's a good question. I think like the research landscape has changed a bit just in the last couple of years. So now kind of like transformer based language models are just like the sort of predominant architecture, not just in NLP. And so that's kind of what I've been focused on more recently. But I'm also, maybe more so in the past, have been interested in kind of like more sort of architecture agnostic approaches to interpretability, if that makes sense.

2:43Sure. Some of the most popular approaches to interpretability, things like Lime have been kind of come out of that architecture agnostic approach. But I think what you are demonstrating in the research that we'll be talking about is that if you are aware of the architecture and what's more, if you're able to manipulate the architecture, you can do a lot more with interpretability. Is that broadly the general idea you pursue with the research? Yeah, I would say so. I think like another dimension of the point is just that the transformer has sort of proven to be incredibly effective architecture for sequence processing tasks.

3:23And so you can also think if it maybe is a way to sort of work backwards and say, could we take this model that works really well and try to like make it a little bit easier to understand the sort of rules that it's learning for all these tasks. And now your research falls broadly under the category of mechanistic interpretability. Can you share what that means and talk a little bit about some of the approaches that we've seen thus far in the vein of mechanistic interpretability for transformers in particular? Yeah, this is a very good question. I don't know if mechanistic interpretability has such a precise definition, but I think in general, the goal is to try to take these neural networks that we train to solve a task and try to reverse engineer what the model is doing in terms of some more human understandable algorithmic components.

4:21right so the models are operating in terms of high dimensional representations and the goal of mechanistic interpretability is to try to see if there's any kind of analogy between what the models are doing in kind of high dimensional space and some more familiar types of algorithms that we can think about and reason about as people the mechanistic distinction is versus approaches again like lime i guess or like observational or um surrogate models i guess is another approach like things that aren't like digging deep and trying to interpret what's actually happening in the model but rather looking at you know outputs or relationships between inputs and outputs and that kind of thing yeah i'd say one example of kind of a whole category of approaches in sort of earlier work on interpretability would be things like feature importance.

5:16So say you have a sequence of texts and you have maybe a classifier and their interpretability method will output, you know, here are the words in your sentence that were important for making that prediction. Or kind of another kind of approach is finding examples in the training data that were influential for that particular prediction. And so these are kind of methods that can give you some hints about what the model is doing, but they don't really tell you, they don't give you what I would call maybe like a sort of algorithm level understanding. So kind of one specific test is you can't, one of these say feature importance maps and then predict what the model will do on a new example.

5:58So I think maybe one goal or ambition of mechanistic interpretability is to give you kind of a representation of what the model is doing that you can sort of look at and simulate yourself. And it strikes me that in addition to just solving the problem of providing some interpretability, this research is also aiming to take on another important task, which is to help us understand these models, which are increasingly important, doing amazing things, and we're not really always sure why or how they're working. Yeah, 100%. Yeah. So in the paper, you talk a little bit about some of the prior work.

6:40There are approaches that look at attention heads. There are approaches that look at various network representations. Can you kind of characterize the prior approaches and their, I guess, shortfallings relative to the approach you're taking here? Which, of course, we haven't gotten to yet, but we will. Shortcomings, I guess I should say. Yes. So there's just a ton of prior work on interpretability. I'm not going to do it justice, but I would give as a very kind of simple summary, I kind of think of two categories of prior work. So one is sort of line of work, just sort of in the NLP world on trying to understand models, like maybe beginning with BERT.

7:23So the BERT model, that kind of sparked a lot of interest into, how is this model so good at all these NLP tasks? You know, that work uses just a whole range of, different kinds of methods for trying to inspect different aspects of what the model is doing. So some of the methods I mentioned, like looking at feature importance or finding important training instances. And there's also some work on peering into the internals of the model, like figuring out what information is encoded in the model's representations, or trying to understand the role of different attention heads. And I think to me, what I see is the main limitation of this body of work, or first I should say that it is, all this work has revealed a lot of very interesting things about how these models are kind of working, but it doesn't give us what I was kind of describing before as sort of algorithm level understanding.

8:14So it doesn't give us a complete or faithful description of how the model is making decisions. And so it's not really a type of understanding that you can use to say, be confident that your model is always going to do what you intend. So for example, you can't verify your model in the same way that you could verify software. That's kind of one category of work. Then I think more recently, there's been a growing interest in this new approach, which comes under the title of mechanistic interpretability. I think maybe responding to some of those limitations of the existing work in interpretability.

8:49And with mechanistic interpretability, the goal is really to kind of reverse engineer an algorithm them that you can sort of read and say, here's how the model is performing this task. It's represented as kind of neural network continuous computation, but actually we can sort of translate it to a format that's easier for us to understand. And so this has largely been focusing on, you know, these large transformer language models, although there's also a lot of interesting work on small transformer language models. And I think, again, yeah, mechanistic interpretability has also revealed a lot of interesting things about these models, but it's quite hard to do.

9:25So it requires a lot of manual effort. That's sort of maybe problem one. And I think maybe a more general problem is that you actually, you know, can't always be assured that there even is a kind of a human understandable description of what the model is doing, right? So these models haven't been trained to be interpretable. And so they could be processing information in all sorts of subtle and complicated ways. And so part to try to have some kind of automated interpretability and also to maybe give us a way of specifying more of our own kind of preferences about what kind of like explanation do we want.

9:58So maybe we can place more sort of bias the model to be more interpretable in some sense that we can specify. And now the approach that you've taken, I forget if you use this term in the paper, but design for interpretability kind of comes to mind for me and kind of brings to mind an effort to design for manufacture in engineering and automotive engineering, for example, or engineering of goods and or design for maintainability from a software perspective. It's like you've got this end goal, treat that end goal as central in a design process and come up with an approach that thus helps deliver it.

10:39What were some of the inspirations for you in taking this approach? So I'd say first there's been some prior work in machine learning kind of advocating for a similar goal of what is also sometimes called like inherently interpretable models or kind of intrinsically interpretable models. And a lot of that has come more in fields like computer vision and for models operating on, say, tabular data and things like that. And so that's kind of a motivation in general, but there hasn't been too much work in natural language processing. Then I'd say a more immediate inspiration was a recent paper from Gail Weiss and collaborators called Thinking Like Transformers.

11:22And so this is a paper that proposes a programming language called RASP, which you can use for writing programs in sort of a Python-like programming language that can be compiled into a transformer network. This was kind of proposed as a way of making it easier to reason about the kinds of things that transformers can learn to do. So it's kind of a computational model for the different types of information processing that this network architecture is capable of. But it also gives us this way of seeing the correspondence between the transformer and kind of a human-readable computer program. How is this type of program different from PyTorch or TensorFlow program, which is also a Python program that represents a transformer?

12:07This is a great and sort of philosophical question. I would say it's an alternative representation of what the model is doing. I mean, okay, so one kind of like direct difference is the PyTorch program, you can't read that code and understand what the model is doing because you have to also look at the weights of the model to understand the actual sort of solution. It's not really doing anything until it's trained. Right, exactly. And, you know, you can plug in different weights and get completely different behaviors. So I think the analogy I'm thinking of is more like you take, say, a trained transformer.

12:39So you have it's defined by the architecture, but really by the weights. And can we kind of turn that into sort of human readable program? So, you know, a human readable description of what those weights mean, if that makes sense. But, you know, we sometimes think of a model as like a function. And in this case, the representation is richer than a simple kind of function transformation type of representation. Does that question make sense? Well, I'm not sure I agree, but I also might misunderstand the question. Yeah, so I guess I'm thinking of if you could express a transformer as like some, you know, if there was some way to mathematically express like some nonlinear transformation between the inputs and outputs as a transformer, that is kind of like what the program is doing, but it's not particularly interpretable and not as manipulable in the same way.

13:38I guess that was the thought. I don't know if it is helpful or gets us anywhere. Yeah. I wonder if it makes sense to circle back to this after we've gotten a little bit further, a little bit further. Yeah. To me, in a way, a fundamental sort of in a way, philosophical question about like, what does it mean? So you have, you could have two representations of the same, the same system that are equivalent in some sense, but one is sort of easier for us to understand. Okay. Let's come back to it. Yeah. Okay. You're inspired by this previous work, including RASP and their Python representation of trained transformer models.

14:14And so tell us now about your approach and what you've done. So RASP kind of is this framework for writing down a program. So like you can write, say, a transformer program that will solve a particular task, say a program that will sort a list. And it's designed so that this program is equivalent to particular instantiation of a transformer. So you can take the program and compile it into a set of weights for a transformer that will implement this program. But there's no way of training a transformer and decompiling it into a program. This is what we want to do, right? So, you know, in Rasp, you have to write all the programs yourself, and then you can turn them into a transformer.

14:55But we'd like to just train a transformer, it can figure out its own solution, and then convert it back into this alternate representation, sort of program representation that we contend is easier to interpret. So that's kind of the high level goal. And without getting into the weeds, basically what we do is we just design a set of constraints on the transformer weights that will guarantee that the kind of the corresponding model will always have a direct mapping to a sort of a RASP style program. So one way to think about it is that every RASP program corresponds to a transformer, but not every transformer corresponds to a RAS program.

15:34And so we want to kind of constrain our optimization to the subset of transformers that do correspond to programs. And so this is a, it's kind of a discrete optimization problem. So these are transformers that have discrete weights. They're kind of like constrained in various ways to keep things sort of clean and interpretable. And so we design a set of constraints and then can kind of optimize this sort of discrete model using a kind of a continuous reparameterization. So we use the Gumpel reparameterization to search in the space of discrete models. Okay. To get a little bit further into the weeds, the constraints seem to be a heart of how you're able to construct this, you know, essentially a space that only contains representable transformers.

16:21The first of those is you require disentangled residual streams. What does that mean? Yeah. So this word, this phrase is kind of inspired by some prior work in mechanistic interpretability. And this is sort of responding to a problem with, you know, if you just train a transformer from scratch, it's not necessarily going to be interpretable in part because all sorts of information is jumbled up together in the representations. So every layer can calculate something and the transformer has these residual connections. So each layer will add its output back to the input, you know, from the previous layer.

16:57And so actually a lot of work in mechanistic interpretability is aimed at just solving exactly this problem. Like how do you trace information through different components of the network, given that they're all kind of like over, you know, writing to the same place. And it's sort of by analogy to a computer program, you can think of the residual, the residual stream, the model's representations as kind of corresponding to memory in a symbolic program. And in an off-the-shelf transformer, all the variables are kind of being written on top of each other. Or you can think of like the same variable means different things depending on like the context.

17:32And so this first constraint is just to make sure that each kind of piece of information is stored at a dedicated address in the embedding space. And whenever a sublayer reads information from the residual stream, it's kind of constrained to only read a specific named variable from a dedicated address. And is this constraint, how's it implemented? Is it like an objective function constraint or is it a constraint between two training steps? It's a hard constraint. So is it architectural in the sense of? Yeah. Well, maybe I can explain kind of more concretely. There's sort of like, you can kind of divide this into two parts.

18:13So one is reading information from the residual stream, you know, so reading a variable, and the other is writing a variable. So when we start out the first layer of the transformer, we only are storing two variables. So one is the token identity. So if one variable that says what is the identity of the word at each position, and then we also have a position encoding that tells us what's my index in the position. So those are kind of two variables, tokens, and positions. Maybe the writing part is easier to explain first. So we have each layer, say we have each attention head and each feedforward layer will calculate something based on the previous input.

18:53And then it will write some new information to the residual stream. And in a standard transformer, this is just done by addition, but we just replace addition with concatenation basically so each attention head will get its own kind of new block in the embedding space you think of its own dedicated subspace in the representation space that's reserved exclusively for its use so that's kind of you could think of it as an architectural modification you implement the constraint by not doing the thing that the constraint that you don't want to happen essentially yeah this is like realized in the you know in the code right so it's it's it's kind of like you could realize this kind of constraint with a regular transformer right so instead of concatenating you could just have you know pad all of your embeddings with zeros and add but you know the model could learn to to do this kind of thing itself but we don't think it has much of a motivation to so that's kind of writing to the residual stream so you kind of keep all the information disentangled and then we also want to make sure that each model is reading kind of a named variable and so in a standard transformer this sort of you can think of reading information from the residual stream is sort of corresponds to taking some kind of linear projection of your embedding.

20:07So you have some learned weight matrix. You can think of it as reading off some information from some subspace of the residual stream. And we want to constrain the model so that this subspace corresponds exactly to one of these like named subspaces. So maybe we'll have some attention head that will read the token variable as one piece of information. And so the way we implement the constraint is by instead of kind of learning this whole sort of a free ranging weight matrix, we just learn a gate over the available variables. So let's say you have two possible variables you can read in the first layer.

20:47Say attention head for each piece of information it has to read, it will learn kind of you think of it as a distribution over two possible variables. And this is the sort of thing that we train. And then once it's kind of like picked a variable to read, you can kind of like project that back up to get the corresponding weight matrix in a regular transformer. Got it. Talk me through the second constraint. You talk about modules in the paper. What does that constraint represent? So another way of describing that is that we have, you know, you can think of our program as we have some number of variables and we have some number of functions and each function will read some variables that have been kind of written earlier in the program.

21:29It'll calculate something and then write an output to a new variable. And so the second constraint is referring to that middle part of each module will have to calculate something. So we also want to make sure that each model is calculating something in a way that we can translate into kind of an interpretable sort of rule-based procedure. So we want to make sure that each sub-layer itself is interpretable as well. Okay, awesome. And so you've got these constraints. How do you put those together to then generate the programs? Yeah, so in a sense, we can kind of read the program off directly from the weights in this constraint model.

22:08so we kind of know what all the variables are and we know what each module is reading. So let's say in our first layer, we can kind of look at the attention head, say. So an attention head, it's associated with three variables that it has to read, a key variable, a query variable, and a value variable. And so during training, it'll learn a distribution over the possible variables for each slot. And then after training, we can just sort of read it off, You know, say, what does it pick as sort of the most likely variable? And we can say, okay, attention head one is reading the tokens variable as the key, and maybe the tokens variable as the query, and the position variable as the value, and so on.

22:51So that's the first kind of bit of information. And then we can look at each module itself. So we have our attention heads and feed-forward layers. And the exact procedure depends on the component, But those also we can kind of read off that are equivalent to the logic that that particular model is, module is implementing. And so the programs remain probabilistic because you're implementing that distribution or the distribution at these various layers in the programs? The program is not inherently probabilistic, although you could implement a probabilistic program within this framework in some sense.

23:31So to explain, we learn this distribution. The way I would summarize it is we are optimizing a distribution over discrete programs. And so we have all these gates, we have these distributions. So for example, for this attention head, we're trying to search for this distribution over which variable will attention head one read as the query variable. After training, we can just sample a discrete program from these distributions. Or what we do in our experiments is just take the sort of the arc max discrete weight associated with each distribution. So just take the most likely program. And so this is a discrete program that is not probabilistic.

24:16Got it. So at the end of your process, you've kind of baked in all of the sampling and there's ultimately a program that has a single output for a given input. Yeah. Maybe one more detail. The final layer of the model is just a regular linear classifier. That's the one kind of continuous part of the model. And the way I would think about it is we have the main kind of guts of the transformer is a feature extractor. And so it's a discrete feature extractor. And so every kind of sub path through this program is sort of answering a question about the input sequence. And so it's saying, yes or no, here's some bit of information about the input sequence.

24:58And that's all discrete and that's all deterministic. But then you have a linear classifier that you can use in kind of the standard ways that you would use a linear classifier in machine learning. So you can have kind of a stochastic decision rule, for example. Okay. Okay. By talking about reading the program from the model, I think was the terminology you used, is there a one-to-one mapping, meaning every representation has a single program and presumably that falls directly from just the disentangled residual streams that gives you that one-to-one mapping? Or is it both of those constraints together give you that one-to-one mapping?

25:40It's both of the constraints together. Yeah. So there's a one-to-one mapping. There's some freedom in how exactly we express the mapping. So you can kind of write this sort of rule-based program in different ways, but the rules in a sense are kind of like one-to-one. Okay. Can you articulate like what the programs look like without visual aids? Okay. So this is a challenge, but maybe I can. And there's a lot of examples in our paper. and we also have uploaded all of our learned programs to our GitHub repository. So you can download these programs and run them yourself. They're just Python programs.

26:16So I'd say the main bulk of the program is in representing attention heads. So what does each attention head do? The attention head, the kind of the role of each attention head is defined by what variables is it reading? So what information does it read? You know, so say what keys and queries does it read? and then what's its sort of internal logic for deciding what token to attend to given a particular query. Maybe the simplest example is if we're kind of starting at the first layer of the transformer. The only possible information we can look at is the token identity. So say what word is at each position and also this position indicator.

26:59And so one common thing that the model learns to do is it will say, well, at each position, I just want to look at the previous position. I want to see what word appeared immediately before me in the sentence. And so in the program that's realized by, first of all, we can see that this attention head is reading the position variable as the key and the query. So that's kind of part one. And then part two is each attention head will also be associated with a rule, which we call the predicate, which tells you for every pair of key and query, should I attend to that? Or for every query, which key should I attend to?

27:40So the way it looks like in the case of the position, the look at the previous token rule is you'll have sort of a sequence of if else statements. And it will say if the query position is equal to one, attend to the key at position zero or attend to the token that has a position equal to zero and so on. So that's kind of attention. And then we also have sort of a simple kind of feedforward layer, which is also a series of rules. So we kind of express the feedforward layers as just big lookup tables. Let's say we're trying to look for a particular bigram. So one example we have in the paper is for, we train models to recognize whether all of the parentheses in a sequence are balanced.