Opening AI's Black Box with Prof. David Bau, Koyena Pal, and Eric Todd of Northeastern University

In this episode, we dive deep into the inner workings of large language models with Professor David Bau and grad students Koyena Pal and Eric Todd from Northeastern University. Koyena shares insights from her Future Lens paper, which shows that even mid-sized language models think multiple tokens ahead. Eric discusses the fascinating concept of Function Vectors - complex patterns of activity spread across transformer layers that enable in-context learning. Professor Bau connects the dots between these projects and the lab's broader interpretability research agenda, identifying key abstractions that link low-level computations to higher-level model behaviors. SPONSORS: Oracle Cloud Infrastructure (OCI) is a single platform for your infrastructure, database, application development, and AI needs. OCI has four to eight times the bandwidth of other clouds; offers one consistent price, and nobody does data better than Oracle. If you want to do more and spend less, take a free test drive of OCI at https://oracle.com/cognitive ODF is where top founders get their start. Apply to join the next cohort and go from idea to conviction-fast. ODF has helped over 1000 companies like Traba, Levels and Finch get their start. Is it your turn? Go to http://beondeck.com/revolution to learn more. Omneky is an omnichannel creative generation platform that lets you launch hundreds of thousands of ad iterations that actually work customized across all platforms, with a click of a button. Omneky combines generative AI and real-time advertising data. Mention "Cog Rev" for 10% off www.omneky.com The Brave search API can be used to assemble a data set to train your AI models and help with retrieval augmentation at the time of inference. All while remaining affordable with developer first pricing, integrating the Brave search API into your workflow translates to more ethical data sourcing and more human representative data sets. Try the Brave search API for free for up to 2000 queries per month at https://bit.ly/BraveTCR Plumb is a no-code AI app builder designed for product teams who care about quality and speed. What is taking you weeks to hand-code today can be done confidently in hours. Check out https://bit.ly/PlumbTCR for early access. Head to Squad to access global engineering without the headache and at a fraction of the cost: head to choosesquad.com and mention “Turpentine” to skip the waitlist. TIMESTAMPS: (00:00) Intro (04:03) Reverse Engineering AI (10:10) Factual Knowledge Localization (16:34) Sponsors: Oracle | Omneky | On Deck (18:27) Future Lens Paper Intro (23:54) Choosing GPT-J (26:54) Vocabulary Prediction (30:36) Sponsors: Brave | Plumb | Squad (33:13) Fixed Prompt (35:57) Soft Prompt (41:32) Future Lens Results Analysis (51:50) Tooling & Open Source Code (56:07) Larger Models & Mamba Probes (1:04:10) Function Vectors Paper (1:09:39) Extracting & Patching Vectors (1:13:30) Encoding Task Understanding (1:15:36) Expert Models Implications (1:18:09) Conclusion

Speaker names marked ? are inferred from speaking time and the show’s known hosts — they are not voice-verified.

• 0:00Speaker 00

  Machine learning worked okay, but didn't really work in profound ways until the last 10 years or so. But now it's really working. It's really working remarkably well. And so we're facing a new type of software that we cannot use traditional computer science tools and traditional computer science methods for dealing with how to ensure that it's correct, that it does the stuff that we want.

• 0:22Speaker 01

  Logit Lens is essentially the tool where it projects these intermediate states into the decoding layer. So essentially, we can just see for that current token prediction, what is the model currently thinking about? We wanted to start with a code and code simpler solution, like having some sort of linearity with decoding future tokens. That would have been like a nice final solution, but we realized that no, it's a lot more complex than that, at least at the moment.

• 0:47Speaker 02

  We tested like 40 plus different tasks and it seems like they're all mediated by this small set of heads, which is cool that it seems like the the model has this sort of path that it communicates this task information. Even though in the prompts, we're never explicitly telling it what the relationship between the demonstration and the label are. It's able to figure that out and communicate it forward.

• 1:10Nathan Labenz?

  Hello and welcome to the Cognitive Revolution, where we interview visionary researchers, entrepreneurs and builders working on the frontier of artificial intelligence. Each week, we'll explore their revolutionary ideas and together we'll build a picture of how AI technology will transform work, life and society in the coming years. I'm Nathan Wenz, joined by my co-host, Erik Torenberg. Hello and welcome back to the Cognitive Revolution. Today, I'm excited to share a conversation with Professor David Bau of Northeastern University and grad students Kewina Paul, author of the paper Future Lens, and Eric Todd, who made a surprise appearance to discuss their recent paper describing function vectors. Professor Bau left Google to establish this group with the goal of cracking open the black boxes that are modern AI models

• 1:55Nathan Labenz?

  so as to understand the internal mechanisms that underlie their key capabilities and ultimately gain robust control of increasingly powerful systems. And indeed, starting originally with vision models and now working more on LLMs, the Bau Lab has done pioneering work in this area, delivering a remarkable series of mechanistic interpretability papers over the last couple of years. In this conversation, we go deep on a couple of the group's recent publications. First, Koyena walks us through her Future Lens work, which develops techniques that show that even mid-sized language models seem to be thinking multiple tokens ahead. Then Eric Todd joins to tell us about function vectors, a complicated pattern of activity spread across multiple attention heads in different layers of the transformer that

• 2:38Nathan Labenz?

  seem to enable in-context learning by encoding the nature of the task inputs and outputs. This work is fascinating for anyone who wants to develop their intuition for how today's models actually function under the hood. And they're a notable step on the path to overall AI interpretability, identifying new abstractions that translate low-level patterns of computation to higher-order behaviors. I was a fan of this work coming into the conversation, but I left as a big fan of Professor Bau as a research advisor as well. Throughout the conversation, he does a great job of connecting the dots between his students' individual projects, motivating the lab's broader interpretability agenda, and emphasizing key open questions in the field.

• 3:21Nathan Labenz?

  As always, if you find this work valuable, please do share the episode with your friends. I'd especially suggest this one for anyone who's interested in mechanistic interpretability and considering a PhD in machine learning. And please, don't hesitate to share your feedback or suggestions via our website, cognitiverevolution.ai, or by messaging me on your favorite social network. Now, please enjoy this illuminating discussion about the inner workings of large language models with mechanistic interpretability researchers Koyena Paul, Eric Todd, and Professor David Bau. Koyena Paul and David Bau from Northeastern University, authors of Future Lens. Welcome to the Cognitive Revolution.

• 4:00Speaker 01

  Thank you.

• 4:02Nathan Labenz?

  Really happy to be here, Nathan.

• 4:03Nathan Labenz?

  I'm excited for this conversation. So Professor Bau leads the group and Koyena is part of the research group. And together with your colleagues, you guys have really developed a remarkable interpretability and sort of editability agenda with respect to large language models over the last couple of years. I have covered a couple of your papers previously, although not in the depth that we will get to very shortly with the Future Lens paper. But it's a really cool agenda that is shedding a lot of light on

• 4:30Nathan Labenz?

  How these new fangled AIs are actually working, what's going on in there, kind of cracking open the black box problem. And certainly, regular listeners will know that that's something that I really appreciate and value. For starters, you guys want to just give us a little background on the overall agenda of the group and maybe highlight a couple of your favorite papers that you've published previously?

• 4:50Speaker 00

  It's a pretty simple agenda. So the thing that motivates our group is the fact that machine learning is really a new era for computer scientists. Machine learning's been around for a while, but up until the last decade or so, it's been this dream of having a way of creating self-programming computers that learn from data, that people don't really have to program every line of code. Machine learning worked okay, but didn't really work in profound ways until the last 10 years or so. But now it's really working. It's really working remarkably well. And so we're facing a new type of software that we cannot use traditional computer science tools and traditional computer science methods for dealing with how to ensure that it's correct, that it does the stuff that we want. All

• 5:35Speaker 00

  of our tools are very statistical, they're very black box. And so what the theme of the lab is, is to ask, how do we confront this new type of software? Are there things that we can do to recover our ability to debug and understand the software that we create? Are there ways that we can take more responsibility for the behavior that emerges from this software? Can we edit the code like the way that you would change a line of code in a traditional program? And so this is actually pretty hard. It would be like asking the question, can you totally control the mechanisms of a tomato or some biological thing that emerged out of a process of evolution? And machine learning is similar.

• 6:19Speaker 00

  That way, it emerges out of a training process. And so just like biologists have to reverse engineer how a tomato works, if they want to really understand what they can do to make it better, we're also trying to develop the science of how to understand the internals of these systems.

• 6:35Nathan Labenz?

  Now, is that an agenda that you have changed course on relatively recently? I assume it has to be, inasmuch as the understanding of advanced AI systems in this way is predicated on their existence. And I would imagine like, at least probably prior to GPT-2, there was almost nothing powerful enough to do this kind of work?

