Localizing and Editing Knowledge in LLMs with Peter Hase - #679

0:06All right, everyone. Welcome to another episode of the TwiML AI podcast. I am your host Sam Charrington. And today I'm joined by Peter Hasse. Peter is a PhD student at the UNC NLP Lab at the University of North Carolina. Before we get going, be sure to take a moment to hit that subscribe button wherever you're listening to today's show. Peter, welcome to the podcast. Hi, Sam. Thanks so much for inviting me. I'm happy to be here. I'm excited for our conversation. We're going to be talking broadly about your research interests, which span topics like interpretability, models as belief storage, scalable oversight.

0:42Take us to the next level of detail. Tell us about how you think about your research and your research agenda. There really are three areas I'd pick out that we focused on during my PhD research. First, broadly, interpretability. You know, this is all about understanding the internal reasoning processes that language models are using as they solve problems, as they answer our questions. I think this is not only really fascinating, but also really important as we seek to know, can we trust the internal reasoning that's going on in these models? Is it the kind of reasoning that we think is good and we will generalize in the ways we intend?

1:16That's the first area. And then the second, this model editing problem has gone through kind of a resurgence in popularity since 2020, 2021. It's all about updating individual factual knowledge and beliefs and language models. Also really interesting, also practically important, if not only just to keep language models up to date, there's also other kinds of use cases like deleting things we might not want them to know. And then lastly, another kind of particularly safety oriented area is this scalable oversight problem. Here, we're focused on understanding how can we supervise and how can we evaluate AI systems as they get better and better at solving tasks to the point where potentially, you know, we might not even know the answers to a problem ourselves.

1:57and we still want to be able to supervise a model and steer it in a certain direction or better calibrate our own trust and like, okay, should I trust, you know, this answer being correct or not? Do you think of the work in the interpretability domain as kind of the most fundamental of the three in a sense to the degree you can figure that out, it makes the other two tasks much easier? Oh, that's a good question. Long term, yes. But in the same sense that like, oh, neuroscience long term is like the explanation for everything we do. Right. Because it's like that's how the brain works. It's it's it can be so low level, especially if you're thinking about mechanistic interability, that if you care about certain safety problems, there can just be other potentially more direct avenues to solving some of these safety problems or thinking about properly supervising models.

2:47You don't always need to know the individual underlying mechanism behind every behavior you're interested in steering in a neural network. And in fact, historically, that's not the way it's done because historically we just set up, here's the objective function that we want to optimize. Here's what we care about. Let's make sure we're optimizing this end-to-end black box system for this objective function properly. And when we're testing the system, we're actually testing all the cases that we care about. really the way progress is typically achieved is not by understanding all the little individual internal mechanisms behind everything.

3:22Fair enough. Fair enough. So talk us through some of the individual projects within interpretability that you've been working on. Yeah. So one in the past couple of years that we put out through some work I did while I was at Google was focused especially on understanding, you know, what do we learn when we say this component in a neural network is responsible for storing this fact. So this is called knowledge localization or localization and interpretability. We want to be able to say this component, even this neuron in a model potentially, we might get that granular. This component's responsible for some behavior or some piece of knowledge in the model that we know the model has.

4:02Some of the work we did here, people are really excited about connections between this kind of interpretability research and model editing research. The idea being, okay, you want to edit a fact in a model. First, you should know where the fact is stored, and then that should help you edit the fact in the model. I still think this is broadly the correct narrative going forward, but we basically carried out an analysis where we found there might've been like a false start in this direction. So we found that some current methods for saying, okay, layer five is responsible for storing this fact in this large transformer model.

4:35But if you want to change what the model says in response to a question about the fact, you can actually go in and edit a different layer inside the model. You could add a layer 10 instead and still change what the model says. Or you could add a layer one, even though you're pretty sure the fact was stored at layer five. So this is pretty unintuitive in this view that facts are stored in specific places and neural networks. Stepping back a bit, we're thinking probably that's not exactly the right narrative for how the models work. And in this case, there was a little bit of a disconnection between the interpretability result and the model editing application.

