David Bau: Oh.
David Bau: recording in progress. I'm gonna have a German cookie. Can I have a German cookie? I'll report to you if these are really fancy.
David Bau: So, you guys, it's our great pleasure you're welcome.
David Bau: Jack Murillo. Not quite Dr. Jack Murillo yet, maybe days away. You know, wrapping up his dissertation.
David Bau: And you know, he, you know, obviously.
David Bau: as you all know, Jack has been doing such interesting work in our field and
David Bau: and so I'm I'm not sure specifically which which papers he's going to talk about here today, or if he's going to talk about things that are not yet published. But but I'm looking forward to it. So welcome, Jack.
David Bau: Yeah, thanks for the introduction. And thanks so much for having me.
David Bau: I think that this talk is gonna be mostly about
David Bau: one paper. So I don't think I'm really going to like answer the
David Bau: or give a very satisfying answer for the title of the talk. I think it's going to be like.
David Bau: partly kind of aspirational. And since this is kind of like a familiar crowd from you guys hoping to like, talk a bit about
David Bau: maybe the answer to like this question, or like the answer to the the title,
David Bau: But yeah, so the the main thing is that
David Bau: I want to talk about like decomposing neural networks into features. This is like one aspect of
David Bau: interpretability that has gotten like a lot of like over the years. A lot of attention. But like, I think, especially recently, like dictionary learning type stuff, has has really also pushed it up again.
David Bau: I think some text got cut off there. But decomposability, meaning just that the the neural network predictions can be like broken down into
David Bau: sort of some of its parts, into like independent parts of the the model.
David Bau: and I think that this is going to become like even more relevant again, when more laws like the EU right to explanation law come up
David Bau: cool.
David Bau: So I think, like as as language models and neural networks become like a bigger part of our lives like we'll see more laws like this. And like
David Bau: answering these kinds of questions becomes more important.
David Bau: So decomposability is like.
David Bau: really easy. When you have a really simple model like this perceptron that does like the Xor task, you can basically
David Bau: decipher exactly what every neuron is doing. This model is like entirely decomposable, but obviously, like
David Bau: Chat Gpt is not so. This is the I didn't use this abstract, but this was the the
David Bau: prompt that I gave it when I gave my thesis proposal.
David Bau: And I wrote a pretty sensible abstract, and there's.
David Bau: you know, even if we had the weights of Chat Gbt. It's a huge
David Bau: ask to try to like, decompose the
David Bau: every little feature that goes into like producing this output.
David Bau: But recent work has shown that we can decompose models a few different ways.
David Bau: one way is like on the neuron basis. So
David Bau: a lot of you are probably familiar with this paper, which is about decomposing, like the next word prediction into like individual contributions from like neurons and Nlp layer. There's like a lot. I could have put it like a ton of papers in here. I'm going to talk about a few of these later. So
David Bau: you know, just pointing out that there's like a lot of work on
David Bau: interpreting individual neurons in like Mlp. Meetings or attention heads.
David Bau: And then there's of course, like the sae stuff that is
David Bau: like, really like, taken off in the past couple of years which is basically like.
David Bau: you just decompose some representation into an even larger vector, that's much sparser, so that you can try to understand, like individual
David Bau: neurons in like a very, very wide basis.
David Bau: So for example, like it's just a graphic. If you're like a little unfamiliar, you train a very, very wide Mlp. And individual neurons in that dictionary correspond to like interpretable features, like words ending in ing or chemistry terms. Stuff like that.
David Bau: So this talk is going to. I like mainly I'm going to be talking about
David Bau: this paper talking heads which I presented on nerves in December.
David Bau: and it's like really good fortune that Mikhil gave the talk right before this, because, like, I, too, care a lot about abstract positional representation. So I think, like, there's
David Bau: a lot of relevance. And they still did a lot of the heavy lifting, talking about like some circuit analysis type stuff.
David Bau: And then I'm gonna be talking about some like work that's
David Bau: in prep that I'm like hoping to finish up soon. And those results are still kind of cooking. So I'm I'm sort of showing you some like
David Bau: early ish stuff. And like, I basically just want to talk about that. And there's like some other things
David Bau: I've been thinking about kind of like, and particularly interested in
David Bau: and more generally at Brown, like the lunar lab, is thinking a lot about like
David Bau: what we can actually get or conclude from different
David Bau: interpretability methods just more generally different like
David Bau: circuit analysis methods or decomposition methods so interested to hear your thoughts on that, too.
David Bau: So this paper talking heads.
David Bau: The main motivation for this was
David Bau: understanding like how different components in a transformer like send signals up through the layers. So like, how does this attention had at layer one communicate to attention at a layer like 20 or something.
David Bau: And this sort of like stemmed from like there's so much circuit analysis work
David Bau: from the past couple of years that
David Bau: just kind of like uses this language really freely like. So in this circuit, if you're familiar with it, and I'll talk about this specific task in a minute. But like.
David Bau: basically, you get some some conclusions like, Oh, this previous token had it looks at the
David Bau: there's an intention that looks at the previous token sequence, and then it passes a signal about that to this later induction head, which uses that to then send a signal to the inhibition head. So there's all this communication through the layers from like layer 2 layer 5 to layer 7 or 8. But it's like to really undetermined like
David Bau: how that's actually happening through the model. So that was sort of the basic question which is like, when you say something like this, attention head is passing a signal this other head like, what do we actually mean by that?
David Bau: And if you know, if we know that there's like this communication happening.
David Bau: can we like find that connection in the ways. And so another part of this was like
David Bau: I was interested in, like weight, based interpretability. So this seemed like a a logical place to look, which is like
David Bau: empirically. We know there's like communication between these weights. So then can we find it in the weights themselves without any data.
David Bau: And I was thinking about this a lot.
David Bau: And I was rereading the mathematical frameworks for transformer circuits. Paper from anthropic from
David Bau: now 4 years ago.
David Bau: And they, like just kind of in the beginning, are talking about like the residual stream as a communication channel. So this is like a kind of like obvious.
David Bau: I guess, like perspective, I'm holding in my head. Which is like residual stream, of course, is passing up these signals. But they also have this like really interesting point about virtual weights, which is that because the residual stream is connected through like additive updates.
David Bau: you can multiply out a 2 way matrices like the output of layer. One
David Bau: writes into a subspace, which is then later, you know, up and up and up through the residual stream is is read out by like the input weight matrix of like layer 3 or something. And they arrived. This like nice composition score formula for this, which is
David Bau: just telling you like intuitively, how much does we matrix at layer 3 reach from subspaces written to by weight. Matrix is that layer one
David Bau: and this is just taken by the the previous norm of the matrix product over like normalized by the individual probeness norms of those matrices. So you get a nice score from 0 to one where one means like
David Bau: this, matrix is writing only into subspaces or reading or sorry. This weight matrix is reading, like only from subspaces written to by the earlier weight matrix and 0 means, they're like orthogonal matrices.
David Bau: And so, yeah, I was like, okay, this should be such a slam dunk. We should just be able to like, take these weight matrices from the circuit and apply the composition score at. Why has no one done this? And we should be able to read off like, oh, there's signals being passed like everything we can just see them without running any data, so we should be able to like see the whole circuit without
David Bau: without doing anything.
David Bau: so I I was particularly interested in this, like one part of that circuit. Just that I thought was cool in the paper, which is from the inhibition heads to the mover heads
David Bau: which is just the example I'm gonna be looking at right now, which is
David Bau: inhibition heads write out a signal which is then used by the queries, the query matrix of the mover head, so that the query
David Bau: of that head is is impacted by this earlier earlier inhibition head
David Bau: And so, before I like, jump into the results. Just wanted to go over
David Bau: some different kinds of composition that can happen. So I showed you this like generalized formula here, which is just abstracted weight matrix, one weight matrix, 3
David Bau: but there's actually like a few different kinds of composition that can happen with attention. Heads
David Bau: query, composition is this one that I'm gonna be talking about, which is
David Bau: between the ov matrix of the earlier layer, which is, of course, like, if you're familiar with this like terminology, ov is just the
David Bau: its own weight composition composition between the value and output matrix. So the attention had bad
David Bau: writes information out into the residual stream. So the thing that writes into the residual stream for an attention. There's always this ov matrix
David Bau: And this writes out information that the the query matrix or the Qk matrix query key matrix reads from. And so the output of this
David Bau: this ov matrix had an earlier layer
David Bau: informs the query like what to look for in the keys.
David Bau: like, that's adjusting the the keys that are most activates to hopefully that makes sense.
David Bau: So then you just get this this formula same, we're just like dropping in the Ov and Qk matrices here.
David Bau: And so you do this between this, this inhibition head to the mover head that I just mentioned.
David Bau: and it looks
David Bau: not good, basically. I think that if you so just just qualitatively, if I asked you to point out like what the
David Bau: so sorry. Let me explain the heat map first, st before I get ahead of myself. So this is the layer and the head. So this is the the composition between every attention head in
David Bau: the Gpt. Too small here.
David Bau: How much it composes with the qk. Matrix of of a layer, 9 head 9, which is our mover head.
David Bau: And so
David Bau: if you were to, if I try to ask you like, okay, we have this circuit, like, point to all the inhibition heads from that circuit
David Bau: like you probably wouldn't be able to find them. They're these ones, maybe 8 dot 10 has is one that you would pick. This, I think, has the highest composition score, but these other ones don't like look too great.
David Bau: and we look at another part of the circuit. So the
David Bau: composition through to we'll look at 7.9, which is one of these inhibition heads. So just taking a step back from the circuit
David Bau: through this time value composition, you get a similar story.
David Bau: Only thing I want to point out here before showing them is value. Composition is the same type of formula, but between 2 ov matrices, so output from a previous head affects the output of the current head back into the residual stream.
David Bau: So this red signal from the 1st ov
David Bau: interacts with the blue later over you, producing like a purple signal.
David Bau: How's that?
David Bau: And again, we see, like the same kind of thing you probably would not pick these like in circuit heads, as like being that you basically wouldn't be able to find these. Probably like
David Bau: has the most composing heads.
David Bau: And so maybe the problem is that like, you know, we
David Bau: we're just seeing like other circuits, like the the model contains like a ton of circuits, and we can't just figure out there's like one from this composition. Short, that's probably true. But just in general, like none of these
David Bau: circuit connections are like statistical outliers with respect to the composition score.
David Bau: Even though empirically, we know they're interacting just from this like circuit analysis work from long at all.
David Bau: And this is the I did find one other mention of the the composition score being used outside of like a toy model setting. And
David Bau: there is kind of the same story, which is that it's it's basically just very hard to interpret from
David Bau: in like these kind of like real models, even though Gp, too small is so tiny unless you have like a toy model. The composition score doesn't seem to be
David Bau: very useful on its own
David Bau: and so the the kind of like
David Bau: hypothesis I had after that was like, well, every matrix is just like.
David Bau: you know, full rank. It's writing out into like every little like subspace, and every matrix down the road is also just like reading from every little subspace. So maybe like, if we decompose the matrix to like, you know, smaller the subspaces that it's writing out into. Maybe we can try to like isolate the signal that way.
David Bau: And so I use the singular value decomposition
David Bau: on the Ov. And Qk. Matrices to do this. Then there's a nice identity. You can rewrite the Svd. With which is, you just rewrite it as the sum like element, wise sum of a bunch of rank, one matrices. So this gives you this like big stack of
David Bau: rank, one matrices that when you like, collapse them all together. You get the original matrix and
David Bau: with attention heads. This is also really nice because you get because the head dimension is usually a lot smaller than like the residual stream dimension.
David Bau: You, you get fewer of these matrices that you actually have to care about, because the rank of that matrix is
David Bau: only full with respect to the head dimension. So you only need to care about like in Gpt, too small. There's only 64 of these matrices. So it's actually pretty easy to like search over them.
David Bau: But like then, what do we actually do with these component matrices? There's kind of 2 things we can test. One is the I just talked about the composition score between matrices. The obvious thing is like, let's just use these component matrices
David Bau: to do the let's just do the take the composition score with these component matrices instead of like the full matrices.
David Bau: And so this is like a weight based
David Bau: metric that we can use. So this is again something like this, where we
David Bau: one at a time. So the the choice here was to
David Bau: take one component matrix. So in this case, this red one
David Bau: and take the composition score between that and just the full like later matrix, you could also just take a component matrix from this one
David Bau: but that felt like, wasn't actually necessary, because we wanted to see if if anything out from coming from here is particularly writing into this one. So
David Bau: that was the choice there. So it's the same same exact formula. We're just using this this individual matrix instead.
David Bau: And then the other thing we can do is that this also gives us a really nice like model editing procedure. So we can test the behavioral effect of a given component matrix by swapping out, or just rather removing one of the component matrices from the sum. So this is just
David Bau: rewriting the the weight matrix
David Bau: but without the ability to write out into that subspace.
David Bau: And so I kind of like
David Bau: very, very briefly went over the Ioi task. But to understand, like what we're measuring when we remove this component, I need to talk about what this task is actually like. And so the goal here is just to to generate, like Mary over John, the intended indirect object over the other subjects token. So we call this s. 1 s. 2, subject. One subject 2. And Mary is the intended indirect object.
David Bau: and so the circuit that they they found, like the the original goal of the Wong at all papers is, let's find the circuit, for like how the models how to do this.
David Bau: And so a key part of this is that
David Bau: the moverhead, which I kind of talked about earlier, like 9.9
David Bau: other such ads in the model. They attend to the name that they want to copy.
David Bau: and so, for example, the
David Bau: move ahead will attend to Mary.
David Bau: and as a result of attending to Mary, the next token is more likely to be Mary. If it attended to John, then the next token would be John.
David Bau: I keep the part of the circuit that I was interested in is the interaction between the inhibition head. Let me go way back. Actually.
David Bau: The interaction between the inhibition head and the mover head. The inhibition head is telling the query to inhibit
David Bau: the John token, and instead prefer the Mary token. So just say, attend less to John.
David Bau: That's the signal that that's being passed there.
David Bau: And so, as a result of of inhibiting John, like I said, the model, the mover has more like a tend to marry.
David Bau: and so the
David Bau: score we can take directly out of this is like the move ahead, attention to the indirect object versus the subject token.
David Bau: So, without any inhibition. The moverhead would attend uniformly to John and Mary. 50 50
David Bau: and thus produce, like.
David Bau: you know, the resulting next token prediction would be kind of like randomly chosen between John married.
David Bau: so the score we can just take it, like the attention to the indirect object minus the attention to the subject at between these positions? Is it just fully based off of the the previous word
David Bau: which found that these interactions were pretty predictable?
David Bau: Does that make sense. Before I keep going, I kind of like glossed over those components.
David Bau: Cool.
David Bau: so again, those are the 2 metrics we're taking as a weight based on the behavioral score. And so what we found was one at a time removing the components, or one at a time testing the composition score between components. We see like
David Bau: there's 1 component in particular that has a like very
David Bau: outlierish high composition score, which is this blue X here.
David Bau: and the result of removing that component at the same time has a very disproportionately high behavioral effect on the inhibition score. So here again, lower inhibition score means like more even attention to John and Mary or the subject and indirect object. This is across average, across a larger data set.
David Bau: And so it seems like there's this one component that that seems to really stick out here.
David Bau: and we can actually get a signal from the composition score that corresponds to a real behavioral shift as well.
David Bau: And there's 4 inhibition heads in the model that they found. And we found this effect. Basically, with all of them, there's 1 or 2 components that look like this particularly. There's 1 that has like a really high composition score.
David Bau: the one exception being kind of like 7.3. I'd say that there is still one. But, like the just the the general effect. It seems like this inhibition is just generally less important
David Bau: for the model entirely so we see less of a behavioral effect. So for future analyses, when I talk about some of these, we kind of cut out 7.3, because it just says.
David Bau: See, when it's convenient, just because it doesn't seem that important.
David Bau: So the kind of like conclusion from this is like when
David Bau: anthropic was talking about. Like, you know, communication channels between the residual stream. What we found is that there seems to be like
David Bau: a rank. One signal in these cases that the in which this is occurring. So
David Bau: another way to say, this is like the communication channel between these these layers is is extremely low rank. It's like just in this one direction, basically.
David Bau: maybe you have already talked about it. But could you
David Bau: describe it once again? What was the motivation for applying the Svd on the weight matrix in instead of doing it on the residual stream itself.
David Bau: Yeah, because we were talking about. I mean that you could do it on the residual stream. So
David Bau: I because I was interested in particular
David Bau: weight, like just weight analysis. I wanted to find the signal in the weights, like somewhere in the weights of the model. We already knew the circuit ahead of time, so we knew which weights to look at which made it really easy and convenient to know which ones to look at
David Bau: and the specific choice of the Svd.
David Bau: But it's kind of like other people have done it. Yeah, it tends to find interpretable things, for whatever reason seems to be somewhat Sv line. But but in that the reason for that one, instead of like other decompositions.
David Bau: was like, yeah, CD, just has a lot of nice properties like
David Bau: it's ranked by variance. Each one of these component matrices is orthogonal to every other component matrix.
David Bau: so by removing it or changing it in some way, you're not actually affecting any other part of the weight matrix, which is really nice.
David Bau: and what else? Oh, yeah, these matrices are just all like, since our outer products are all ranked one. So you're only affecting like one direction. So it's also, if you do the similar analysis on the residual vector itself rather than the weight weight matrices, I feel
David Bau: we may get up end up getting these 3 results as well.
David Bau: I think you would get the same results. Because,
David Bau: yeah, yeah, because you're getting just the accumulation of all of the outputs of all weight matrices. So they have to be in there. And I imagine, just like looking at a task specific level, like
David Bau: like the residual streams on this task.
David Bau: This must, as I think, be a high, like a
David Bau: a large component, like in the activations that. I basically.
David Bau: I'm not sure about that. Jack. Yeah, yeah. So I, you know. So I I feel like it's a
David Bau: you know. It might be a good appendix for this paper to like measure against baseline, because I kind of suspect that your wave based decomposition might be superior in some ways to an Svd. Just based on pulling out major components of the variance of the of the residual stream. Were you? Were you here for the? I think it was what's his name? Pan, the pan paper, the the one from transformers and vision
David Bau: where he's doing the decomposition of attention heads in a
David Bau: vision transformers. You guys remember that talk right? What? You maybe miss the talk. Yeah.
David Bau: yeah. Yeah. But my question was similar to that. I mean, if you can get to the same residual, the same result by analyzing the residual vectors. Then what? Why did you go for the weights? So you need no data.
David Bau: You need no data. You don't need the data. Yeah, I see. Yeah. Okay. But but besides that, right? So the the query key business, this is, it's it's this weird thing. It's a bilinear form.
David Bau: and taking an Svd composition of a bilinear form is a special thing that's different from taking Svd of of the data
David Bau: and so it might be that there's something that like, there's some new insight about doing that. It's
David Bau: yeah.
David Bau: Okay.
David Bau: And yeah, I guess you did mention like,
David Bau: you know, killstock also, like you would find clusters for like gender, which is not necessarily like which is not
David Bau: not the task, not the task. Yeah, but it is creating a lot of variance in that. Yeah. So that's
David Bau: that's true. Actually, I might be, I might be assuming too much.
David Bau: okay, so yeah, it'd be it'd be interesting to take a look at. But for sure.
David Bau: Yeah. So I'm gonna talk a bit more about have these like individual
David Bau: components. So we have again, like 3, maybe 4 inhibition heads that I just talked about. So it's like
David Bau: a few different components that seem to all be coordinating together. I'm going to talk a bit more about that coordination in a bit. So did you have a question? Oh, yeah, I was just wondering what you think the direction represents like. So if it's just one d, it's just saying, Yeah, it's just like a 10 more to
David Bau: like, like, is it like a positional thing again, like Nikhil's stuff. Yeah, yeah. So I'll talk about that. Yeah, totally.
David Bau: Just before going into that. The
David Bau: yeah. So before talking about. Yeah, okay, I'll talk about that right now. So since it's each of these component matrices, I kept them in matrix form because it's like, really just convenient. You can just plug it into the composition score formula. So we'd have to change that at all. But really like, since it's rank one. It's really just like outputting onto one direction. So it's really just a vector representation. And so what we can do is like intervene on
David Bau: this space. So we can overwrite the output of the attention head to just be like
David Bau: the vector that's this direction in space. That's being written out times, some scalar. And this is our just like intervention, like.
David Bau: How does scaling up and down this direction? Space affect the prediction, affect the query specifically.
David Bau: and we found that it is a positional signal
David Bau: and it's not really shown in these results. But I can say that it's a relative positional signal that tells you to either attend to name one or name 2.
David Bau: Basically what this is, showing that if we
David Bau: flip the names, the order the names, and we just scale up and down, we can control whether. So this is the inhibition score. Still, you can actually get the model to fully attend to like one name or the other, or not fully fully, but like basically fully prefer one name over the other. And it's it's purely based off of attending to like the name in the 1st position or the second position. And so this weirdly smooth curve. I was really kind of surprised by this result that it's like so cleanly interpolating in this head.
David Bau: So yeah, do you know any any questions about that?
David Bau: So that's what we we found that this I was doing. I have some. I can show the other heads as well. They have the same pattern control heads, like other random heads, did not have this pattern
David Bau: or random components within other heads, and even within these heads didn't have this pattern.
David Bau: So
David Bau: So the next thing, I guess, is that like, okay, so it's just going back to like this number. Specifically, the composition score.
David Bau: Oh, that was again a behavioral result on like scaling it up and down.
David Bau: but I showed those heat maps at the beginning which were like not looking so great. For the composition score. But if we, if we think of the composition score now, instead of the component basis. We can look at for each head. We can take the composition with like 9.9 again, but for each head only plot, for, like Max, composing component within that head. So for each head
David Bau: we take the composition score between each of the 64 components, and there's only plot the highest one.
David Bau: so I'm gonna again just go back to showing heads composing into component 6 of 7.9, which we also know to composer, is trying to do the mover head.
David Bau: This is what I showed you before. So these are the duplicate token ads and the induction ads composing the 7.9.
David Bau: Then we decompose it
David Bau: specifically into this component, which we know to be important for the other part of the circuit. We only get the in circuit
David Bau: heads like really popping out
David Bau: And so what this kind of tells us is,
David Bau: yeah. So this is just showing that
David Bau: I need to come back to this. I need to. I want to jump to this this slide. Actually. So
David Bau: the nice result from that is like we can go to 7 dot 9 dot, 6 to find things leading up to like this point, and then continue onwards to like 9.9 doing the same kind of strategy, 9 dot 9 component 0.
David Bau: And we find a few out of circuit heads really composing with this
David Bau: from other circuits we actually look into those heads in the appendix.
David Bau: But basically we can find, like the whole.
David Bau: almost the whole Ioi circuit, like directly from the weights without any without any data at all. Of course we know that the circuit is there. So that's how we know.
David Bau: Let's verify it. But there is, in fact, like a weight based signal that we can like get out of this just using the the components.
David Bau: If I go back a bit. If we focus just on
David Bau: the duplicate token, head 3 dot 0,
David Bau: you can also look at the components within 3.0 that communicate with 7.9 dot 6 layer, 7 head 9 component 6
David Bau: and you can look for if there's well, is there another like low rank communication channel. I kind of only talked about one up to this point, which is inhibition heads to mover heads.
David Bau: We found, like more evidence, for, like low rank communication channels.
David Bau: And in this case I'm showing a slightly different visualization. But this is each component and
David Bau: and the composition score between between these we found 2 that are like really large outliers, not one this time.
David Bau: And if you plot the activations projected along these component directions for
David Bau: random tokens where some are duplicated and some are not
David Bau: They seem to like nicely
David Bau: separate duplicate from non duplicate, like. Both of them seem to have some signal for this.
David Bau: Why, it's 2 instead of one I don't know. I don't really know if it's that important, but, like
David Bau: with some noise, it. It is seemingly separating these quite a bit.
David Bau: And then again, we can influence the inhibition score all the way back from layer Z. 3
David Bau: by just changing this signal. So in this case, since it's 2 dimensions or the 2 directions. Rather, we're kind of like going.
David Bau: I'm scaling them both at the same time, it's kind of like diagonally, like along both axes, either positive or negative.
David Bau: And again, you can get this like kind of crisscross shape.
David Bau: which is controlling attention to either name one or name 2
David Bau: although the absolute impact that you can get from doing it so early is a little bit less.
David Bau: probably because there's a few other signals that we're not accounting for. So before we could get like really, really high attending to one name or the other, and in this case we can only get it to like kind of slightly prefer one or the other. But we're at least getting the the cross shape that we expect
David Bau: versus. This is the same head and 2 different components, the zeroth component, which is actually the highest variance component, and the 4th component, like our index, 3 component. And we don't see that effect at all.
David Bau: Why, there are these like little bumps I don't know, but like something cause that.
David Bau: So just here's another just a zoomed out to be with that signal
David Bau: cool.
David Bau: So I don't think I usually don't have time to talk about it.
David Bau: but I showed you this, these components to 9.9, like the highest composing components with 9.9.
David Bau: And I I pointed out that there's a few here that are not in the Ioi circuit, but are still highly composing.
David Bau: So we do some investigation. We do some like kind of preliminary, like circuit analysis
David Bau: on these heads to see if they're anything interesting. And it does seem like they are
David Bau: part of a novel circuit that
David Bau: has something else to do with like context, retrieval similar to this inhibition mover, self circuit but has more to do with like
David Bau: retrieving, like semantically related items from context, which is so, it's a little. It's like, qualitatively like, very different from inhibition or
David Bau: mover had this land a little bit.
David Bau: So the the deals for that already in the appendix we didn't
David Bau: arrive at like a super conclusive thing. It would require like a whole circuit analysis, I think, to do that. And we just didn't. Didn't have the space. I think so.
David Bau: But yeah, some cool results there.
David Bau: So that takes us to like. Okay, this is nice.
David Bau: We, at least I I mentioned this like relative positional signal in one dimension. I think this ties back nicely to like Mkhil's result on the like low, rank signal controlling positional relative positional signal
David Bau: in the like the least task.
David Bau: So
David Bau: we only showed this for like 2 names, then the question was like, Okay, like, What is this? What are these heads doing in general like these specific components? What are they doing in general, like
David Bau: on arbitrary data like.
David Bau: what is scaling this up and down? Do like, what are they generally responsible for?
David Bau: Every basically just looked at like attention
David Bau: patterns on these heads on like a wide range of like just open
David Bau: like text corpora, and just looked for patterns. I specifically looked at like the top
David Bau: like the most active patterns when the model is when these heads are attending to like the non beginning of sequence, token, and the pattern that was like. Very, very consistent was that there's attending to words
David Bau: from like things like the word, and to previous items in a comma separated list or
David Bau: multiple choice questions. Kind of like anti repetition patterns.
David Bau: and things like that. So basically, the the kind of intuition here is that if these are in inhibition, it's more generally
David Bau: what these are doing is like
David Bau: attending to previous terms to avoid just like repeating them. So you don't want to say dopamine, serotonin, and dopamine. That wouldn't be good. So we kind of think that's what these are responsible for.
David Bau: and to kind of study that more we came up with like a kind of toy task which we call the laundry list task.
David Bau: and the idea is
David Bau: to predict it's a leave one out task. You give it a list of items.
David Bau: and then you kind of repeat the list in some naturalistic way, and then the model needs to predict the like remaining item that hasn't been repeated yet.
David Bau: And so it's like, you say, like, Okay, I'm I'm gonna go to the store. I need to buy a plate pair and a pen.
David Bau: 1st look at the pen.
David Bau: the pair, and then I'll get the black like the plate. Is the remaining item, and these are randomly shuffled each time. So the second order is
David Bau: is always different. So it's not just a you can't just use like an induction head to complete the pattern.
David Bau: And it's really easy to scale the number of projects you can have, like, you know, 20 items on your on your laundry list, or something.
David Bau: So
David Bau: so we we take a look at this task, and the model is pretty good at predicting the last item up to some number of
David Bau: items
David Bau: but now, since we kind of like know something about these individual components within the weights, we kind of do the same test that we already did, which is like scaling up and down
David Bau: the activations along these directions again, going back to the idea that there's like coordination happening between
David Bau: the the the component, the individual components across the inhibition, the multiple inhibition that we have.
David Bau: And so what we do is like, basically plot
David Bau: each inhibition head singular vector which corresponds to the important component
David Bau: on some plot, you know, like on this, like 3D space.
David Bau: like, I'm just like, move around the activation in this 3D. Space and see how that affects the
David Bau: move. Your heads attention to the the remaining, the items in the list.
David Bau: And what we find is that like, we get this kind of smooth interpolation between, like, you know, item, one, item, 2, item, 3 up to 10 items. So basically, this subspace, this 3D subspace is controlling.
David Bau: the attention to like the item in the list. And this is again a relative. It's not. It's not an absolute position, because
David Bau: these prompts can have, like varying number of tokens. So like the the prompt formats like
David Bau: are different every time. So they're they're varying by a few tokens
David Bau: and it's always moving the attention between the like list items. So like lighter orange in bottle, or like
David Bau: pair, pen and plate like things like that.
David Bau: Can I ask how you find the part in the 3 displays? I
David Bau: So we? I did a really large it. I have a figure that that might like visualize it nicely for you. But basically, like, I just swept this whole space like, you know, took different, like.
David Bau: like every every 10 or something. And I found the one spot that's like most attending to. Item 3, most attending to item one, item 2. And I just like, Drew the triangle. Then interpolated.
David Bau: yeah, so this is that that figure basically like this is each.is a point activation space that I tested, and the color is the most attended to object in that part of the space.
David Bau: And so this is with 5 objects. And you can see, it's like, actually, really nicely
David Bau: grouped where I this is all attention to. Item one item, 2, 3, 4 and 5 kind of like goes around
David Bau: as you add more items
David Bau: so this is now, this is the one I just showed you. And this is with 6 6 items, and it like splits it at a certain point, something that seems to happen. So people have talked about like primacy and recency bias in model, which is like preferring the 1st item or preferring the last item. Humans have this bias as well, which is maybe how the humans got it because they're trained on our data. But that's speculation. What seems to happen is
David Bau: item one. This space seems to like, never get broken up as you add more items, and this last, this space. So this orange space in this purple space.
David Bau: both cases respond correspond to the very last item in the list, regardless of like the number of items. And it doesn't seem to get broken up.
David Bau: basically until it like really needs to be. I guess what needs to be means here is kind of weird, but like
David Bau: it always comes at the time. Seems to like split things up in the in the middle items.
David Bau: And I'll show some more examples of that in a minute. But basically, like
David Bau: this plot is showing with Gpd, 2.
David Bau: The accuracy as you scale up the number of objects
David Bau: it decreases like very, very rapidly, even though the difficulty of the task is essentially the same.
David Bau: I mean, there are more distractors to be fair. But like
David Bau: you know, there's this attention. Mechanism is not like losing information. It always has perfect information of the whole sequence. So it's kind of
David Bau: not surprising that it gets so much worse as you increase the number of objects. But there's not really a good explanation for why
David Bau: and so one way we tried to explain this is that, well, okay, if you go up to like 20 objects like, here's item one again, it's taking up this huge portion of the space, and it's like splitting up like this little corner into, like, you know, 15 different objects. So the error, tolerance for attending, for the mover had to attend to.
David Bau: Any one item is like super low now, like you can only ask to like.
David Bau: you know, precision. Strike this one little like yellow spot, if it wants to attend to like item 14 or something so that's kind of what we think is is why
David Bau: one reason why they actually goes down.
David Bau: We're able to somewhat validate that with this like really manual intervention where we
David Bau: we hard code the if we know that the answer is going to be like. Item one. We hard code the attention to be in like this part of the subspace correspond to item one or like item 14 or something, and we can like pretty significantly improve the accuracy on this task across the number of items. We don't fully fix it. But, like for example, the item 10,
David Bau: there's 10 objects.
David Bau: Accuracy
David Bau: after the intervention is about, is actually higher than the original models, like 5 objects accuracy. So it improves quite a lot.
David Bau: The kind of summary here is that, like
David Bau: we find evidence that models are
David Bau: communicating through like these low, specifically low rank communication channels
David Bau: and an inhibition mover sub circuit example, which is just one that I was particularly interested in.
David Bau: It seems like this communication is like entirely content, independent. It's basically a mechanism. For, like acting like as an indexer in a list when the model makes the keys, the key vectors for the as in the list. It's
David Bau: assigning something like, Hey, I'm item one, I'm item 2 or something like that. And then the query is able to index those specific keys. And it's sort of like position, independent way.
David Bau: her absolute position independently.
David Bau: So yeah, so that's that's that paper.
David Bau: I have. Some ongoing work that I want to talk about, but I'm happy to take. If there are any like questions right now about this one
David Bau: have to answer that
David Bau: it's in earlier paper. I think they mentioned the best innovation. Heads also include some form of
David Bau: at least to some extent, the token value information as well.
David Bau: Yeah, how do you? What? How do you take that? How? How do you think about that?
David Bau: Yeah, right? So like, under like answering it, maybe only under the perspective of this I would maybe expect, like the token value. So the the, I think, is what you're referring to also finds there's like a, they find that it's like strongly passing a positional signal as well as a token value signal. Yeah, although predominantly positional. Yeah, and so if I had to
David Bau: guess if we if we wanted to like find the like
David Bau: token value signal, that's also going going through these heads. And maybe it's just through other components. And
David Bau: How large that subspace is that's storing like
David Bau: the token value, like the identity of that token. I'm not sure. I think that's like a cool one, cool like you were talking about doing like a binary mask like learning procedure. I think that would be a cool way to actually find that
David Bau: I remember I was talking to Sheridan about like
David Bau: the sub embedding space being in like again, induction had like moving token values around. So it's like, I think, like, still, a lot of work on that. I just don't. I don't know the exact answer, but that that'd be my guess is like. That's where I'd look at least, is like, try to find out
David Bau: some subset of the the components that are passing like token value.
David Bau: Okay.
David Bau: how important that is for the prediction I like. How the model decides what's important for a prediction I don't really know, though.
David Bau: between those 2 signals.
David Bau: just out of curiosity for the rank, one matrix that you find. Is it normally in like the top, like 3 singular values? Or can it be like much further down? That's a good question. There. It seems like there's a bias towards. Okay. So here's like the duplicate token ad example. This is like.
David Bau: so we index at 0. This is
David Bau: which 1, 1 and 2, yeah, one and 2. So the second and 3, rd which are high
David Bau: high variance, but not the 0, and then these ones are
David Bau: this is index one. I know you can't like. Read these, but this is index one. This is index 6. So like this, the 7th most important. This is index 2. And this is index one again. So it seems biased towards the beginning, but not necessarily like the 0, or something like the highest variance.
David Bau: So I don't really know if there's any rhyme or reason to that, it might just be.
David Bau: What are your thoughts on using the weighted term based methods.
David Bau: the blind section where you don't know the
David Bau: right? So good. Yeah. In a totally blindside my thoughts, for I'm gonna I. I'm gonna bring that up again. And I am also just curious to make other people's thoughts on that.
David Bau: My thought is that like it's a good way to reduce your search space in a completely blind setting.
David Bau: It would be great to be able to like.
David Bau: like reach some some solid, like
David Bau: interpretable insight just from without running the model.
David Bau: So the nice thing that you can find all the connections between all the ways this way,
David Bau: and like, particularly like strong connections like that you would actually be interested in.
David Bau: There are other things that you could like mix into the process without running the model, like, for example, doing like logic lens on
David Bau: rows of the weights, or even even the components themselves. Like the singular vectors, you can do larger lens on actually, people have done that, and they're highly interpretable.
David Bau: Tend to be
David Bau: so you could say something like, Oh, this head is attending is connecting really strongly with this specific component and also this specific component outputs like
David Bau: the names of mountains like really strongly in the output space, or something like that, you know whatever that means of. But like that, then you have some idea of what that circuit is already. So I'd be curious about like in like
David Bau: introduce, like introducing the like, the connections between the weights with other like weight based capability stuff.
David Bau: For example, there's other stuff on like you can find a mover head by looking at the diagonal, the ob matrix, if it's like really positive, or if it's really negative, and it's a negative mover had like demotes, whatever it attends to.
David Bau: So that's like.
David Bau: I think that's about as far as you could get in a totally blind setting. I I think there's probably like a lot of room to grow
David Bau: curious about. The other people have ideas on that later.
David Bau: So there's a you think that there's more work that needs to be done for validation of the circuit.
David Bau: Yeah, I think in general, like, you're definitely going to want to
David Bau: actually run them all at some point to validate your your findings right? But like
David Bau: there's maybe settings where, like, you
David Bau: have some feature vector like like from an essay or something. You know that this is like a feature you care about that you can look at without running the model. What components
David Bau: interact with this singular vector like, which ones are. If I were to steer this, what components would actually just like react.
David Bau: So that's another 2 stuff
David Bau: I'm gonna go like, really, quickly through this. There's not like, actually that much like I
David Bau: mainly just want to like. Talk about what the idea is, and like some things I care about studying. But there's not so much to get into. This is work that's being led by Alan Chan, who's a student at Brown Undergrad.
David Bau: And the idea is to like we want to study like universality of of model features, particularly across model sizes.
David Bau: And one thing that people
David Bau: have done since like 2015, or something, although I don't think people do it, that much, for whatever reason is, is model stitching, which is the idea that you learn a projection between the representation spaces of 2 models. So this can just be like the middle layer, residual stream of a 1 b model and a 7 b model, let's say, and you just learn a projection matrix, p, maybe plus a bias term. So it's like an affine transformation from one representation space to another.
David Bau: With the idea of that.
David Bau: you can run like the 1st half of the 1st model and the second half of the second model. And you get something that's like coherent like you get something
David Bau: maybe like the performance of like something in between the 2 of them, in terms of like your distribution of tokens or something right?
David Bau: And so we train this projection in this specific project, we train this projection layer with like an Msc. Loss like, I've just reduced the mean squared error between the 2 representations like the stitch and the
David Bau: ground truth. If you were to run like the 1st half of the 7 B model.
David Bau: Once you have this projection matrix. This is just mapping from
David Bau: hidden size A to in size. B,
David Bau: you know, there's a lot of stuff you can do with that by itself. But
David Bau: sparse autoencoders are a really nice way to decompose your residual stream.
David Bau: And so one thing we can do is just pass the weights of this sparse autoencoder through the projection matrix. So like the the encoder matrix rows like the decoder matrix rows or columns, or whatever you can just pass those through this projection layer. And now you have a sparse autoencoder
David Bau: for model B, basically without training it on model B at all.
David Bau: And so if this works even like a little bit, then you can say something about and like the the activations like, correlate to each other. Then you can say something like, Okay, dude, are learning really similar features, or maybe really dissimilar features.
David Bau: What we found is that generally you're don't worry about the difference between stitched and seed stitched. Seed stitched is kind of like a little bit like a skyline. It's like
David Bau: 2 random seats of the same model stitched into each other versus stitched is between a model, A and a model B, which are 2 different sizes.
David Bau: All these models are trained on the same data. By the way. That's also an important thing. I should mention.
David Bau: meaning the question we're kind of asking is like, specifically about features learned by models of different sizes. So that's something to keep in mind.
David Bau: Generally we find that the features correlate really, really strongly. So compared to like a random baseline, which is a very low baseline.
David Bau: we got, we got a pretty positive correlation between the activations of individual neurons
David Bau: in the Sae. This also lets us do it. Yeah. So your correlation is between you. You stitch an sae over through your stitch transformation, and then you're correlating them with what with an sae just trained in
David Bau: in the new setting. Is that right? So it's I glossed over that. So you have sae on model a which is trained on model A, and you can get some distribution of activations over just text for all the neurons in your sae, and then you stitch it into model B,
David Bau: and then run the same text through model B through your sparse auto encoder and get the distribution over neurons.
David Bau: Oh, so each each of this, each of the transferred
David Bau: directions is responding to the same
David Bau: text in in about the same way
David Bau: as it was in the original model. Yeah, yeah, I see.
David Bau: That's good. Yeah. I see
David Bau: and it's it's it's hard to really evaluate that. I think that's currently a a point of
David Bau: of it can be flushed out a bit more. And then for your random baseline, yeah, are you taking like an untrained sae and doing the same thing? Yeah, you can do a couple of different things, like random sae pass through a train stitch or a random like a trained sae through a random stitch.
David Bau: Yeah. So to me, the interesting thing would be like a random essay through the same train stitch. Yeah. And is that what you plotted here is that different? These are
David Bau: random. It's both. Actually so random essay is this orange 1? 0, I see
David Bau: through the same stage. Yeah. Question conceptually, because I'm not really as familiar like, what's the difference between this and a sparse cross. Coder. Yeah, we started working on this before cross coders came out. They are different, though, like one way. You could answer this question now is with probably with the cross. Coder although there is, there are kind of important differences.
David Bau: Where you are
David Bau: specifically optimizing for a shared feature space for the cross code. But they're they're conceptually pretty similar.
David Bau: Yeah, there's there's a couple of nice things that you can get out of
David Bau: this approach, though, which I'll talk about.
David Bau: So one thing you can do which we were
David Bau: really interested in the fact that this worked is that if you have a
David Bau: a vector that successfully steers model A to produce Spanish text, you can pass this through your stitching layer. Oh, okay, I'll answer that question in another way, actually, right now. But keep that in mind. But
David Bau: you have some vector
David Bau: which maybe you get from your sparse auto encoder which which steers Spanish text. You can pass that through your stitching layer. And now you also have a steering vector on model B, that also produces Spanish text.
David Bau: And so another part of this project I didn't mention. We we I've been kind of focusing as far as I don't coders.
David Bau: but this works for any any way that you can get a steering, vector
David Bau: you can pass it through this stitching layer or whatever feature vector, you want. So
David Bau: you know, you can.
David Bau: The the per. The point is that the stitches is able to like map these representation space as well. Not that if you had a
David Bau: a difference in mean steering. Vector, you could also pass that.
David Bau: But yeah, that's still something. We're kind of figuring out.
David Bau: So we're interested in that kind of result. We also found that if you wanted to have a sparse autoencoder on both models, you actually get a really large flop savings by doing it this way. We're still evaluating how good these ses actually are. So we're running on like sae bench and stuff. But basically you get
David Bau: if you initialize
David Bau: your sparse auto encoder and then continue training it. You get about a 2 times or about a
David Bau: 33% flop savings, training the stitch and fine tuning your initialized essay rather than training the essay on the Small model and the big model separately.
David Bau: Oh, fuck.
David Bau: which is maybe a nice way to to kind of like Bootstrap a bunch of essays like, very quickly.
David Bau: but still looking into that.
David Bau: sorry. I know I'm running over time. I'll kind of end here, the kind of like things I'm really interested about that that kind of fell out of this
David Bau: are these like features that don't map? Well,
David Bau: We call them like bookkeeping features, because we think they're like storing like metadata. A good example would be like the kind of indexing feature I just talked about like in
David Bau: in the talking heads paper.
David Bau: So we found a few that are like this.
David Bau: And they map well, to, yeah, sorry. What's the effective null space of the unembedding matrix? Yeah, I can explain it. I just don't want to go way over. But basically the idea is like
David Bau: from this paper, this paper? And also actually, there's a West Kearney paper about it, too, that also talked about as like
David Bau: the
David Bau: effective null space is, basically, you take the unembedding matrix. You do the Svd on that. And there is a long tail of singular values that are like basically 0. And so if you write into those directions, you're not affecting the logics at all. So
David Bau: if you are outputting features like into that space
David Bau: that's doing nothing for your next token prediction. And the idea is that well, they're probably doing something.
David Bau: This paper showed that they're controlling like confidence regulation
David Bau: specifically in like the neuron basis.
David Bau: so, for example, like changing the entropy of the next token prediction without changing the like ranking of the next token distribution. Basically.
David Bau: And we found that these features that are like high norm and like high density.
David Bau: are also like high cosine similarity with this effective null space
David Bau: and seems to be doing something kind of similar, like there's a mix of them that are
David Bau: confidence regulation. There's maybe some that are like deactivating attention heads.
David Bau: So
David Bau: you know, we found this through the essays. But you can also be finding the neurons. So like, really is like, there are better the kind of like kind of tangential question is like, is this better to look at the se features or the the neurons? I don't really know. I'm kind of like curious about that. But I'm also just curious about these features in general that exist.
David Bau: Anyways, I'll wrap up here instead of like rambling on. But back. It kind of like
David Bau: goes to like the kind of broader point of like.
David Bau: it's really nice to look at these features that if you were to actually like lost them look totally uninterpretable, but are, in fact, doing something.
David Bau: Essays can give us that. But maybe also neurons can, too, and like maybe weight decomposition can as well like with the indexing thing. I don't know like what we can get out of each of these methods. That's really different.
David Bau: It'd be really nice, I think, to
David Bau: figure out like the advantages, disadvantages of all of those, and
David Bau: like what kind of guarantee that we get us on, not others. So that's kind of, you know.
David Bau: leaving off on kind of a like
David Bau: shrug point like, this is kind of what I'm thinking about now?
David Bau: Yeah, that's it.
David Bau: It's amazing about your work. Yeah, yeah.
David Bau: are you gonna be sticking around this afternoon to chat with folks? Yeah. Cause I think that you know all this. All this stuff that we've talked about today is really
David Bau: relevant to a lot of the projects. And so I don't know if we have, do we have, like a spreadsheet to be organized about chatting with Jack? We can make one, and while we make it, Jack, can I steal you for a few minutes and just chat with you, and then maybe in the meantime, you guys can make a little plan for the rest of Jack's time here. How how long can you stick around? Oh, okay, great. Yeah. That's great. So yeah, maybe you can help. Thanks. Thanks a lot.
David Bau: Thank you. Thank you, Jackie.
David Bau: Bye.