• 6:55Speaker 00

  Oh, no, I've been at it for a while. So my personal history is I worked at Google for many years and I actually left Google to pursue research when it started to become clear around 2015 that the field was going to change. And the speed at which it's changed has surprised me. It's surprised almost everybody in the field, but it was clear that there was a sea change happening. The moment for me was also the same moment for a lot of people when AlexNet came out in 2012 and showed that this really simple, really classical machine learning procedure, something from an idea from the 1950s, right? And an idea from the 1980s, training these neural networks, which really the community

• 7:36Speaker 00

  at large thought that such a simple procedure probably wouldn't be the thing that solves machine learning in the long run. People thought it would be something trickier or something more complex or something like that. But this idea that all you have to do is scale up neural networks and it works really well was a real eye opener. So the things that I studied at first, I studied generative vision models a lot, like GANs, which were really surprising when they came out. And I thought, I'll study these GANs for a long time. Of course, after GANs, we have diffusion models. And of course, we've got these amazing language models and other things going on. So it's been this incredible moving target.

• 8:13Nathan Labenz?

  It does seem to be a pattern that the earliest to the game started with more vision models. I know that's true for Chris Olaf, for example, as well. So yeah, fascinating kind of shared lineage there. You want to touch briefly on the Rome and Memet papers? This is one that I've covered a little bit in a previous episode. And the way I remember those papers is as a demonstration of the locality of factual knowledge in a language model, and also the ability to then edit that language model in ways that preserve a lot of the key properties that you would want to preserve if you were editing a knowledge graph, where, for example, if you were to, I think that the example in the paper is

• 8:54Nathan Labenz?

  Michael Jordan played basketball, of course, is the real sentence. But if we wanted to edit that world knowledge to reflect an alternative reality where Michael Jordan played Baseball, then we would want that to be not just for that specific sentence, but robust to different ways you might ask the question, different ways you might set up an auto completion and in a process of kind of systematically working through the layers to figure out where you're corrupting bits as you go through the layer and figuring out where it starts to fail and then using that as the place to edit and then achieving that kind of robust editing. It's been a couple of months since I went deep on that one. Tell me how I did.

• 9:32Speaker 00

  No, no, that's right. And I would back up a little bit, maybe half a step further, which is the fundamental question that we're asking, and it's actually related to Future Lands 2, is, is there like a physical reality to the idea of a concept or an idea or a unit of knowledge, right? We have these things in our head. We think you can go to somebody and you can say, that's something that you know, like you know where the Eiffel Tower is, or, oh, you've never heard of the Eiffel Tower? You have no idea where that is? No, it's something that you don't know. So we have this idea that people can know things, or they can not know things, or they can have an idea or a concept, or they can not have one. It just seems so crisp intuitively in our own heads. So the question, is there a physical reality to this? When if

• 10:15Speaker 00

  an artificial neural network seems to know something or doesn't know something, is there a physical reality to that knowledge? Or is it just diffuse and spread out in all the neurons and all the computations through the whole network? And I think that there's definitely a school of people who would generally believe that most things are just unimaginably diffuse, that everything's just spread out. There's no reason that there would be any locality. But there's a bunch of neuroscientists in the biological world and some computer scientists in the artificial neural network world who have taken a look over the years and noticed that a lot of times you see locality. I saw that kind of thing in GANs, but our question was, do you see this kind of organization in really large models? And do you see any

• 11:00Speaker 00

  that corresponds to the facts? And so what Rome was about is a bunch of experimental methods for narrowing this down and finding out where facts are localized in a large language model. And it sounds like you had another Pytest on it. So maybe you can direct people to listen. To that one, but yeah, that's the general setting. And of course, there's a bunch of clever methods that you can use to really pin that down.

• 11:22Nathan Labenz?

  What is the state of that work now? I know that the successor to Rome was the Memet paper, right? And that one showed that you could edit up to like 10,000 facts at a time.

• 11:33Speaker 00

  Sure. And then there's another paper that we've written. It's actually also coming out this coming year at ICLR. It's called the LRE paper, the linear relational embedding paper. And that paper looks at a different slice of the problem. That paper really asks, if you have a relationship, like something is located in a city or somebody plays a certain instrument or somebody plays a sport or somebody holds a certain position or something like that. There's this general relationship. What's the representation of that general relationship? The Rome paper was asking, what's the representation of like an individual association? But we really wanted to break that down. We can see that there's an individual association.

• 12:15Speaker 00

  We can see that the representation of the subject and the object seem pretty clear. But we wanted to understand what's the representation of the relation between the subject and the object? And the LRA paper explores that and finds it in a bunch of the cases. The relation can actually be understood as a linear mapping. But it's actually sort of an incomplete story because we found that for a lot of the cases, there's something more complex going on. There's something that looks nonlinear. And so it opens the door to future research, new tech, to ask what's going on in the more complicated cases? These relations are sometimes really non-trivial. And I think that there's going to be some interesting, exciting work to do to figure out like,

• 12:57Speaker 00

  why is it that the models are modeling some relations in a really complex way? What are they getting out of all that complexity? And is it enabling new types of reasoning that we haven't characterized yet?

• 13:09Nathan Labenz?

  Cool. That's awesome. I look forward to that one in more detail in the future as well. Koen, now let's get into the Futurelands paper where you're the lead author. This one really called out to me pretty much immediately because I think you were speaking to it a moment ago as well, where you're saying there's a school. Of thought out there that's like, this is all stochastic periphery and there's no real understanding and it's just next token prediction. So I'm always interested in things that explore like when mere next token prediction can in fact be a bit more than meets the eye. And this paper is a really good example of that. So we want to start off maybe by just talking about the the specific motivation for this particular investigation within the broader agenda of the group.

• 13:54Speaker 01

  Yeah, so actually, the way this problem started was like a little different from how it was before. Initially, I was interested in looking at memorized content in large language models. Cause when you use chat GPT or any other models, some things that I've noticed like license codes, for example, if you just give couple tokens in, you kind of know the next 11, 20 tokens that tend to be exactly the same from its training data. So initially, I was actually first exploring that. And while I was exploring that, I like tangentially went to this point where like, okay, well, I can see that it's saying these tokens pretty much like word to word, but like, how about for smaller phrases like New York City and stuff like that, which

• 14:37Speaker 01

  for us, we consider that as one whole entity. Yeah, that's how it started off in terms of just looking into seeing whether one, a single hidden state. Can we actually decode such information from that? And yeah, we were exploring ways that you have seen in the paper. In terms of to what extent can we decode that future information about a particular entity or like just anything in general at this point.

• 15:00Nathan Labenz?

  So the question, as it's, I'm reading directly from the paper, in this work, we ask, to what extent can we extract information about future that is beyond the currently under consideration token from a single hidden token representation? So let me again, just pitch you how I understand this work and you can correct any misconceptions that I have. I think the vocabulary, especially for people that are not in the research, can sometimes be a barrier to entry. So if I understand correctly, what you're doing here is looking at activations within the forward pass of a large language model. As the numbers are getting crunched, I sometimes like to use the most basic terms I can. The numbers are getting crunched from beginning to end.

• 15:44Nathan Labenz?

  And between the layers where the actual intensive computation occurs, there are these kind of middle states, which are known as activations. You are taking an array of numbers, basically, that sits there at different places. And we can get into more details of the different places that you can look. You can kind of look at every given layer, right? Through the transformer to see a different activation. But the question you're always asking is, given this array of numbers, these activations, can I predict not just what the current token is going to be as I'm working through that forward pass, but can I also see what the future tokens are going to be just from the information that's in that one particular vector of numbers? Is that right?

• 16:25Speaker 01

  That sounds about right. Yeah, like we just termed those activations hidden states because those are like hidden and they are part of the intermediate computation and stuff.

• 16:34Speaker 05

  Hey, we'll continue our interview in a moment after a word from our sponsors.

• 16:38Speaker 01

  So Logit Lens is essentially the tool where it projects these intermediate states into the decoding layer. So essentially, we can just see for that current token prediction, what is the model currently thinking about? Like, based on the distribution and when you decode it to the highest probable word, what is the model currently thinking?

• 16:58Speaker 00

  I recommend anybody who's listening to the podcast who hasn't seen Logit Lens before. This is not our work. This is somebody else's work. So there's a blogger, a nostalgia brist, who created this idea of Logit Lens. You should Google it and take a look at what Logit Lens looks like. It's these beautiful grids of words that are overlaid on all the hidden states of the transformer, where every hidden state, every one of these vectors that's inside a transformer, all of these intermediate numbers are all translated into a word. And then at the bottom layers, at the very earliest phase of the transformer, you can see what words the transformer is basically internally thinking. And then as you go through more and more layers, the words evolve until you get to the end.