5:13Are you saying that you question the fundamental idea that knowledge is stored in a particular place in a complex neural network, but rather it's distributed more broadly within that network? Well, really, yeah, that's a fair criticism, I think, because previously one of the notions of like why neural networks were successful is the view that, OK, there are distributed representations in neural networks. And particularly, we know that residual layers, residual connections were really important for scaling neural networks up. And transformers are built with these kinds of residual layers such that information just tends to flow across the entire network.

5:53and each individual layer in a model can add or remove information from what's called a residual stream. The overall system here ends up looking quite different from how you might imagine a cell or like an engine, like a motor engine, where this component has this job and it's both necessary and sufficient for carrying out this function. And if it breaks, then the whole thing definitely breaks. Neural networks don't really work like that. You know you can eliminate a component and the system is surprisingly robust. You can add noise and the system might still be surprisingly robust. Yeah, as much as the parallels between neural networks and the brain are sometimes overplayed, that is a parallel.

6:35There is information or capability that's localized, but we've also seen that the brain can compensate if a particular area is damaged in some way. Yeah, that's an interesting connection. And in fact, even some research is really leaning into this localization angle and asking the question, well, maybe the information isn't localized by default, but maybe we can encourage it to be more localized. So that way, we build a system that better resembles other more mechanical systems we're used to working with. And so maybe it'll actually be easier to go in and edit the system's behavior once we've designed it to be more editable to begin with.

7:11And so when you talk about localization, what's the granularity that you looked at? Are you looking at particular weights or activations or modules or layers as a whole? How far did you go with that? But yeah, we went down to the layer level, thinking about the outputs of specific, in this case, MLP or attention layers in a transformer, which is, I would say, a very intermediate level of granularity. You can go all the way down to the neuron. There's been a lot of research showing that the neuron isn't necessarily the right unit of analysis, but there's something that's called a direction in latent space or kind of a latent direction or a feature direction that is more so like the right unit.

7:55that is closely related to, you know, just the outputs of layers themselves. Yeah, the more granular you get, kind of the bigger the search space becomes. So when you're talking about knocking out, you know, earth, you're doing a kind of causal intervention on a layer versus a neuron, you know, your model might only have 30 layers, it's going to have like, you know, maybe 30 times 1000 or 30 times 4000 or more neurons. So at that point, when you're thinking about, you know, searching through intervention space, particularly in any kind of combinatorial way, things could get a little bit more complicated.

8:29The paper that we're referring to, does localization inform editing? Surprising differences in causality-based localization versus knowledge editing and language models. Did I get the right one? That's right. That's right. Yeah. And the surprising part there was, you know, we definitely saw it seemed like this fact was stored, you know, basically between layers like five and 10 or so. And you can get the model to say something different by editing layer one or like layer 15 sometimes. And we think, you know, we're speculating, actually, we can't run the experiments immediately. The key piece there is these residual layers where the model between layers five and 10, the model wants to say one thing.

9:09You can go in before or after and inject information into that residual stream that ultimately changes the model's downstream answer. even though you haven't touched like the original information store, so to speak. And you referenced in the title of the paper references causal modeling based approach. Can you talk a little bit about the way you set up the problem from a causal perspective? Yeah, I mean, this is this is the fundamental workhorse of interpretability research right now is being able to do causal interventions on components and models and exactly measure the effect on a behavior you're interested in.

9:47The specific method we were looking at for localization is called causal tracing. That's kind of a denoising ablation, like knocking a component out. So that's like taking a neuron and setting its value to zero. You knock that neuron out from the network. So you can think of that as a zeroing ablation or a noising ablation. You could add noise to the neuron's activation such that maybe the rest of the model can't really, like, you know, it's not a typical value to the rest of the model any longer. So causal tracing, kind of a new and fancy approach and successful approach to this particular form of analysis of localizing facts, focuses on denoising.