• 17:41Speaker 00

  And you can see the predictions of words get more sophisticated as you get to the end. It's always guessing. What the next word should be that it should say. And those guesses get better and better the deeper you go into the transformer. It's really interesting to look at the logic lens.

• 17:57Speaker 01

  Yeah, the biggest advantage is you only see what it's currently thinking, which may not be representative of what it might be actually thinking about a couple of tokens ahead. It's not as informative in that sense. And so we are trying to like make it more informative and seeing like, okay, it's currently thinking of the word new. What word is it thinking after new? Like, is it thinking of New York, New Jersey? These are like some examples of things that could follow new, but that's kind of like not possible to see with logic lens for that particular hidden state, at least.

• 18:27Nathan Labenz?

  Yeah, so there's two directions that we can consider. Tell me if this intuition jives with your understanding. In the course of the forward pass, basically, this is what the logic lens looks at. It's like, as we're going layer by layer, we're sort of gradually getting more sophisticated in the way that the model seems to be understanding the input. And it's gradually converging to the prediction that it's actually going to make. So you can see at the beginning, it like has a low level of confidence and maybe wrong. If you just took that hidden state and went directly to decode. And that would be basically like, okay, let's say we do two layers of computation, then we skip all the rest and then we decode whatever came out of the two layers instead of all the layers. What would we get, right? What we get is not very accurate at two layers.

• 19:10Nathan Labenz?

  It gets more and more accurate as you go through the layers as it finally reaches its output, which makes sense, right? Because that's presumably why we need all the layers is to be working through the information and gradually getting somewhere. And then there's also the token direction, right? We've got like a sequence of tokens and information can only flow forward because that's the nature of the look back attention mechanism. So it can only flow forward from tokens to future tokens. But that means, or at least it stands to reason, right? And this is what you will then set out to prove that as this information is built up from prior tokens to the current token, it probably means that there's more there than just what would be needed to predict the immediate token.

• 19:54Nathan Labenz?

  And there's been like lots of interesting intuitions around this with even just. Articles in English, right? Like if you set up a situation where it's very clear that the actual noun is going to begin with a vowel versus a consonant, then you can see that it predicts the article correctly. And then you're like, huh, that's pretty suggestive that maybe there's more information there. Otherwise, how would it be choosing the article? It seems to be anticipating what's coming beyond this current token as well. So I think of the information processing proceeding like vertically and the token information flow proceeding horizontally. And now we're in this moment where it's okay at each layer through the forward pass through the vertical dimension. We can check to see like, would this be enough already to predict the current token?

• 20:37Nathan Labenz?

  And again, that's the logit lens kind of based on what you're building. Now we get to, okay, can we also detect that future token information is there, right?

• 20:47Speaker 01

  Yep.

• 20:47Nathan Labenz?

  Okay, cool. So tell me how you did it. First of all, I noticed that this was done on GPTJ. That was kind of surprising, but maybe I'm just not sure what I don't know about GPTJ. It seems like these days this is like Lama would be the default for this sort of thing. So how did you choose the model?

• 21:04Speaker 01

  Again, since this was like a spin off from the earlier project that I was talking about. So I was using GPTJ and I just continued with it, but it was definitely like one of those models where I could run it on like a single.

• 21:17Nathan Labenz?

  Do we have a future lens on Lama now?

• 21:19Speaker 01

  I think people are building it. If it has the same vector size, like four, zero, nine, six, then you can technically do the whole transfer. But yes, generally this method should be workable for other models as well.

• 21:31Nathan Labenz?

  Yeah, let's go back to the other models in a little bit because I have some questions around like how you think the different models might behave differently under this sort of examination. But it'll help to establish exactly what it is we're doing. And you said the activation vector, this was actually one of my questions. The vector from which you are predicting, that's a four thousand ninety six size vector?

• 21:54Speaker 01

  And like the number of tokens, but then that would be one for us. So, yeah.

• 21:57Nathan Labenz?

  Okay, cool. So then tell me about the data set. So this is the first, as is so often the case, first challenge in doing something like this is you've got to collect a data set. This one is kind of a subtle one to describe.

• 22:09Speaker 01

  Yeah. So since I knew that GPTJ was trained on pile, I was like, okay, why don't I take the pile data set and sample? Couple of texts from it and just get the hidden state from each of the sentence in there. And yeah, I'll run it. So I collected like hundred thousand of them. One of my initial condition was if it is GPTJ itself, if it is reading this one token, that it's predicting the next token correctly in that sense. And so that's one of the few sample condition that I had. But otherwise, yeah, once I got like hundred thousand of them, I stopped and then I was experimenting on that.

• 22:45Speaker 00

  So it's a hundred thousand samples of texts with what what's the filter that you used in that?

• 22:50Speaker 01

  Oh, yeah, the filter that I used was that given the last text of the prompt, that the next one it predicts correctly, the model predicts correctly. Yeah.

• 23:00Nathan Labenz?

  Okay, so there's a hundred thousand prompts identified by the condition that the model must get the next token correct per the actual source material. So what are you actually collecting for each of the hundred thousand? It is the activation, right? Or the or the hidden state for each of the layers throughout the entire transformer?

• 23:20Speaker 01

  Yeah, that's what eventually happens. But in terms of just storing purposes, I had just a CSV file of these hundred thousand sentences. But then later, during the training process, I would gather the last hidden state of it at that particular layer for whatever like probing or a method that I'm trying.

• 23:37Nathan Labenz?

  Because all of the analysis downstream of this is going to be at a layer specific level, right?

• 23:43Speaker 01

  Yes, ultimately, the actual input itself would be the four zero nine six vector. But yeah, when I'm collecting all of them, I'm collecting them based on the strings and stuff. So, yeah.

• 23:54Nathan Labenz?

  I always like to get very clear in the inputs and outputs. So now we've got to the point where, okay, we got the hundred hundred thousand texts. And our goal is to say, for a given layer, can we figure out a way to take that four thousand ninety six length vector and use that as the input and then predict not just the current token, but the next several tokens from that single input vector?

• 24:19Speaker 01

  Yeah, almost there. So definitely layer, but also the last token. So if my prompt was 10 tokens, if it was layer wise, it'll be 10 by 4096. But since again, it's just one token and a single hidden state, it'll be the last token of the input prompt. And that would be so that's one by 4096. Yeah.

• 24:37Nathan Labenz?

  Okay, cool. And then just as validation, there's another thousand that you'll do the the testing on once you develop the methods. Okay, cool. So we've got five distinct methods where the first one is the null hypothesis, which is just the bigram, which is, okay, given the total frequency of word pairs in the entire pile dataset, to what degree can I just look at the current word and predict the next word solely from that? And if you're guessing at home what that number is, like, give me a second, it is 20%. Apparently 20% of tokens in the pile, you can guess just from what the knowing the current word gets you basically to 20% accuracy. That was kind of surprising to me. I wouldn't have guessed it would be that high. There you have it. 20%.

• 25:21Speaker 00

  Yeah, this language modeling is not so hard after all. We just get 20% from counting. And that's it. The San Graham thing was the state of the art method until not too many years ago.

• 25:31Nathan Labenz?

  Yeah, but it's not state of the art method anymore. So let's talk about the four methods that you developed. And first one is direct vocabulary prediction. So here, you take that same for 96 vector and just specifically try to jump directly to the logits, right? The the final percentage, it's not exactly percentage, but it's within a very short transformation from percentages of the actual tokens to be predicted for the current and then the the next few. What is the nature of the thing that you're learning? This is because it's a learned method, right? You're going to optimize and essentially learn a way to make this prediction across the 100,000 samples. What is the nature, though, of the

• 26:16Nathan Labenz?

  vocabulary prediction thing that is being trained?

• 26:19Speaker 01

  Yeah, so like the goal with this direct vocabulary prediction is essentially can we like literally decode the token from that hidden state? I imagine the linear models would be the most basic model for probing in terms of just being able to grab that token. And so, yeah, that was essentially the goal. We wanted to see whether there was some sort of linearity in like being able to just decode what a future token is. Like, is that stored linearly in the hidden state?

• 26:47Nathan Labenz?

  So that is, to be very concrete, a one-layer transformation.

• 26:52Speaker 01

  No, one layer, one token. So in some sense, it's like context-free, but yeah, it's one hidden state to one other hidden state. And each of them is like one by four zero nine six.

• 27:01Nathan Labenz?

  The thing that is being trained, that is doing the transformation from the hidden state to the prediction, is just a simple linear transformation.

• 27:10Speaker 01

  And sorry, just to correct myself earlier, since we were talking about the vocabulary prediction, it's actually one by four zero nine six to five zero four zero zero. So essentially, the way this vocabulary prediction is that the GPTJ, if we were to tokenize the words, the size, the vocabulary space of a GPTJ is about fifty thousand tokens. And we are essentially trying to get that vector of like fifty thousand tokens and choose whichever is the highest probability as the token that would be predicted. So that's for the linear direct vocabulary prediction. And so, yeah, that's why for this particular method, we grab in the hidden state, which is again four zero nine six vector.

• 27:50Speaker 01

  And then that's to five zero, like that fifty thousand size of like one token.

• 27:56Speaker 00

  You're literally training a decoder, your own decoder head on this problem, right?

• 28:01Speaker 01

  Yeah, essentially, yeah.

• 28:02Speaker 00

  Training a little linear classifier.

• 28:04Nathan Labenz?

  So the number of parameters in that linear classifier would be four oh nine six times the fifty thousand vocab size?

• 28:13Speaker 01

  Yep.

• 28:13Nathan Labenz?

  So essentially, each number in the hidden state makes a contribution to each output token. And the intensity of its contribution is determined by the parameters which are learned in the training of this linear.

• 28:27Speaker 01

  That sounds right, yeah.

• 28:28Nathan Labenz?

  Okay, cool. So that's direct vocabulary prediction. So we got a hundred thousand of these four ninety-six vectors. We know what the actual next tokens are. And we're training this mapping of four oh nine six by the vocab size to do that mapping. Hey, we'll continue our interview in a moment after a word from our sponsors.

• 28:52Speaker 05

  Hey, y'all. I'm hearing more and more that founders want to get profitable and do more with less, especially with engineering. Listen, I love your 30-year-old ex-Fang senior software engineer as much as the next guy, but honestly, I can't afford them anymore. Founders everywhere are trying to turn to global talent, but boy, is it a hassle to do at scale, from sourcing to interviewing to on-the-ground operations and management. That's why I teamed up with Shawn Lanahan, who's been building engineering teams in Vietnam at a very high level for over five years to help you access global engineering without the headache. Squad, Shawn's new company, takes care of sourcing, legal compliance, and local HR for global talent, so you don't have to. With teams across Asia and South America, we can cover you no matter which time zone you operate in. Their engineers follow your process and

• 29:36Speaker 05

  use your tools. They work with React, Next.js, or your favorite front-end frameworks. And on the backend, they're experts at Node, Python, Java, and anything under the sun. Full disclosure, it's going to cost more than the random person you found on Upwork that's doing two hours of work per week but billing you for 40. But you'll get premium quality at a fraction of the typical cost. Our engineers are vetted top 1% talent and actually working hard for you every day. Increase your velocity without amping up burn. Head to choose squad.com and mention Turpentine to skip the waitlist.

• 30:13Nathan Labenz?

  Then there's kind of two variations to this, right? The first one is the direct vocabulary prediction, the second one is the linear model approximation. And if I understand correctly, the difference is whether you're jumping straight to the output tokens is what you're predicting, or keeping the existing decoder and instead trying to predict something that then the decoder would properly decode.

• 30:35Speaker 01

  Exactly, yeah. One is to the vocabulary distribution, and then the other one is to the hidden state, which the decoder layer would decode it to that ultimate token. Yep.

• 30:44Nathan Labenz?

  And that's pretty interesting. I didn't have too much of an intuition for like, why one might be thought to work better than the other. And then it seems like it basically turned out that they... Roughly work equally well, like almost uncannily similarly.

• 30:56Speaker 01

  Yeah, actually, for me, that was a little surprising because I would have assumed that because for the vocabulary distribution, like it's also trying to learn how to decode, like it's creating its own decoder thing. I thought it will take probably more number of iterations or basically a lot more data for it to come to that stage, but it worked similar to how if you just predict the hidden state and have your decoder layer handle that. Yeah, it works similarly. So that was pretty cool to see.

• 31:24Nathan Labenz?

  Yeah. Is there something that we can understand about the broader system based on the fact that those two approaches work the same? Like, it seems you basically have two linear classifiers and you can't like two is not better than one is kind of an interesting way to see that, right? I guess one way to interpret that would be that you really are getting toward the max of the information that is there. I guess another thing you might try would be like a bigger, more, you could train a full transformer to do this sort of thing. And maybe you did, but I guess if two linear classifiers aren't better than one, then maybe you'd think we're already maxing out. Is that how you would interpret that?

• 32:01Speaker 01

  Yeah, that's how it interpreted. We wanted to start with a quote unquote simpler solution, like having some sort of linearity with decoding future tokens. That would have been like a nice final solution, but we realized that no, it's a lot more complex than that, at least at the moment. And yeah, that's what we tried. We tried our best to max out the linearity aspect of our experiments.

• 32:21Nathan Labenz?

  Okay, cool. So that's one class of thing, training these linear classifiers. And then the other class of thing is basically an approach where you take the activation out of the context in which it was generated and port it over to and kind of stitch it into a broader forward pass context, where that other forward pass context is engineered to be as neutral as possible, or I guess ultimately engineered to allow that activation to shine through and get its content through to the predicted tokens as effectively as possible. So the first way I thought was like super intuitive to do that. You just set up the. The prompt with hello, could you please tell me more about?

• 33:05Nathan Labenz?

  And then port in the activation from the other context where however much information has been aggregated at that final token is now being ported in and the theory there is like, well, because hello, could you please tell me more about is so generic? Then like in theory, whatever you put there should set it up for a kind of nice, clean continuation. Unfortunately, that one didn't seem to work so well. So I was surprised by that. What do you think is up with that?

• 33:29Speaker 01

  It's because perhaps it was too generic in that sense. While yes, we were targeting generic prompts so that it has the possibility for whatever the intervened hidden state is to like ultimately have that token shine through. But it was probably still too general. The prompt wasn't really optimized for it to be like, hey, I wanted the future tokens out from this particular state. And so that's why I'd feel like it didn't work as well.

• 33:56Speaker 00

  We tried a few, right? Didn't find any that really worked that well.

• 33:59Speaker 01

  Yeah, before I gave a couple sentences, I literally started with start of text or just a dot of opening bracket. Yeah, I tried a bunch of them, but they didn't seem to like work as well either. Yeah. But it was working a bit better than linear models, of course, but still not really as great.

• 34:16Nathan Labenz?

  Yeah, that's interesting. Okay. So now what comes next though, is perhaps even more surprising in view of the fact that you tried multiple actual natural language prompts and none of them worked. The next approach is saying, okay, well, you know, where we're going, we don't need language. It's the soft prompt technique, right? Where this time you're learning what the embeddings should be from a sort of hypothetical abstracted prompt, like what would the ideal prompt be if we didn't have to represent it in tokens? And we could just go straight to the numbers and engineer it at that level so that we could get the tokens out the other end that we want. So in this case, the optimization is tweaking those embeddings across all of the activation to future token prediction pairs.

• 35:01Nathan Labenz?

  And by optimizing those embeddings, now you can really start to get somewhere.

• 35:07Speaker 01

  Yeah, that sounds right. That's exactly what we did.

• 35:09Nathan Labenz?

  That's surprising. Well, it kind of jumps out at me, obviously. Now, looking to the results, that method works by far the best. And how do you understand why that works so much better? Is there any ability to say what the ultimately learned soft prompt means? Can we translate that back into something that I can understand?

• 35:26Speaker 01

  We were actually attempting to do that. We were like, okay, we found a soft prompt that works great. Let's try to like, get a discrete one, but it was pretty gibberish. So it was like, okay, we'll keep it in like, the vector space for now.

• 35:38Nathan Labenz?

  How gibberish was it?

• 35:39Speaker 01

  It had like, a couple of ads somewhere, and it had some like, they looked like words, but then they didn't make sense kind of words. Those were what was there. Yeah, it was pretty gibberish. So we just went on with the vector space soft prompt. But yeah, in terms of the results, it was nice to see that there was a pretty significant increase. And the reason I imagine is because this prompt is trained so that like, whatever intervened state we have, it enhances whatever the future context is encoded in there. And so because in the manual versions, we were thinking of what the prompts could be, like what we think would make sense. But over here, this is what made sense to the computer. Yeah.

• 36:19Speaker 00

  It's amazing that basically, there's this little capability in the transformer to solve this problem, right? To tell you what the future tokens are. It has some little machine in it that knows how to decode it. And Koyena, with her soft prompt training, basically came up with a prompt that is like a key. They're like, it's sort of this capability to do it. At least that's the way I see it. But the key didn't seem to correspond to any real piece of text, other than this weird gibberish that she saw, right?

• 36:47Speaker 01

  Yeah.

• 36:48Speaker 00

  It wouldn't have been something that we could have predicted.