10:26This is still a causal intervention, but it happens that the twist here is that first you noise the input to the model. So now the model's kind of getting this noisy information input that's given to it. But if you have a question that's like, where is the Space Needle? Or you might input a prompt to the model that says the Space Needle is located in, and that's the whole prompt. You would go in and when the model sees Space Needle, what that transformer sees are word embeddings or token embeddings. And you'd actually take those token embeddings and add Gaussian noise to them. It's as if the model sees this fuzzy subject entity, this noisy thing is located in blank.

11:12And so it's like, okay, well, I'm clearly being asked where something is located or I need to figure out where something's located, but I don't know what the thing is. And so the denoising part here, it says, okay, where can I go in and substitute the clean space needle representation inside the model such that the model arrives at its original answer of Seattle? And so it'll happen that if you go in and substitute maybe between layers 5 and 10, you go and kind of denoise those representations, the model once again understands, okay, now I need to say Seattle. And so you think that, okay, when you do that, it's like these components seem to be sufficient at arriving for the original answer.

11:54the connection to knowledge editing is kind of the one we alluded to earlier that if you know where the information is then that gives you some leverage to edit it and and change the the output that's the intuition and i think that's a really reasonable intuition and i'm optimistic about this intuition in the long term but we we need to be building uh you know i'm going to be saying the word causal probably in a few different senses but we need to be building our our own better causal models of how these transformers work, which is to say, we need a clear overall picture about what kind of system the transformer is.

12:32When we're asking a question like, where is this fact stored? That question presupposes a lot about how the system actually works. Facts could be represented by individual vectors, or there could be individual buckets or bins inside the system or storages or registries that are responsible for storing those like you know whatever language the fact is written in like it's a vector it's responsible for storing that vector um and and it's like it's uniquely necessary and sufficient for for that function inside of the model where we're realizing just increasingly that's not really how transformers work so we actually need to figure out okay what are the questions we're supposed to be asking such that we can get like a sensible answer.

13:17The second broad area of research interest for you is model editing. The implication is that, you know, knowing how things are represented or knowing where a particular piece of knowledge sits in a model is not the only way to do it. And there are other ways to do it. And you're looking at some of those other ways. Tell us a little bit more about how you've approached that area. And other ways are really more traditional because traditionally how machine learning has worked is you treat your model, your neural network, at least in deep learning, as a black box, end-to-end, differentiable function.

13:53And the way that you train your neural network to do something you want it to do is you give it data and you write down an objective function and you optimize the objective function such that the final system is just better at being able to give you the outputs you want for the inputs that you're giving it. And so this is exactly a recipe you can follow with updating, you know, for instance, factual information and models. So if it happens that, okay, some of the most common examples are like sports players, you know, professional athletes get traded between teams all the time. So, you know, okay, well, now we need to update, you know, this player plays on this team and this player plays on this team and this language model.

14:31We don't, you know, this language model costs between$10 million and$100 million. We don't know. It's proprietary to train. but when we want to update this fact and the model well you know just traditionally the way you would do that is you say okay well here are a bunch of questions that involve that knowledge and you know we want to rather than the model saying the name of the old team we want the model to say the name of the new team so let's just train the model to do that and and this is this is exactly you use the off-the-shelf optimizers that have been a ton of work has gone into making these optimizers really well suited to the architectures we have and the kind of data distributions we have, you don't necessarily need to know anything about the internal mechanism that the model's using to get the answer correct.

15:15It's just black box optimization. And so when you say train and meaning, you know, training with the new questions and answers, are you referring to, you know, pre-training or ground up training, or are you referring to fine tuning? really more of a fine tuning step you'd be working with a pre-trained model and the part where you're updating it is doing just a little bit of last tweaking of the way it's just a little bit of fine tuning to to be able to update that knowledge inside the model another area where this desire to edit comes up is in the case of you know models that have factual data that you want to delete i'm thinking in particular of like you know more visual models that have information about an artist and that artist wants to be taken out of the model.