• 36:50Nathan Labenz?

  It reminds me very much of the Universal Jailbreak paper, which we did an episode on. And there was sort of a similar, now, they were operating in token space, but still often found these Universal Jailbreak prompts to be gibberish sequences. And sometimes they did have some that were like, more human readable. Now, this was something that you guys did a lot, right? Like, all of these techniques are happening at each layer. So for each of these techniques, you had to train, like, for the soft prompts, there's what? How many layers in GPTJ? Is it 28? So you had to train 28 different soft prompts, one for each layer at which you're going to patch in the activation.

• 37:30Speaker 01

  So that was what the method was initially, but like the actual future lens, the tool, we just use one soft from, we use the best model, best layer that worked. We did that for every state when we actually built the tool. So in the tool, it's like one, but then we were using this method to find which was the best soft from to use, if you may. Yeah.

• 37:49Nathan Labenz?

  Gotcha. Okay. So yeah, let's talk a little bit more in detail about the results. I think these are suggestive of many things. Certainly it's a good graph to just sit and ponder for a little while. Again, folks can refer to the paper to see this is figure four from the paper. Basically, what you've got here is four different graphs. Each one corresponds to the token that is being predicted. So N equals zero, the first graph is like predicting the current token. Then it's N equals one, the next token, two, the one after that, three, the one after that. And then for each of these predictions, you're able to look at, okay, put the layers across the X axis and the success rate on the Y axis. And you can see that as you proceed through the layers, the prediction success changes.

• 38:33Nathan Labenz?

  For the first token, the current token, basically you recap the logit lens results, right? Where you can see that as you move through the layers, the predictions get more and more accurate by the time you're at the end. By definition-ish or close to definition-ish at the end, are you getting the answer right? Close to a hundred percent of the time.

• 38:49Speaker 01

  Yeah, that was our sanity check for N equals zero. We were like, okay, is this method actually working? Let's make sure it's actually decoding the current tokens. Yeah, that was more of a sanity check.

• 38:58Nathan Labenz?

  So then several things jump out at me about the next three. One is that, again, we've already said that the two linear projection methods basically worked almost exactly the same. Like you can see that those lines for tokens one, two and three are in almost lockstep for through every layer. Those did not work that well. And in fact, and this is kind of surprising too, on the first token, they beat the 20% baseline a bit. But then on the second and third tokens, they're actually worse than the baseline. I guess maybe the baseline is out of domain there at that point anyway. But I wasn't expecting to see things go lower than the baseline, really.

• 39:33Speaker 01

  Yeah, I assume so. We created like the bigram baseline. And so for N equals two, N equals three would be like trigram and. Like the four gram models, they were too huge for us to compute at that time. And so, yeah, that's why we didn't include the 20% in there. But yes, it's more to do with the fact that it was trying to predict two tokens ahead, three tokens ahead.

• 39:53Nathan Labenz?

  So the baseline is not so relevant past just the immediate next token. So you probably don't have a paper with those methods, right? Because they didn't work that great. The next one is the fixed prompt. This one actually does go below the baseline on the first one. So again, I'm like, man, that seems weird. Why would it be lower than the straight baseline? And even to make that more confusing, that method seems to get better for the subsequent tokens. Like it's worst at predicting n equals one and it gets better at predicting n equals two and n equals three. And that really surprised me.

• 40:28Speaker 01

  So to provide a bit more context, so when we were checking n equals two, n equals three, we assume that we give like the correct token for the previous ones. So let's say we are talking about New York City, for example. Let's say n equals zero is predicting New York, one is York, two is City. And so let's say, because there are so many possibilities after New York, that it doesn't really predict York. But if we give New York, if not, if we give New, if we give York, can it actually then predict City? And so that's the idea behind checking for two and three. Because generally, if one is wrong, then it's pretty sure that two and three will be wrong if you check whole phrase wise. So that's why, yeah, we gave one extra token context in that sense.

• 41:10Speaker 01

  But with that, it's essentially the same thing with one additional token, but nothing else, like whatever the current token was.

• 41:17Nathan Labenz?

  Can you help me understand that a little bit better? When I am imagining porting an activation over two different contexts? How should I understand that, like mechanistically, interacting with the idea that there's another kind of subsequent token included?

• 41:33Speaker 01

  Yeah, so when you do transplant that hidden state to the fixed context, yes, it's actually completely from a different step at this point. Let's say we did this transplantation at like layer fourteen. So from layer fourteen onwards, it would have the context that is present in that hidden state. But if sometimes that's not enough, as we saw in the fixed prompt result, like N equals one. But let's say if we gave that particular hidden state plus what the next word could have been. So like York, in that sense, maybe the reason why it's performing better, it's okay. Cause it now has like a direction of like where it can predict next. But even then, it's just that hidden, single hidden state and that particular one token. So.

• 42:13Nathan Labenz?

  Interesting. So with that method, it's getting more information. Does that also help me understand? Cause another thing that I was quite surprised by is that the performance of that fixed prompt method seems to be pretty clearly declining as you go through the layers. And whereas the first two, the linear transformation ones are like flattish, climbing a little bit. The fixed prompt one is declining. And I was like, hmm, that seems like a very odd reversal from the main logic lens. Do you have an intuition for that? And does that have something to do with the extra tokens that it's given?

• 42:49Speaker 01

  So yes, in some contexts, but usually the way I imagine, just as the first like last layer is that from beginning, it's trying to build an understanding of what it's gonna say. And then in the middle, it has these set of ideas that it wants to say. And towards the end, it probably wants to say immediately what it thinks it should be next. And so in that sense, I imagine there could be some information lost along those layers when it's ultimately trying to predict literally the next token in that sense. Yeah.

• 43:16Nathan Labenz?

  Yeah. And what you just described is definitely my intuition as well. The sort of gradual working up of inputs into higher order concepts through at least half. And my general sense is that in many models even goes deeper than that. Cause rough picture in my head is more through like eighty percent of layers. And then, yeah, in the final layers, it's like now we're condensing again. Like now it's time to actually cash this out, all this sort of high order or high concept work that happened in. Middle, now we got to cash that out to a concrete prediction. And so you see a collapsing toward the end. And that is very much what we see in the soft prompt one. So a couple of notable effects about the soft prompt results.

• 43:55Nathan Labenz?

  First of all, they're just a lot better than everything else. And second, they do more follow this pattern that we both each just described of seemingly more and more information buildup that's more useful, rising line from the beginning up until the middle layers and then declining. And just the sort of peak here is in the middle. And I've seen later peaks in other models. My guess is that's probably just a function of model size. Where you have to have that end time and that limits how deep into the model you can continue to build. Does that seem right to you?

• 44:25Speaker 01

  Yeah, that definitely seems right to me. And the whole idea of like the middle layer importance, it is something that I've seen in other intro papers as well. And yeah, there's definitely something along in those middle layers, middle late layers as well.

• 44:37Speaker 00

  It is interesting how it declines, right? It definitely suggests that at the late layers, it's erasing some information that is using that space for something else. And this sort of explains this decline over time. So I feel like the whole phenomenon of a transformer erasing its own information is an interesting one. And we see hints of it here. It's another cool thing to study.

• 44:58Nathan Labenz?

  Yeah. What else do you take away from this set of results that I haven't landed on myself?

• 45:04Speaker 00

  One of the things here is that the information is present, but we weren't able to decode the information directly. Somehow the transformers' own circuits are essential for decoding the information. That's a pretty interesting insight. The mental model I have is like the transformer has its own dictionary in its weight somewhere of how to decode concepts into sequences of words. I love the example that you chose in the final version of your paper, Back to the Future, right? Like the movie title, Back to the Future, this is longer sequences of words. And I imagine that it's not actually directly in the hidden state, the instructions of how to decode Back to the Future.

• 45:46Speaker 00

  But it's more like a pointer. It's just this compact vector that you have to look up in the model weights somehow. And it's that soft prompt that Koyana learned that is triggering some mechanism that causes this lookup to happen that eventually rolls out to the whole phrase back to the future. So I thought this is pretty interesting. If the decoding was simpler, we would have found it through one of these direct decoding methods. But it's something that's the encoding is complex enough that the transformer itself is needed.

• 46:14Nathan Labenz?

  Yeah. So something I said earlier, maybe I'm now thinking is perhaps wrong. Cause I had said, okay, two linear transformations don't seem better than one. And I jumped to the conclusion that like you could train a whole transformer, but it probably wouldn't work. But now I'm thinking maybe that's exactly what would be needed. What does your intuition say if you were to scale up the direct, to go back to those first two methods? But instead of using a linear transformation, what if you used a full transformer to try to do the processing? Do you think that would work?

• 46:43Speaker 01

  Yeah, I think that should work better. Again, with when you introduce transformer, there's like the non-linearity thing, which was required for the soft prompting, like instances. So ideally it should work better, but that's not something we tried out.

• 46:56Speaker 00

  I don't feel like we had enough data to try it.

• 46:58Speaker 01

  Yeah.

• 46:58Speaker 00

  So, you know, this was a relatively small model training exercise you did, right? Where, where at every layer and every token, all these different settings, you're like training lots and lots of models. If you wanted to make a prediction on a large model like this, you'd probably need to train on millions or hundreds of millions of examples. And at that point, it's a little bit different. It's almost like the task that you're trying to learn, you might be best off just memorizing the whole thing. You could just say this vector corresponds to these tokens a few ahead and just memorize the whole problem. The thing that really surprises me about what Koyena found was that she didn't need to do that. With a really pretty small amount of data, sort of a hundred thousand samples, she could unlock the dictionaries

• 47:41Speaker 00

  for all of these phrases. And it wasn't like we're really training the transformer to do something new. That's, that's why it feels like we're unlocking the knowledge that's already in the transformer. Does that make sense?

• 47:53Nathan Labenz?

  Yeah, that's really interesting. That's super helpful. And I understand exactly what you're saying about 100,000 data points would not be enough to retrain a transformer. Maybe it could do the memorizing, but it certainly couldn't possibly relearn all the stuff that the full GPT J 6B knows because there's no way to get there from such a small sample size. So in theory, it could maybe work if you had the full pile and you ran the full pile this way. But obviously, that's not feasible on the resources available here. So you can't recreate that magic.

• 48:25Speaker 01

  When we basically add in another transformer, that kind of ends up becoming like a chicken and egg problem. Like, we're trying to understand this transformer, but then we now have to try to understand the method transformer.

• 48:34Nathan Labenz?

  Yeah, it's a bit absurdist, no doubt. And it becomes almost anti interpretability at some point as well. Yeah, that's definitely an intuition building exercise, more so than a way to create insight into how the things are working. How big is the soft prompt? I assume that individual tokens are embedded in the same 496 space?

• 48:54Speaker 01

  Yeah, so the soft prompt size was like 10 by 4096, assuming 10 tokens and each of this token having 4096 vector size. Yeah.

• 49:03Speaker 00

  How's that compared to your linear models?

• 49:05Speaker 01

  Oh, linear model, like the input was like 1 by 4096.

• 49:08Speaker 00

  Yeah, but those models were like 4096 by 4096. Like, the soft prompt was way smaller.

• 49:14Speaker 01

  Yes, that is true. It is much smaller indeed. Yeah.

• 49:16Speaker 00

  So even though it was the best performing thing, it was like the smallest number of parameters.

• 49:21Nathan Labenz?

  It'd be like 1000x difference, right? Because the linear transformation were 4,000 times 50,000 parameters, and this is 4,000 times 10 parameters.

• 49:31Speaker 01

  That's true. But I guess like the trade off is the linear model is trying to understand everything what we technically do have.

• 49:36Speaker 00

  We have the transformer helping us out, but the weights are frozen.

• 49:39Speaker 01

  That's true. Yeah.

• 49:40Speaker 00

  That's right. It's the power of using the transformer to explain itself.

• 49:43Speaker 01

  Yes.

• 49:44Nathan Labenz?

  Yeah, okay. That is really definitely very insightful. Did you try also varying the length of the soft prompt? My mind goes to what if there were no soft prompt at all?

• 49:53Speaker 01

  Yeah, no, my co-author, Jude, he definitely tried out a couple of variations of the prompts and found out that the size was like good enough in that sense for this. And I'm not sure if we tried zero prompt, but I assume that would work similar to a fixed prompt size, just because it is like a fixed prompt.

• 50:11Speaker 00

  Yeah, it's something we don't know in detail. We didn't see anecdotal evidence that it made a huge difference to make the prompts at different sizes. So we just picked a relatively small prompt and tested it.

• 50:20Nathan Labenz?

  Can we go back to the fixed prompt again for a second? I'm still a little kind of not content with my understanding of why that would actually perform worse than the bigram. I don't know, I would expect it to do better. I'm struggling. That feels like such a natural idea that it seems like it would work better than just such a simple statistic. And yet it's worse. And all of the ones you tried were worse? Like there were none. I assume that you didn't pick the worst of the.

• 50:46Speaker 00

  No, no, this is the best one. My intuition is this. So we show these various successes. The function vector paper is an example of this, where if you like go and intervene directly in the hidden states or the Rome paper, you go directly intervene in the weights of the model and we show, oh, it does something amazing. But actually, this isn't generally the case. Imagine sticking a random probes into somebody's brain and then saying, hey, you know, I wonder if this will make you smell the scent of bananas. But you know, generally, it's not possible to have a success with an experiment like this. So despite our best efforts at prompt engineering, some very clever thing for the model to be set up to do this.

• 51:28Speaker 00

  Just jamming in this hidden state at the end of these engineered prompts, it just didn't work. It's just doing some sort of brain damage. It's really not getting it to predict anything very useful. Obviously, like you said, like a natural thing to try.

• 51:43Nathan Labenz?

  Yeah, fascinating. Certainly, we can say it's out of domain, right? You're putting it into a state that is clearly very unnatural, where all of a sudden, and you still have more computation to go from that particular place of patching. So yeah, it still feels weird, but it's far enough out of domain that it just can't handle it.

• 52:03Speaker 00

  Possible that there is some way to make this work that we weren't able to figure out. So we tried our best with a bunch of different setups for this sort of manual prompts. And this is the best example that you can see. But it's certainly possible that there's some structure that we're missing or some clever way of setting it up that would make it work better. In fact, so we could talk about the function vectors paper also later if you want. There are some other tricks that you can find to get these kinds of interventions to actually work pretty well.

• 52:29Nathan Labenz?

  I was curious if people like me are still feeling like they want to go try another fixed prompt and see if they can't find the the magic words to make it work. What does the tooling look like for all this sort of stuff? What does the coding look like?

• 52:42Speaker 01

  There are no variations or ways to do about this, but when I was working on the fixed prompt stuff initially, I was using actually David's bow kit, a tool which allows me to look into a model, trace it and grab the hidden state and put it in like another run. But then there is a very recent tool called insight, which also does this. It's basically like a successor of bucket and it allows us to do all these like patching interventions and all these just basically few lines of code. Yeah.

• 53:14Speaker 00

  You have the open source release for this project, right?

• 53:16Speaker 01

  Yes, we call it future dot bio lab dot info. And yeah, over there we do have the code as well. And yes, basically in the code I have like notebooks and scripts and yeah, people can refer to that in that sense.

• 53:29Speaker 00

  So yeah, so definitely check out Kait Kait's repo for trying to people can try their hand at seeing if they can get a clever prompt that works better or some variation on the method. We'd be delighted to see if there's some clever trick that we missed. There's definitely room for creativity in this direction.

• 53:44Speaker 01

  Absolutely.

• 53:45Nathan Labenz?

  So you you give a pretty good intuition already for what that would look like. You've got the open source code. You've got a library that's specifically designed to facilitate this sort of activation, extraction and patching and the the loops surrounding that. To actually run a test, I should be able to do this probably in just even like a Google collab notebook, right? I I would think.

• 54:05Speaker 01

  So because of the model and everything, you might need a Google collab probe. I haven't run the training part, but again, the testing part with the future lens, I've tried it on collab pro and it works. Yeah.

• 54:14Speaker 00

  Have you used the end of backend?

• 54:16Speaker 01

  Yes, I know. I use it on the server actually, but not necessarily remote equals true because I was like, I'm on the server anyway.

• 54:22Speaker 00

  You're a little bit of a cheater because you're you're sitting next to the implemented.

• 54:25Speaker 01

  Yeah.

• 54:26Speaker 00

  But actually, a lot of the students who are playing with future lens and these kinds of things, that library was developed by Jayden. Good friend Jayden. And the library supports a remote backend that we're trying to get funding to to provide the scalable free backend for people to do interpretability research. We call it the N-DIF backend. It's like a deep inference fabric. And so if you code your interpretability experiments using this idiom, this little library, then, well, if you have the GPUs to run your model, then you can just run them locally. But if you don't, if you need resources to run a model that is too big for your laptop or for your co-op environment or something, then you can just flip a flag.

• 55:10Speaker 00

  You can just say use a different backend and then the experiment will be run on a shared remote backend for you. So, yeah, there's, we have GPTJ running and we've got different Llama models.

• 55:23Speaker 01

  Yeah, I think the seven billion model as well is.

• 55:25Speaker 00

  Yeah, like Llama 70B, which is a pain in the neck to run on your own, but we have that running.