16:02Same broad problem, like you don't want to have to retrain the model. The approaches are ultimately similar. And I really care about this problem. And this is something we've done some work on. The idea of deleting information from models is increasingly important, particularly because of these huge pre-training costs. But it's just even if you're training things from scratch, we're still going to want to be able to particularly avoid copyrighted information, information that is simply not owned by the model trainer or model developer, maybe accidentally made it into the pre-training data set, these large-scale web scraping activities.

16:37Other kinds of sensitive or potentially dangerous information, we've seen some early studies from OpenAI and RAND suggest that models probably aren't that helpful at helping people build bioweapons uh but at the same time the models are learning a lot more about biology uh every year and and doing better and better at answering all kinds of questions there and and i think there's some legitimate concern um in the short to long term somewhere in the middle there uh about models just simply having like potentially dangerous uh information and you know the other half of that obviously is being able to explain it really well to people if they're jailbroken or if they're, you know, in this kind of chatbot form.

17:23So I really care, you know, this, I think this deletion problem is just pretty critical going forward and how people are doing that. Yeah. So, so again, you know, there's nothing really new under the sun and machine learning as far as I can tell. This, this machine unlearning area has been really active for a long time. People have cared about privacy in systems for a really long time. You know, going back to the 2000s, early 2000s, there's seminal work on you have a database and you train a model on the database, but someone simply requests that their records be deleted from the database. So now how do you change the model to respect their privacy in this setting?

18:00And a lot of the methods here, you know, it's really all some form of fine tuning. I would say that the thing, if you want one umbrella term that ties together the traditional fine-tuning methods and the new state-of-the-art model editing methods, it's constrained fine-tuning. You're doing some kind of traditional fine-tuning process, but you might have some additional constraint on which weights can I edit? What is the overall size of the edit I can make to the model? What are other constraints I need to respect in order to not change the model too much, make the edit more targeted or surgical?

18:38Yeah, but I'd say constrained fine-tuning is the umbrella. What's the Venn diagram between constrained fine-tuning and parameter-efficient fine-tuning? No, this is a good observation because really they're the same thing. They just often have different goals, I think. So the parameter-efficient fine-tuning world, for a while... This is LoRa, Q LoRa, Dora. And the main goal there is like doing fine tuning in a more performant way. More performant, less memory, faster, potentially more sample efficient. These are all current, basically observed benefits of a lot of these fine tuning methods like LoRa.

19:20I remember when some of these methods were initially being developed, again, going back to 2020, 2021. I was a little confused because often they might update only 1 % of a model's memory, sorry, model's parameters, but they use slightly more memory and are slightly slower. So I didn't really understand, but I think some of the more ML engineering people basically really figured this one out. And now they're faster and more memory efficient. I mean, one of the huge benefits here is being able to, If you want to train 100 different models, you don't need to store 100 different models. You only store 100 sets of LoRa parameters.

20:00And it's super lightweight and storage efficient, too. A lot of the big providers like AWS and Google call these LoRa parameters adapters. And so the idea is that individual customers can fine-tune on their own data without them needing to have tens of thousands, millions of copies of these large models. Yeah. But in fact, some of the modeling methods work almost exactly the same. You might be doing – so Laura does these low-rank updates to some subset of layers in a model. you know some of the method this modeling methods like mimet is an example where if you want to update multiple facts over time you can use this method called mimet mimet does low rank updates to specific matrices in the model and so it's almost exactly the same idea i think at the end of the day as as laura yeah and but mimet is an approach that's used more in the context of editing than in efficient tuning.

21:00Yeah, and this goes back to difference in basically stated goals, when in fact the underlying approaches are quite similar. You reference existing state-of-the-art model editing methods such as ROAM. What is ROAM? How does ROAM work? So it's actually quite similar to MIMA. There are these super constrained kind of updates you can make to the model. I mean, when you're talking like rank one for a matrix that is like, you know, it's actually like 16K by 4K. So the rank's at least 4K. So it's like quite a high dimensional kind of weight space you could be editing. And you're editing, you know, more or less one of those dimensions.