• 55:30Nathan Labenz?

  Cool. So you have a single library that kind of abstracts away the subtle, I mean, all these things are transformers, of course, but they have differences in terms of their implementations, vector sizes, et cetera?

• 55:41Speaker 00

  Oh, no, they're not all transformers. Like, we have folks using it to study Mamba, for example, these state space models.

• 55:48Nathan Labenz?

  That's definitely of strong interest to me.

• 55:50Speaker 00

  Yeah, it's pretty fun stuff. Okay, so people who are interested in this library, it's embryonic, but it's pretty cool. And there's a really nice dynamic community around it. So the library is called NSIGHT, nsight.net is the URL to get to it. We haven't really promoted it, so I don't know if you search for it on Google, if it'll come up even, right? It's a little bit of a secret alpha stage project right now. But if you go to nsight.net, there's like an icon in the corner to link to a Discord channel to join, and there's a bunch of tutorial and documentation on it. But the really valuable thing is to just join the Discord and ask around, and there's a really friendly community that's developing it.

• 56:30Nathan Labenz?

  That's awesome. Is there anything that we can say about the bigger or more intensively trained models as compared to GPTJ? My hypothesis would be that you would probably find higher success rates with bigger models because you would, in theory, hope that they have more semantically meaningful middle state representations.

• 56:52Speaker 01

  Yeah, I would go along in the similar direction as well. It's just that with bigger models, you also need bigger data set for it to potentially train and be more exhaustive in that sense. But like in terms of takeaway wise, at least till now, for example, the whole middle layer stuff, I think it should still be something that is pretty transferable like other models as well.

• 57:12Speaker 00

  As the advisor person here.

• 57:13Speaker 01

  Yes.

• 57:14Speaker 00

  I'll say this is a gap in our knowledge and we we need to run some more experiments. So you're right. I think the the intuition that you have, which is that there might be more structure as you get larger models. And it actually might become more interpretable in this way is really interesting. It's a little counterintuitive. Traditionally, people would imagine that as models get bigger, they might become less interpretable and harder to understand. But this intuition that you're sharing here that maybe it becomes more structured, more predictable, easier to understand. That might be true, but it would be quite a claim. So more research is is definitely in order to systematically take a look at these kind of questions.

• 57:51Nathan Labenz?

  Is there anything interesting that you can share about the probes that have been sent into the Mamba universe so far? I did for reference a two and a half hour monologue about the Mamba paper in December. And I'm counting actually along with another like minded states based model fan. I think we're up to very close to if not hitting twenty papers now published downstream of the Mamba paper with different remixes and hybrid versions all over the place. So yeah, have you found anything interesting there?

• 58:22Speaker 00

  You know, all these attention free architectures, you you could basically ask, how are they doing it? You can try to lay things like logic lens on top of them or try to build future lens on top of them. And to some extent, there maybe in surprising ways, they show some similarities to transformers. Now, we're we're sort of halfway through doing some research to try to figure this out. And so I I don't want to lead any of your listeners astray by saying something that's not actually true. But yeah, even though they are missing a traditional attention mechanism, they seem to be able to do some of the sophisticated types of computations that transformers use attention for. So exactly how it all works, we're still. Trying to untangle it.

• 59:06Nathan Labenz?

  Yeah, cool. All right. I'm looking forward to that. If you don't have a striped hyena in that menagerie already, definitely throw one of those in there at some point as well because I don't know if you've seen just from the last few days, there was a usage of the striped hyena to do a DNA foundation model, which is starting to show some really interesting properties. First of all, super long context and strip pain is it's a hybrid. So it has attention and also the state space. And my sense definitely is that some sort of mashup is going to be the winner for a lot of scenarios in the future. Certainly in this DNA paper that just came out, they showed that the transformer was the least effective. And then Mamba was a bit better,

• 59:47Nathan Labenz?

  but then the hybrid attention and states base was the best performing. And it is kind of uncanny because it's sort of like, if you believe the theory that the language models are learning some sort of world model representation through next token prediction, by analogy, this seems to be learning some sort of like life model, or maybe better said would be like cell model at the level of just next base pair prediction. It's a whole other world that's got obviously its own tremendous interpretability challenges, but definitely fascinating space. So plus one for the striped hyena in your next round of experiments.

• 1:00:23Speaker 00

  We will be investigating these things carefully. It's a very exciting time as the world potentially moves on to whatever comes after transformers. It'll be very interesting.

• 1:00:31Nathan Labenz?

  Yeah, I always say it's not the end of history. Do you want to talk a little bit about the function vectors paper?

• 1:00:36Speaker 00

  So the the basic question for function vectors was Eric was very interested in asking, how does in context, learning work? It's just this amazing thing. I mean, it was the foundation for all the excitement around GPT three was this ability for a model to seemingly learn how to do a task after seeing a half a dozen examples of the task being done. Then you give it another half worked example and it'll work it the rest of the way. It's just amazing that these models are so good at doing that. And so what Eric was looking for is, is there some localization in the in context, learning task? What he found was that there is the bottleneck in it between

• 1:01:17Speaker 00

  when you give a set of demonstrations and when you ask a model to generalize. Yeah, come and join.

• 1:01:24Nathan Labenz?

  Eric Todd, welcome to the Cognitive Revolution.

• 1:01:26Speaker 00

  So yeah, Eric is the author of the function vector paper. So actually, maybe now that you're here, you can just sort of. Explain what it is.

• 1:01:33Speaker 02

  Yeah, so we found that there are like a small set of attention heads that when you process an ICL prompt that mediate the identification of a task. For example, let's say you give a bunch of examples of like antonyms of big and small, short and tall, and then you give it another query word like a bright. And the model, you can extract the output of a few attention heads from some other context where they saw a bunch of antonym pairs and take this output and stick it into this new context and it'll give you the antonym of the word that it's processing.

• 1:02:10Speaker 00

  So if you take the same attention heads.

• 1:02:12Speaker 02

  Yeah.

• 1:02:12Speaker 00

  And you read out what those same attention heads are saying. If you said Paris, France, Moscow, Russia, Washington DC, United States, and then you said Ottawa, blank, right? But then you read the attention heads at that moment and then instead you have a totally new context where you have Ottawa or something like that, or you have Madrid or something like that. Even though those attention heads were the ones that seemed to cause you to say antonyms in one of the experiments. In this experiment, if you read out the attention heads for the country capital task and you stick it in to the new context, it seems to encode the new task.

• 1:02:56Speaker 00

  It'll actually say Madrid goes to Spain instead of doing antonyms or something like that. And what are the types of tasks you tested? You tested a comment?

• 1:03:03Speaker 02

  Yeah, we tested like forty plus different tasks and it seems like they're all mediated by this small set of heads, which is kind of cool that it seems like the the model has this sort of path that it communicates this task information or like this bottleneck, I guess, that it's communicating what the task is that it's doing. Even though in the prompts were never explicitly telling it what the relationship between the demonstration and the label are, it's able to figure that out and communicate it forward.

• 1:03:32Nathan Labenz?

  Yeah, this is cool. I'll just resummarize it, make sure I have it right. So you, first of all, set up a task implicitly with a few-shot prompt, which is the original, GPTs are few-shot learners, right? The original paper title, "Few-shot learners" refers to the ability to understand or infer and actually do the task just from a few examples. So, pretty fundamental, notable behavior, without question. So, okay, set up a couple of these examples, thus creating a few-shot prompt. One example is antonyms, another example could be translation. So, "arrived to part, small, big, common, rare." And then, at that point, as the pattern has been established,

• 1:04:13Nathan Labenz?

  you go looking inside the transformer for, is there some... I'd love to hear a little bit how you think about the term bottleneck, but you're going and looking for some sort of relatively small dimensional thing that you could extract and patch over to another context. And now, when you find that thing, and obviously that's where all the hard work is, but then the demonstration is that when you have found that thing and you do in fact extract it now in a clean setting, without any examples, you can set something up and say, "The word 'fast' means..." and then inject this representation of, "Give me the antonym." And instead of doing what a language model would normally do, which is define what fast means,

• 1:04:55Nathan Labenz?

  it now says, "The word 'fast' means 'slow,'" because you've forced the antonym behavior with the patching from the few-shot context now over to the clean context.

• 1:05:07Speaker 02

  Yeah, you're right about bottleneck. So, attention seems like a natural place to look, just because that's how transformers move information between tokens. And we figured, in order to gather information about the task from your context, you'd have to do that via attention. Like, if you never explicitly tell it what the task is, you need to go back and look at all the other tokens to figure out what you're doing. And attention is, it's nice because attention heads act on a small subspace, and so they're natural bottlenecks for sparse activations.

• 1:05:40Nathan Labenz?

  So what is it exactly that you are extracting and patching over? Is it the output from an attention head?

• 1:05:47Speaker 02

  Yeah, that's right. So you can take the output of these, the small set and just add them all together. And that's the actual thing that we call a function vector is this sum of attention head outputs. And then you can add this vector into some other prompt at a particular layer in the model.

• 1:06:03Nathan Labenz?

  So it's one vector that represents the sum of multiple different attention heads, but you're selecting which attention heads to look at. And obviously a big part of the challenge is figuring out which ones to look at. Once you have that sort of thing, does the relationship between tokens need to have some sort of length similarity to the original length relationship or not really because all of the information passing between tokens has already happened by the time the attention head is done, right?

• 1:06:32Speaker 02

  It doesn't need to. It doesn't seem like that matters that much. Seems like fairly robust to different settings. So we tried sticking it into like natural text settings, different templated settings. It seems like it works just the same for long or short. Yeah.

• 1:06:48Nathan Labenz?

  Isn't that cool? How do you figure out which attention heads to use? And what can we say about the relationship between the different attention heads that together make the vector? It's nice to think of this like single vector, but then it's weird to think of these different heads contributing to like part of this sort of single conceptual vector.

• 1:07:08Speaker 02

  Yeah. So we found that a lot of these heads have had a similar attention pattern, which is kind of interesting. Most of them were attending to previous label tokens, which is like you have demonstration label, demonstration label. So they're basically like looking at all of the answers of all the pairs, which kind of suggests that maybe, and some other papers have found, like when you do ICL, these label tokens seem to be sort of places where the model stores important information. And so there must be something about the pairs where it's able to figure out the task and it's transporting that task information from those labels onward.

• 1:07:47Speaker 02

  The way that we found these heads was we do a process called causal mediation analysis. Also, people sometimes call it activation patching. So basically, the setup is you have one sentence with a bunch of pairs of words demonstrating a task. And you have some other sentence where you can set it up where it's maybe demonstrating some other task. We decided to scramble all of the labels so that it wasn't really clear what the task was. And then what you do is you patch the activations from one attention head into this other setting, where it's not clear what the task is doing. So you take it from the clean setting into the jumbled setting

• 1:08:30Speaker 02

  and measure how much it influences the model's output of the word that you would expect the behavior to induce. So for example, if in your jumbled setting, the query word is like big, and you paste in an attention head output from an antonym setting, you'd expect that the model's output for small would go up slightly. And like, just doing one attention head isn't enough to flip the prediction. So we started trying multiple attention heads. And it seems once you get to around like 10, you can start having the significant causal effects of getting the model to flip its prediction and start understanding what the task is.

• 1:09:09Speaker 00

  So to me, the big surprise was Eric got this working for about 10 attention heads for like antonyms. And then he looked at something else. I think in the paper, we looked at like English to Spanish translation or something. He had a bunch of different tasks and he went and he started looking at the attention heads for these other tasks. And we noticed that they were largely the same attention heads. And so if you just picked up your attention heads from one task and used them on a different task, that whatever information is in those same attention heads would tend to induce the new task too. Which was like, it was really surprising. Like, I could imagine that there's some antonym attention heads. But it was very interesting to find that there are these attention heads which are just like general function.

• 1:09:52Nathan Labenz?

  Task understanding.

• 1:09:54Speaker 00

  Yeah.

• 1:09:54Nathan Labenz?

  Yeah, that's really interesting. And it looks like from the figure three in the paper, they're found mostly in the middle layers, but not like all at the same layer. They're kind of mixed at different parts of the middle layers.

• 1:10:05Speaker 02

  Yeah, most of them in the middle. Maybe there was like one towards the end, but yeah, it matches our intuition. I sort of see transformers as having two processes. One is figure out what I want to from the context and then use that to predict the next token. And so like, You could see maybe the first half of the network doing the figuring out and the second half using that to predict a token. And like, it's cool that right in the middle, right before it like starts to predict a token is where this process is happening.

• 1:10:33Nathan Labenz?

  Yeah, that's fascinating. Obviously, recently, mixture of experts has been a huge trend, right? With your GPT-4 allegedly being a mixture of experts model and your Mixtral, and of course, now Gemini 1.5 has been declared to be a big mixture of experts model as well. So I've been looking into what is known, what is understood about mixture of experts. And it's like, most of the mixture of experts papers are focused on swapping out the MLP blocks, because that's where like, most of the computation is, at least if the context window isn't huge. But now, obviously, with context windows getting huge, you know, that's starting to even out, or maybe even tip the other direction. And so I guess people are also increasingly starting

• 1:11:17Nathan Labenz?

  to do mixture of experts where the attention blocks are also switched out at runtime. So I wonder how this kind of finding should help me think about that. It certainly, at first blush to me, sounds like, geez, that's a pretty complicated structure. Yeah, there is like, you've got 10 different heads, and you could have had more, right? If I understand correctly, you kind of drew up arbitrary line and said like, this gets us most of the 80, you know, Pareto principle type of thing. So it definitely seems like it's a sort of distributed process. The process of identifying the task is a distributed process. It's kind of like mixed throughout layers. Yeah, it is interesting. There's like, two in the first layer where they exist, two in the second layer, one in the third, one, one, one, one. And then it goes quiet, at least in your top 10. I wonder, maybe one way to think about it is, maybe

• 1:12:02Nathan Labenz?

  there's like, one to two of these per layer? I don't know. It's an interesting just mashup. I wonder if you have any intuition for how this would help us understand mixture of experts?

• 1:12:11Speaker 00

  Let me answer in my interpretability professor advisor way. I guess the way I look at it is this, is that we don't know yet what is the right level of abstraction that we should be looking at to understand these things. So we know the physical things that we made them out of, that we have dimensions, we've got nonlinearities, we've got modules of different types, we've got these MLPs and attention layers and cool. State space recurrences and other neat things that people have suggested. And so in the end, it's all just matrix multiplications, right?

• 1:12:48Speaker 00

  We know what we're building these things out of, but we're not sure yet exactly the right way of looking at all of these computations. And the attention heads, so Rome was a way of looking at the MLPs and saying, hey, there's this interesting association structure in these MLPs. It seems to be pretty informative, although the amendment sort of qualifies that and says, well, you know, might be distributed on a few layers and it's a little bit more subtle. And all of this great mechanistic interpretability work, I think, has been a way of looking at all these attention heads and saying, oh, gosh, these attention heads also reveal a lot of interesting computational structure. And I think that motivated a lot of Eric's work. He says, oh, is there an attention head? Maybe there's a few.

• 1:13:34Speaker 00

  And so there's definitely a structure in the attention heads. To me, it's not 100% clear that's the end of the story. That it may be like when Eric finds that a dozen attention heads are working in concert with each other, then it leads to the question, is there some abstraction? Is there some way that we should be looking at these attention heads altogether? And I actually see function vectors as one proposal for a way of doing that, right? You can summarize all of these attention heads in a single vector. And even though we've sort of blurred them together and we've lost a bunch of the fine grain distinction, that that vector has pretty strong causal effects. It has some interesting properties that Eric investigated in the paper.

• 1:14:16Speaker 00

  Like to some extent, these vectors can be composed and do interesting things. And so I don't know if function vectors are going to be one of the abstractions in the end. That's like the right way of looking at what these things are doing. But what it hints at is that there's potentially more structure to be found, maybe at a higher level even than what we've been looking at so far. And to me, the big game, the research game, the chase that we're all on in the interpretability world is to figure out the answer to the question, what is the right level of abstraction?

• 1:14:47Nathan Labenz?

  That's great. That might be the perfect note to end on. And you guys are doing admirable and certainly laudable job of chipping away at this massive problem. I'm a big fan of all the work that I see. Anytime I see a Val Lab paper, I'm excited to get to the website and you guys do great visualizations as well. So definitely encourage people to go to the individual paper websites and see the visualizations, the animations.

• 1:15:09Speaker 00

  For visualization, let me give a shout out to Nikhil, who just archived his paper that's also going to be at ICLR. We didn't talk about it today, but he has the best circuit visualizations out of anybody. He really has done a really great job at putting those together. So check out his paper. It's called "Fine Tuning Enhances Existing Mechanisms." So it's this whole study of the interplay between circuits and fine tuning. And it's also an interesting paper, but really beautifully presented as well.

• 1:15:34Nathan Labenz?

  Yeah, we will save that one for next time. This has been a great conversation, guys. Eric Todd, Koen Napol, and Professor David Baugh. Thank you for being part of the Cognitive Revolution.

• 1:15:42Speaker 00

  Thanks, Nathan, for having us.

• 1:15:44Nathan Labenz?

  It is both energizing and enlightening to hear why people listen and learn what they value about the show. So please, don't hesitate to reach out via email at [email protected], or you can DM me on the social media platform of your choice.