21:40This is an extremely surgical edit to the model. And, you know, basically it's improved a ton over all the off-the-shelf black box approaches. we've been talking about before in terms of its ability to update the facts of the model. And it strikes me that it's a little bit of, it's not that surprising that you could edit surgically, but it's maybe surprising that you can figure out where to apply the surgical edit. Yeah, I mean, this comes back to the localization part here. So because they're working at, you know, Rome was at a layer level. In fact, you can just try all the layers. So this, I think, turns out is the reason, you know, basically, this is the tuning step that led it to work, you know, as well as it did.

22:27It works well across different layers, but, you know, just getting that last little bump. The second step, that's when you're editing an individual fact. So you're saying, okay, the Space Needle is located in Seattle, or maybe we put the Space Needle in London. So you do a little bit of optimization in order to change where the Space Needle is located in the month. And so there's a little bit of a learning thing going on there where you're actually saying, okay, here's the matrix and its old output was this vector, but we want the new output to be this other vector, which basically gets the model to say London.

22:56And now we figure out how should we change the matrix to reflect that new information? So that's like the inference time kind of learning. There is, interestingly, this pre-computation step, which I think maybe some more experiments are showing exactly how valuable this pre-computation is. But the model actually, and what they do is they scan over some data from Wikipedia and they say, okay, what does the internal distribution of these latent states typically look like? So, you know, you could have all these different vectors coming into this matrix and you want to get a sense of what that distribution of vectors looks like.

23:34And that basically forms one term in their update equation. So that's kind of just like a mathematical necessity. They're borrowing, David Bowen's borrowing an older model from neuroscience, a computational neuroscience called linear associative memory. So it's kind of, it's actually a comp neuro inspired method. And linear associative memories rely on understanding the distribution of inputs to your memory matrix. And so in order to do that, they run the model of Wikipedia, compute what that distribution looks like. And then that distribution always that's stored in memory and that gets called during the inference time optimization.

24:14And so that having been said, what did this particular paper of yours go on to explore? Yeah, so on the practical side here, we are adapting these model editing methods, not for just really the predominant way these things have been used and tested, too. Let's say we know where the space needle is. Let's change it from Seattle to London. And well, it's a somewhat artificial thing to do. I mean, why would you want to do that? Like, why would you want to just like get the model to believe a bunch of false things? Believe it or not, that is the way academics were often benchmarking these methods.

24:47So it's not necessarily the most practical thing, but we definitely do want to be able to delete information from models, you know, for all the reasons we talked about before. So, you know, the very first thing we did was just change the problem and say, okay, let's think about deleting all these pieces of factual information from the model. Like, they're not exactly the kinds of dangerous information we care about, ultimately. In fact, very recently, the Center for AI Safety just released a data set of a bunch of, like, questions about, you know, biology and protein synthesis that they are encouraging researchers to work on in this area of machine unlearning.

25:24But prior to that, we were just deleting, you know, random trivia facts from the models. um and that's the first thing we did was you know kind of change the problem to this deletion problem but adapt existing model editing methods for this problem um and and i could say a little bit more about kind of the threat model too although we've already talked about the context of yeah sometimes models know things you don't want them to i'm imagining then you set up some evaluation criteria um you know that is almost maybe like you just want negative performance against a set of known questions for the model?

25:57Yeah, almost exactly like that because you just don't want the model to get the answer correct anymore that it previously knew. But so here's our adjustment to the evaluation because if someone's working with a chatbot and they're asking for, okay, someone's personal information or this particular or copyrighted data, I want you to give me this thing that normally I have to pay for or something or just all the other reasons we were thinking about. You know, are they just going to try once and then give up if they don't get the answer they want? Probably not. They'll probably try a bunch of times.