David Bau: Who are more like.
David Bau: Oh, say, like senior in this room I am looking for like brutal, honest feedback, because I'm hoping to turn some version of this into adventure job talk. So
David Bau: that'd be helpful.
David Bau: Okay, so I'll get started. Thanks for the Intro. My name is Andrew. I'm a Postdoc at Harvard with Professors Martin Wanberg and Fernanda Vegas.
David Bau: And today I'll be talking about how we can reverse engineer like recent phenomena. We've been seeing in language models.
David Bau: So I'm gonna start with this comic.
David Bau: It's a bit of a lag.
David Bau: We're on the left. This man is talking about how AI can write emails for him and on the right this woman is talking about how the I can read and summarize your emails, the point being that language models are becoming ubiquitous.
David Bau: And on a more serious note. We're seeing this in multiple applications or domains like education, healthcare etc. And so as practitioners, it's a very exciting time.
David Bau: And as researchers, it's also a very exciting time, because we're seeing all sorts of interesting phenomena show up like paradigms, like alignment and jailbreaks, reasoning models.
David Bau: In context, learning and inference, time scaling.
David Bau: But the truth is, even as researchers, we don't fully understand why these phenomena happen.
David Bau: And
David Bau: on a more. And it's not that we're only seeing like mysterious behaviors show up. We're also seeing unintended behaviors, whether it's things like bias
David Bau: toxicity, hallucinations, models, taking shortcuts or hacking their rewards.
David Bau: And so, if we were to
David Bau: try to understand, like why we see either like mysterious behaviors that we don't understand or unintended behaviors.
David Bau: we can revisit how these models are being trained or built. So during training, we rely on the models inputs.
David Bau: We come up with some metric to monitor and like Monitor, this metric over time.
David Bau: and once we think this model is trained, we again rely on the model's outputs, benchmark it against some held out data set to evaluate this model.
David Bau: But what about all these intermediate computations?
David Bau: Unfortunately, we often turned a blind eye because we're dealing with a few challenges, such as the fact that we're dealing with high, dimensional features, complex nonlinear operations, and an intractable number of parameters.
David Bau: But my goal for today is to hopefully convince you that we can actually interpret a lot of what's going on in these intermediate steps.
David Bau: By kind of understanding the geometry of their weights and activations and whatnot
David Bau: and so today, I'm gonna talk about 2 of these paradigms
David Bau: and an example of how we can reverse engineer what these models are doing, and then I'll spend a bit of time talking about more about like the geometry of language models.
David Bau: So that's kind of the agenda for today.
David Bau: And so the 1st paper I'm going to present is called a mechanistic understanding of alignment algorithms. A case study
David Bau: with these amazing mentees and mentors.
David Bau: So, as you likely already know.
David Bau: Often these language models
David Bau: learn like undesirable behaviors from pre-training. And so they go through this post training process like rotf to kind of align their behavior to match human preferences.
David Bau: But it turns out like, even after this post processing step you can give these models, these adversarial prompts to jailbreak them and revert or, yeah, revert their behavior back to and like, demonstrate undesirable behaviors
David Bau: and Dp, right. So so like, this is a case where, like we have to understand the models, internals to understand why this might be as opposed to just relying on model outputs.
David Bau: And Dpo direct preference. Optimization is a popular choice of an Rohf.
David Bau: Up with them
David Bau: where it assumes you have, like pairwise preference data. So given the same prompts. W. It you assume
David Bau: you have, like a preferred continuation y plus and a non preferred continuation y minus
David Bau: and dpo works by promoting the likelihood of your positive samples from being generated while suppressing the likelihood of your negative samples.
David Bau: So in this work we were
David Bau: curious. We we ask, like, what exactly is happening to these concepts being quote unquote, unlearned by Dpo.
David Bau: And why is it that after you unlearn these concepts, these models can still be jailbroken.
David Bau: And so we take Dpo and toxicity as a case study. To understand these questions.
David Bau: and the procedure to answer
David Bau: these questions is actually pretty straightforward, which is to 1st understand how toxicity is represented
David Bau: before alignment, and your pre-trained models in the paper
David Bau: we look at Gpt. 2 and Llama 2. But all the results here in the slides will be with Gpt. 2.
David Bau: The next step is to unlearn toxicity with Dpo, and then the 3rd step is then to inspect how the representations for toxicity has changed.
David Bau: But 1st I want to give some very quick background information. In particular, I'm going to talk about the mop box.
David Bau: and many of you might have already seen
David Bau: this. Because I know, David, I think you have a paper on this. But like typically, you might have thought of Mops as like having a linear projection with Wk. Followed by element, wise non-linearities, followed by another linear projection back down.
David Bau: But instead, I want us to think of these operations. Row wise and column wise.
David Bau: And so I'll start referring to each row of the 1st linear projection as key vectors.
David Bau: And then I'll refer to the columns of the second linear projection as value vectors.
David Bau: And you can think of the forward pass as taking some input vector X
David Bau: and doing a dot product of X with all of your key vectors.
David Bau: followed by applying some nonlinear non-linearity.
David Bau: And this, whatever resulting value you have here gets multiplied to a corresponding value. Vector and it's so that your value vectors are being scaled, either larger or smaller, depending on these earlier steps.
David Bau: the point being that, like, you can think of the output of your mop block as some weighted sum of your value. Vectors, where the weights are again determined by these earlier steps.
David Bau: Any questions so far?
David Bau: Okay, so yeah. So let's talk about some findings, the 1st one being
David Bau: understanding how toxicity is represented in pre-trained models.
David Bau: and our analysis will start by training a linear probe model. I'm sure a lot of people here are familiar with this concept, but the idea is to basically pass in a bunch of input sentences that have been
David Bau: labeled as either toxic or non-toxic.
David Bau: And during the forward pass we're going to grab the activations from the last layer and try to learn a linear probe model that I'll call W
David Bau: to classify whether this hidden state belongs to a toxic input sentence or a non-toxic sentence.
David Bau: And it turns out that you can learn a like probe model with high, pretty high accuracy.
David Bau: So we believe, like this
David Bau: means, like the language model, has a linear direction to represent toxicity.
David Bau: Pretty standard, I think, for this audience.
David Bau: The next step we do here is we can take all of our value vectors from all Mlp blocks
David Bau: across all layers. So like in the case of 52, there are like 40 96 value vectors per
David Bau: layer. And you have, like, say, like 24 layers.
David Bau: So we grab all these value vectors. And we basically just sort them, based on how similar they are to this W toxic fraction that we found from cobine.
David Bau: Does this make sense so far?
David Bau: Right? And so, because, like the transformer architecture, is additive with the residual connection, I kind of think of.
David Bau: please toxic value vectors as weights that contribute the most towards this toxic direction that we found
David Bau: and it turns out that if you take some of these toxic value vectors and quote unquote unembed them like, if you project them back to the models, vocabulary space, and look at their nearest neighbors. You see very toxic words show up often with a coherent theme. Yeah, are these the activations of the actual weights? Actual weights? Yeah.
David Bau: so we identify some value vectors that seem to encode toxicity.
David Bau: And so this gives us a hint as to where
David Bau: and how toxicity is represented. But we can take a step further with causal interventions. Again, I'm sure you are familiar with this idea. But the idea is during the forward pass we take some of these vectors that we believe are representing toxicity, and then we'll simply subtract these vectors from the activations
David Bau: to see kind of what causal role they play.
David Bau: and so given prompts like, I didn't give a blank, or she's full of blank or fully blank.
David Bau: Gpt. 2 would have originally outputted a offensive word.
David Bau: Whereas if you do this causal linear intervention instead, you get something like I didn't give a thought. I didn't give a suggestion, or she's full of it. She's full of trust or holy cow. Holy Grail! Holy!
David Bau: So
David Bau: these are like the resulting toxicity scores before and after our interventions. Given 1,200 prompts that are known to trigger toxicity.
David Bau: So like W. Toxic is the linear probe.
David Bau: the second one is one of the toxic value vectors.
David Bau: What I'm missing here is a baseline. But yeah.
David Bau: like. Maybe we could have subtracted some random vendor or something for a baseline. But
David Bau: anyways. So then, let's take. So now that we understand how toxicity is represented in our language model, we can apply Dpo to unlearn toxicity.
David Bau: I'm gonna skip over the details here. But one thing I'll note is like Dpo requires pairwise preference data meaning like given for each of your training samples. You need pairs of continuations that share the same prompt
David Bau: which
David Bau: like. At the time of this writing, we didn't have such a data set. And so we built our own. And the way we did that was we basically grabbed a bunch of random Wikipedia sentences.
David Bau: And we let a language model generate a continuation.
David Bau: so we had 2 runs like a normal generation, and a second run, where we used an approach called Pplm to purposely generate like toxic continuations.
David Bau: So this gives us a pairwise preference data that we can use to for Dpo and unlearned toxicity
David Bau: and doing so results in this quote, unquote, aligned model with a much lower toxicity score.
David Bau: So now that we have a aligned model that's unlearned toxicity, we can revisit how the representations for toxicity has changed after this Ppo step
David Bau: and previously we've like localized how toxicity is represented in these Mlp. Value vectors, and so we can see how those weights have changed.
David Bau: and to our surprise we actually find that those weights don't change at all, so you can compare the weights before and after Dpo with like cosine similarity scores, we see that they have a cosine similarity score of one.
David Bau: and you can also look at their norms, and they don't change at all, either.
David Bau: so that's kind of surprising. But we do find that during the forward pass the activations of these value, vectors do drop
David Bau: which means, like, basically the model, the ability for the model to elicit
David Bau: toxic outputs is still sitting within the model. It's just that they're no longer being activated.
David Bau: So we can do a deep dive into why, this is
David Bau: and in order to do do so, I need to explain what I've
David Bau: what I call an activation region of a key or key value vector pair.
David Bau: which is basically just a set of vectors that will trigger a value. Vector
David Bau: so imagine, like, if the input to this mop block was within this quote, unquote activation region, it would have a high enough dot product to trigger the nonlinearity and therefore have some positive output and scale the corresponding value. Vector
David Bau: whereas the input vector is outside of this activation region. It won't have a high enough dot product. And it won't trigger the non-linearity. And so your value vector won't be activated and won't get scaled and won't be outputted.
David Bau: So this is kind of a 2D depiction of the same idea.
David Bau: and what we see is like the hidden state at some layer where toxicity is encoded. Is original. Oh, right? So let's pretend like in this 2D depiction, the second key and value vectors are
David Bau: responsible for toxicity.
David Bau: What we see before Dpo is that the hidden state at some layer is within this activation region, and therefore triggering toxic value vectors.
David Bau: whereas after Dpo we see the slight shift of the hidden state outside of this activation region. And therefore, like no longer triggering toxic value vectors.
David Bau: And so there's there's this small shift introduced that I'll call Delta in the residual stream.
David Bau: And so this is kind of a empirical visualization of what I just said.
David Bau: and there's a lot going on here, so I'll break it down for us.
David Bau: So here, like every point in this graph, is a hidden state sample from either the pre-trained Gpt model or Dpo's model.
David Bau: where the circles are from Gpt. 2. Before Dpo. The triangles are after Dpo.
David Bau: and the dotted line indicates pairs of hidden states that have been given the same prompt.
David Bau: and all of these hidden states have been projected down into 2 components. The 1st is the main principal component of all of these points.
David Bau: and then the X-axis direction is the mean, like
David Bau: direction in which all of these hidden states have been shifted in. So it's like the delta direction that I showed you in my 2D picture.
David Bau: And then, lastly, the colors indicate how strongly each hidden state is activating toxicity like the toxic value vectors.
David Bau: And really the takeaway here that I want to emphasize is that as all of these hidden states are being shifted in this like X-axis direction, you see toxicities dropping so like all the red points, turn either yellow or green, and likewise all the yellow points turn green.
David Bau: So that was like very low, level technical details. Let so let me
David Bau: take a step back and summarize what I've said so far.
David Bau: What I'm saying so far is that the weights for toxicity don't change from Dpo.
David Bau: but we do see their activations drop. And this drop is because we see the hidden States being shifted to avoid regions. That
David Bau: trigger toxicity. Yeah, so maybe I missed. How did you actually get the Delta?
David Bau: Right? So you can simply just pass in the same prompts to the Pre train model and the Dpo model. Just record the activation. And then
David Bau: any other questions.
David Bau: So the toxic weights don't change. So the toxic weights were the value vectors.
David Bau: So your key vectors also do they change? Yeah, in in general, like most weights don't change like their norms. I forget the exact number, but
David Bau: like there are tiny changes, for, like all weights.
David Bau: including key vectors, value vectors, attention head rates, and so on.
David Bau: And I think people have been seeing similar observations and just
David Bau: a lot of post training 10 years. As well.
David Bau: Yeah, that basically, the changes are very diffuse. Yeah, diffused
David Bau: minimal. But somehow, like very strategic, if that makes sense.
David Bau: And I always like show you an example of that here as well.
David Bau: Yeah, just another thing. Where do you find these like toxic value? Vectors are they? In the later layers?
David Bau: They're like scattered like, the one that seems most
David Bau: most like causally relevant was like at layer 19, I think. Like, right? So like we're we're looking at in this case. Gpt, 2 medium. So there's 24 layers, the earliest I think I saw was like Layer 12 ish, but mostly around like layer. 1819 20, ish.
Nikhil Prakash: Oh, I have a question.
David Bau: Yeah.
Nikhil Prakash: So what is the thing that changed the most? If the value actors did not change that much, what was the thing that changed the most.
David Bau: Yeah, let me go through a few more slides and then revisit.
Nikhil Prakash: Yeah.
David Bau: Yep.
David Bau: And
David Bau: so there's this notion of sometimes outlier values within the weights. Did you notice anything outlier in the weights or activation, or both, both, I'm working on weights. So I saw, I see. So like sometimes, if I if IZ squared normalize, there are a few, maybe 5%, 2% of the actual weight entries within the matrices
David Bau: which can be like 100. So that's like an outlier, which is not probable, I think. Since the universe began, or something like that, if it's a Gaussian distribution.
David Bau: I haven't seen any in this work, but I have heard that sometimes, you see, like, really abnormal
David Bau: activations. I didn't realize you saw the same in weights. But yeah, yeah, I'll send weights as well. It's definitely interesting. Why, that might happen?
David Bau: Okay? Any other questions?
David Bau: And the thing is that, have you seen these toxic vectors getting activated in non toxic scenarios like, are they only encoding toxic vectors, or are they like encoding other information as well?
David Bau: Oh,
David Bau: I don't think so. I I haven't seen such cases. It's it's possible, especially if, like.
David Bau: if we consider, say, like superposition, it might be that
David Bau: say, like some linear combination of these toxic value. Vectors also encode some other concept, perhaps. But we didn't look that far.
David Bau: Yeah, if if you look at like some of the unembedded.
David Bau: like nearest nearest tokens in them on embedding space for these vectors.
David Bau: I want to say, like, roughly, like after, like 10 ish really toxic words. In the beginning you start seeing
David Bau: other.
David Bau: it seems random words show up. Yeah. So I don't know if those random words somehow form some other pattern. But
David Bau: yeah.
Can Rager: I was wondering on the previous slide you chose to plot the principal component on a y-axis. What do we learn from the principal component here?
David Bau: It's mostly like here. I'm mostly just trying to emphasize that like
David Bau: like, there's there's a single direction as indicated by the X-axis, where, like, if the hidden states
David Bau: move towards this direction like you see activations dropping.
David Bau: and this I don't think this direction necessarily
David Bau: matters. It's more so that
Can Rager: And it's not a principal component.
David Bau: No, no, no.
Can Rager: Yeah, okay.
David Bau: It was just a convenient way to all these questions.
David Bau: Yep.
David Bau: okay. So right.
David Bau: weights are barely moving. But we do see the hidden state bypassing these regions that trigger toxicity.
David Bau: So we can go even deeper. Well, then, where this? Where did this delta? Where did this shift come from?
David Bau: And to enter this we can look at even earlier layers.
David Bau: And then, in particular, we can look at again the value vectors from earlier layers and see how they've been shifted.
David Bau: and what we see is
David Bau: like, I said before, these weights are also being shifted ever so slightly, and
David Bau: such that, like the models overall end behavior doesn't get affected very much.
David Bau: But these small, tiny shifts.
David Bau: like across weights across layers, do accumulate during the forward pass, such that by the time you reach, say, like layer 19
David Bau: these small shifts accumulate to introduce that delta that we've been talking about.
David Bau: and it it allows the model to avoid these regions.
David Bau: And one hypothesis that I have is that
David Bau: what I just described is kind of perhaps a effect of kale divergence, regularization, where the goal of a lot of these post training methods is to keep as much of the preaching behavior as possible. And perhaps the model is doing this by introducing minimal shifts per weight.
David Bau: But these minimal shifts accumulate in a strategic way such that like it can, for instance, in this case reduce toxic outputs and like collect its rewards.
David Bau: So I know there's a question on zoom about like which weights shift the most.
David Bau: I think, if I remember correctly like it might have been like the unembedding weights. But even those were like very minimal.
David Bau: like always shifts were very minimal. If I, if I
David Bau: remember correctly the one that did ship the most were, I think, the unembedding rates.
David Bau: Does that answer your question?
David Bau: Talk soon?
David Bau: I have a question. Yeah.
David Bau: have you tried just holding all the way? It's constant and just manually manually shrinking the
David Bau: the toxic value vector? Weights. And what happens when you do that? Yeah,
David Bau: we have. Let me revisit that question in a few slides.
David Bau: Oh.
David Bau: okay.
David Bau: there's this slide that I never get to share, because it's too like low level. But I think this audience will appreciate it.
David Bau: And so there's a lot going on here, and I'll walk through it again. So for now just only pay attention to the blue curves.
David Bau: and what I'm showing in the blue curves is like per layered
David Bau: like each layer has, like 44,000 ish value vectors. I'm comparing the cosine similarity. Scores of how these value vectors shifted versus
David Bau: the shift at, say, like layer 19 in this residual stream.
David Bau: So the cosine similarities were between this pink arrow and all the 4,000 pink arrows in the earlier layers.
David Bau: So that's what the X-axis is showing the cosine similarity scores. And on the Y-axis, I'm showing, like the percentage of those 4,000 vectors with a corresponding cosine similarity score.
David Bau: So this is basically saying, like as we get closer to say, like layer 19
David Bau: most of the value vectors are being actually shifted in the opposite direction of the shift we saw in
David Bau: the hidden state the Delta vector that I become.
David Bau: Does that make sense at just the blue curves?
David Bau: Can you please repeat that. Yeah, okay, so
David Bau: let's like, only focus on this layer. 18 subsplot.
David Bau: What the blue. So like, what I'm doing is at layer 18,
David Bau: we have 4,000 value. Vectors that all got shifted ever so slightly.
David Bau: Take, like those shifts.
David Bau: Take the cosine similarity of those versus
David Bau: the shift at Layer 19 in the hidden state
David Bau: like, only the shifts. Yeah, okay, just the shifts
David Bau: and compare it against the value vectors at layer 19,
David Bau: comparing how the value vectors shifted. Okay? And and there at earlier layers.
David Bau: it's it's the shift in value vector weights versus the shift in hidden states. So it's both shifts, just one at like layer 18 with the weights and the layer. 19. Yeah, or it's both, yeah shifts. Yeah.
David Bau: does that make sense? So it's like one to many. Yeah, and you're plotting. The histogram is like all of the Yup Yup. So then the Y-axis is indicating the percentage of those 4,000 shifts, the value vectors that have a corresponding score. So this is saying, like most of the shifts seen in the value, vectors
David Bau: are shifted in the opposite direction of like Clear Nineteen's hidden state.
David Bau: Does that make sense?
David Bau: And we see this?
David Bau: Like, the closer you get to layer 19 where say, like, toxicity is encoded. You see this trend.
David Bau: and if you looked at like earlier layers like layer 17, or whatever like you see the same trend
David Bau: when you say layer 17, do you look at layer 18 or so later?
David Bau: So like if if you plotted
David Bau: the shift in value, vectors from like layer one all the way to 16, and then compared to hidden shift in layer 17.
David Bau: You kind of see the same trend.
David Bau: Okay. So like the value vectors are being shifted in the opposite direction of like the
David Bau: ratio streams shift, and like that seems to
David Bau: something must be going on. This seems
David Bau: kind of suspicious. Just be a coincidence, and so
David Bau: to understand why they're being shifted in the opposite direction, we can
David Bau: take 2 things into consideration. One is that a lot of times the activations of these value vectors are.
David Bau: they tend to be very sparse.
David Bau: meaning most value. Vectors are typically inactive.
David Bau: and you have to also take into consideration what activation function is being used in the model. In the case of Gv 2, it's using a glu function which has this corresponding shape.
David Bau: which means that value vectors that are inactive end up taking a very small negative
David Bau: scalar value or coefficient, which means, during the forward pass
David Bau: all of these value. Vectors are being
David Bau: multiplied by some small negative value, meaning they end up flipping directions and then end up contributing towards the
David Bau: Delta direction that we've seen at layer 19.
David Bau: That makes sense.
David Bau: And so that's what this orange area is showing. It's showing the average activation of all 4,000
David Bau: value vectors. And you do indeed see that most of the time they take on a very small negative value.
David Bau: Oh, wait! So this is so blue is
David Bau: pairing the shifts and orange. It's just showing
David Bau: the contribution of the value vectors
David Bau: to the residual stream. It's mostly to indicate that they typically do and take on the negative value. So they're mostly turning things off. Yeah. But when they are turned off because they get multiplied by a negative value, they like
David Bau: end up flipping directions and contribute towards that top
David Bau: delta value. So in this, in this picture, like, I kind of purposely
David Bau: had these pink arrows pointing in the opposite direction.
David Bau: like these value vectors and earlier layers
David Bau: because they're typically inactive, they get multiplied by some negative value, flip directions, and then they accumulate during the forecast to build this
David Bau: shift at Layer 19.
David Bau: Does that make sense? I don't get it.
Chris Wendler: That was really hard to follow, but I.
David Bau: Okay, is the orange, the jelly activation? Yeah, yeah.
David Bau: yeah. So they are typically negative. Oh, wait. So wait. So what's the difference between orange and blue. Are they very closely related? Are they measuring totally different things? They're measuring totally different things. The blue is indicating the cosine similarity score of
David Bau: the shifts.
David Bau: Yes, in in of the of the
David Bau: Mlp. Outputs versus the residual stream. Shift Yup right? And then. So so this amazing negative one thing is very cool where the change in the Mlp. Output is almost exactly the opposite of the change in the residual stream for these late layers.
David Bau: and the reason it's almost exactly opposite is because
David Bau: typically most value. Vectors are unquote, inactive. But when they're inactive,
David Bau: when they're turned off, Geelu assigns them a very small negative activations for
David Bau: and because of this negative coefficient these
David Bau: shifts in value vectors that were in the opposite direction retire stream they're like flipping directions during the forward pass.
David Bau: and then
David Bau: so like. Now what was cosine, smoothly, score of negative one end up became coming close to positive one.
David Bau: See that I that I don't get
David Bau: right, I mean, so it's like the the residual stream is
David Bau: the sum of a bunch of jelly outputs. Also.
David Bau: right? I mean, if there's a i'm not sure why
David Bau: the flip would affect one and not the other.
David Bau: Do you know what I mean?
David Bau: Not quite sorry. Could you repeat that so your residual stream output.
David Bau: So you're you're arguing that has to do with the
David Bau: like. I'm not convinced, like like I would. What I would be more convinced by is, if you
David Bau: you know, if somebody had a big LM
David Bau: that was trained with relu instead of jelly, which doesn't have this inversion, and this effect, like goes away like, I'm like, I'm like, yeah, I mean, yeah.
David Bau: cause it feels like it. It. It feels more likely to be something different than this? Right? Because, like the
David Bau: yeah, you have this slate, you know, sign in version.
David Bau: But you know all of these slightly net. You know negative
David Bau: domain things, you know. Presumably they're all over the place. They're also contributing to the residual stream.
David Bau: Right? So the residual stream is this, you know, represents the sum of a lot of this stuff which has a sign in version. I'm not sure why the sign in version would only affect.
David Bau: you know, the individual Mlp. Outputs and not affect the residual stream in a similar way.
David Bau: So let me try to put it differently.
David Bau: so. But let me. But maybe I maybe I don't understand what the orange thing is that so that what you graphed in the orange, right? So I understand what you graphed in the blue thing. The point is, they're all all the blue, like all the shifts they look like in their exact opposite direction. But actually, they're being multiplied by negative value. So they're in the exact direction of the Delta. Right? That's the point.
David Bau: So like, they're like the reason why. But the value vector so wait. The value vector, is.
David Bau: So there's a, the jello is before the value vector, usually right. Yup, and so
David Bau: so it would be waiting. All these guys that are in the opposite direction of the Delta.
David Bau: But then, when it's waiting them, it's actually negative value, right? So it's flipping them so that now you're adding the pod the opposite direction
David Bau: to the residual stream.
David Bau: Oh, I see. So your argument so so it's the Mlp. Weights.
David Bau: not the Mlp. Outputs that you're measuring.
David Bau: Yes, I see. So that's your argument. I get it.
David Bau: Good.
David Bau: Yeah, good.
David Bau: I think I think I get it. Okay. And and what is it? What are you measuring in the orange? So the orange is the like average activation of
David Bau: the value vectors.
David Bau: So oh, so the orange is what I thought that you had. So the orange is the
David Bau: the Mlp. Outputs.
David Bau: It's like the Geelu.
David Bau: Wait right before
David Bau: So let me go to this slide.
David Bau: So in in this
David Bau: diagram, like orange, is basically this scalar value. It's the it's how much each of the value vectors are being scaled by. Okay.
David Bau: And I'm saying, typically, they're like very small negative values.
David Bau: And so, like all all of the value vectors that are inactive end up flipping directions.
David Bau: Okay, yep. And so here
David Bau: Like all of the the value vectors are being like. Let's go. Let's look at it the other way around. Like all of the value vectors are being shifted in some direction.
David Bau: But all these shifts
David Bau: during the forward pass, even even though the weight has like moved, say, this way
David Bau: during the forward pass. Their contribution to the ratio stream is actually going to be in the opposite direction, right? Because they're all off. Yeah, exactly. And when they're off, they get multiplied negative value. And so then eventually, when all these small shifts and activations accumulate, it's what we see at, like some later later layer where it allows the model to bypass. Okay.
David Bau: does that make sense sort of does? Okay? I get it. Now, okay, I see. I see that you mean, okay, yeah. Okay.
David Bau: yeah. No. Chris raised his hand earlier and also has a couple of comments on the chat. If Chris, you want to chat. Talk about it now.
Chris Wendler: Oh, no, it's resourced. I just wanted to say that it makes sense.
Chris Wendler: like maybe maybe to the plot you can add, like the product as well. That would make it easier to read.
Chris Wendler: to add like a 3.rd
David Bau: Income.
Chris Wendler: Or maybe you start with the product, and then you decompose it into this plot.
David Bau: Yeah, there's a reason why I always skip this slide.
David Bau: But yeah, we we can revisit this later if there are still parts that are unclear. Okay?
David Bau: Okay. So last experiment from this paper.
David Bau: And this is where I'll answer your question hopefully. So
David Bau: once we have this aligned model to
David Bau: check our understanding. What we can do is we can scale some of these key or value vectors to undo unalignment.
David Bau: In particular, if we take the key vectors corresponding to
David Bau: toxic value vectors. We can scale these key vectors larger, which in turn will also scale their activation regions larger because remember, the activation region is defined by like the dot product with each key. Vector
David Bau: and we can. So we can scale these key vectors larger, such that even after the shift in the models
David Bau: like hidden state. This hidden state is now once again within this activation region, and triggering our toxic value. Vectors
David Bau: And so in the paper. We show that as by scaling as few as 7 of these
David Bau: key vectors, we can revert the model back to its toxic behavior.
David Bau: So to answer your question. We can do this the other way around, where we
David Bau: play with a value vector like scale them smaller.
David Bau: But what we found was like working with the key. Vectors is more effective because if you scale these value vectors either
David Bau: too small, too large, like they end up affecting the model and kind of unexpected ways, especially when you combine like 7
David Bau: puppies.
David Bau: Value vectors.
David Bau: because, like the value vectors, I suppose I think of as writing directly to the register stream, whereas the key vectors, you're just deciding what to read. That makes sense.
David Bau: yeah. So
David Bau: both. Probably there's probably some sweet spot of like, how you much you carefully scale these value vectors. But we found key vectors to be stakehold.
David Bau: Okay, so that's like the 1st project where we show by really
David Bau: understanding the representations of some concept like toxicity, we can understand, like, what exactly these alignment algorithms are doing, what they're doing to these unlearned concepts, and why things like jailbreaks can happen.
David Bau: And since the publication of this work numerous researchers have suggested other ways of doing alignment to address this type of issue.
David Bau: Any questions before I move on?
David Bau: Yeah. So maybe this is covered in one of the the citations there. But
David Bau: is, I guess, is your intuition that we should be doing Dpo on like the frozen value vectors, for example, and that would prevent jailbirds. That's a
David Bau: I think that's 1 way like. So some of what some of these approaches do is like they do 2 passes where in the 1st pass they identify weights that are most relevant for some concept. And then they freeze all weights except for those
David Bau: okay, those weights that seem to matter most and only update those weights.
David Bau: One thing that I wanted to try, or that I tried very briefly, is to augment the say dpo loss with some sparsity term, so that the model is encouraged to
David Bau: update as few parameters as possible, but update those aggressively.
David Bau: But I couldn't get them out of the train like it was too unstable when I tried that. And
David Bau: maybe if I maybe it will work. But I've only tried it for like a year or 2.
David Bau: Cool. Thank you? Yeah, okay.
David Bau: I think. Chris resistant. Yep.
Chris Wendler: Thank you.
Chris Wendler: Have you also looked at like language models that are using something like a giglu where you have like these 2 branches
Chris Wendler: in the thing that would, so you wouldn't just have like the keys, as in your Mlp. But you would have like
Chris Wendler: a separate set of.
David Bau: Yes. So this next project we'll talk about that.
Chris Wendler: Alright, cool.
David Bau: Yep.
David Bau: And so I'll get into that. Actually, Chris is on this paper. Okay. So
David Bau: you've probably been seen like these reason models and the way they work is it generates some chain of thought, sequence?
David Bau: and so there's been
David Bau: questions on like whether these chains of thoughts are faithful meaning. Do these chains of thought actually reflect what the
David Bau: like? The computations the model is doing internally, and
David Bau: especially as as we give these models harder problems.
David Bau: it might be useful that these chain of thoughts sequences are faithful so that we can kind of understand how it's doing reasoning.
David Bau: But apparently researchers have found that these chains of thoughts are often not faithful.
David Bau: And so in this work we were
David Bau: what motivated this work was to understand, to like study if we can.
David Bau: instead of like monitoring these chains of thought, if we can monitor their inner computations and latent space instead.
David Bau: And so this is again a problem where you have to understand the models, internals, and not just rely on the models, outputs.
David Bau: And so to kind of scope our work, we picked a very specific
David Bau: skill. So like reasoning entails all sorts of skills. But we, we focused on a very specific one, which is self verification
David Bau: and we study like, how the model does self verification?
David Bau: Again. So some very quick background information. And this is where I'll talk about gated linear units which is kind of the same thing.
David Bau: Except
David Bau: it's kind of the same thing in that gated linear units are also just some weighted sum of its value. Vectors, it's just that. Now, what determines the scale of each value vector, goes through some additional steps. So now, there are 2 key vectors.
David Bau: I guess K. Gate and K up.
David Bau: I'll just probably still refer to these as key vectors.
David Bau: So that's 1 difference. And like previously I was describing what I call an activation region. It's a subspace that triggers a value vector
David Bau: And so similarly, your gated linear units will also have some activation region characterized by these 2 key vectors.
David Bau: And it's perhaps a more like complicated subspace. So I might just refer to these as polytopes
David Bau: but other than those differences like they're kind of the same thing.
David Bau: I'm going to skip over these slides mostly because I think this audience is familiar with this idea of Qk circuits and ob circuits where the Qk circuits decide what information to read from the residual stream ob circuits decide what to write to the residual stream.
David Bau: Okay. So the setup that we came up with is, we're looking for a task
David Bau: that requires reasoning and self verification. And this task, called countdown meets both criteria. Where?
David Bau: What the task is.
David Bau: We basically give the model set of operands, for instance, 2014, and 40,
David Bau: and we ask the model to come up with some arithmetic combination that reaches a target number in this case, 28
David Bau: and
David Bau: so so we we train a language model. We train plan 2.5 3 b using like deep seek r, 1 recipe, which is basically to
David Bau: just simply apply grpo with 2 rewards. One is you reward the model if it finds, if it reaches the correct solution as well as another reward for following some specified format.
David Bau: In this case, like the format is to produce its chain of thought between these, think tokens, and provide the final answer between these answer tokens.
David Bau: One thing that we take advantage of is what's called node collapse.
David Bau: So people have noticed that with preference tuning
David Bau: Often these models run into mode class by which they mean these models become overly like focused.
David Bau: I don't want to say,
David Bau: yeah, these these models become too like focused on just satisfying what preference
David Bau: it needs to satisfy according to it's like Rl loss, such that even at the cost of diversity it will always generate things that satisfy the preference in your
David Bau: either data or a reward model, which I think, perhaps, is why
David Bau: most language models sound the same these days.
David Bau: And we leverage this as a feature and not a bug. And so I think this next slide will actually make more sense where
David Bau: given a training sample like this.
David Bau: after a few gradient steps. The model will learn the task, so it will generate some long chain of thought where it will try to try and attempt and say it will verify its answer. This didn't work. Let's try something else until it eventually reaches the right answer.
David Bau: But even after this point, if you train for longer, for, say, like 300 gradient steps, you run into mode collapse, where the model converges to always generating this very systematic, structured chain of thought
David Bau: all the time.
David Bau: So that's what I mean by boat clubs, and of course, with like.
David Bau: you might have experienced this where, like with Mechinterp, it's very tedious to try to get this formatting right to know where to pro for and etc. And so we take advantage of
David Bau: one claps here, and so like in your projects. Like as you're looking for appropriate setup.
David Bau: keep in mind that mold collapse might be a suitable option for you to make your lives easier.
David Bau: okay, so now that we have a model that's doing like both reasoning and self verification, we can look at how it's doing self verification. And so I'll provide a top, down and bottom up analysis and show how they might meet in the middle.
David Bau: So the top down analysis is very similar to what we did before same recipe. We're going to start with a linear probe from the last layer.
David Bau: In particular, we're going to take the hidden States for all the time steps where the model has made an attempt at a solution.
David Bau: We're gonna take the residual streams at those time steps and train this linear vector and
David Bau: with with this probe. This, I guess the shape of the probe is like 2 by d, where it's mapping from a d dimensional space to
David Bau: 2 outputs, and so I'll refer to the direction, for like trying to classify that the attempt is valid as w valid, and the opposite opposite direction as w invalid.
David Bau: But once we train this probe like
David Bau: typical, like typical thing like, you can steer the model with this linear probe, where, even given some attempt where the model has not found the solution. In the in this left side.
David Bau: You can like add this vector to make the model think that it has found a solution. And so we'll say, Oh, yes, we found the answer and output
David Bau: the attempt, even though it's not correct.
David Bau: Well, what we do here is again same thing as the toxicity setup. Where, given this w valid direction.
David Bau: we look for all value vectors once again. That is most similar to our wal, direction.
David Bau: and so kind of the same thing. Well, then, what happens if you unembed these probe directions or value vector, directions. And you see these
David Bau: words related to verification show up.
David Bau: So in the case of W valid, you see where it's like. Exactly yes, equal signs show up.
David Bau: Interestingly, we see this like pound, pound, pound, pound, sign, show up any guesses as to why or what that is.
David Bau: So my guess is like, if you look at the training data for Gsm, 8, K. The training data will have some chain of thought, followed by pound, pound, pound, pound, and then final answer, and I think that's kind of why, that might be. Why, this is showing up
David Bau: here.
David Bau: Same thing for w invalid, you see where it's related to
David Bau: verification. For some reason we see, like Chinese words show up, though.
David Bau: which translate to unfinished, unable, and etc.
David Bau: same thing with the value. Vectors. You cannot embed them. See words like is success accurately, and so on.
David Bau: Same thing with invalid.
David Bau: You see you again. For some reason, Chinese words for that translate to difficult, can't reach doesn't exist.
David Bau: And so on.
David Bau: So
David Bau: yeah, have you checked their full logic distribution like, are they systematically promoting these concepts? Or does the logic distribution kind of like seem uniform?
David Bau: like? Why, I'm asking that is that even if you see small differences in the logic distribution, it is, gonna get expanded when you project it with Softmax, of course. Right? Yeah. So
David Bau: and also like, if something has the right shape.
David Bau: you can, of course, project it in the popular space. So you'll always end up seeing something, and because of how
David Bau: word embeddings are trained, you are going to see similar things clustered together on the top.
David Bau: I'm not sure I understood the second half, but we haven't looked at the logic distributions.
David Bau: I understand how like Softmax, can
David Bau: explode like some, the likelihood of some token from being outputted.
David Bau: Like, even if the value is like, let's say, 2 values like
David Bau: with logic values like 20.1 versus 19.9. If you project it with exponential like geometry, it is going to be the 20.1 is going to get exploded. But if all the or most of the values in the vocabulary get a get assigned a distribution between 22 and 18, then maybe the model does not really have a very strong preference for those concepts. Only
David Bau: does that make sense?
David Bau: you're saying, if the logic distribution is somewhat uniform. Yeah.
David Bau: Sorry I didn't understand the question like,
David Bau: okay, maybe I'm not asking it. Right? Okay?
David Bau: So
David Bau: So you're not asking about the logic distribution given some chain of thought.
David Bau: I I'm asking about the
David Bau: kind of like this approach in general to localize a concept in certain value. Vectors?
David Bau: I think that like if it has the right shape.
David Bau: you can, of course, project it to something, it will say something right.
David Bau: but if you do not see, like systematic
David Bau: promotion of certain concepts and systematic suppression of other things. You cannot be really sure that this value vector is actually doing something that is very specific, like promoting this concept.
David Bau: And that is why I'm asking that if the distribution looks kind of uniform.
David Bau: then this this is kind of like a little bit counterintuitive.
David Bau: I think I understand better. We we haven't looked at the logic distribution.
David Bau: But like that's kind of where you can do like to understand
David Bau: what role these value vectors play, we can do causal interventions.
David Bau: So like, for instance, here, this is exactly what we do. Yeah. So here we we give the model like 300 samples test cases.
David Bau: And we basically
David Bau: like right after the model has found a solution where it needs to do self verification, we turn off a set of
David Bau: Tlg valid vectors Valley vectors
David Bau: in particular, like we turn off like 50 per layer, or those 50 are selected, just based on how similar they are to the W
David Bau: valid direction, the pro direction. Yeah, yeah, I understand that part. But my so sorry for taking some time my criticism, for that it will be that
David Bau: it will always work this causal intervention. Because because they're specifically turning off specifically turning on or off specific words which
David Bau: maybe accidentally get decoded in the space. And since it is an editing update, yeah? And of course, in later layers. If you add something in Layer 15, it is going going to get exploded. I see what you mean now. I guess 2 things. One is like that, this exact point where we intervene. Yeah, like, it's, it's
David Bau: it will either like predict, like not 28, or this works.
David Bau: but like we purposely intervene right when an attempt is made where the next character being predicted is the open parentheses, and not like
David Bau: answer or incorrect, that makes sense. But so that's at the last token position, right? We intervene at the last token position. But it's not directly predicting yes or no at this point it's predicting, like the open parentheses.
David Bau: That's like one thing. The other thing is like, i i i see what you mean, and I think often
David Bau: like this is kind of going to be hand wavy, but like, if you imagine the token embeddings
David Bau: like we see like clusters like in the case of toxicity, we saw like cuss words. In this case we saw words related to verification.
David Bau: and my like working hypothesis is that, like the model, build some like
David Bau: higher level representation of some concept.
David Bau: But it just happens to be that when they're projected to the embedding space, you see, like relevant tokens.
David Bau: so like to to
David Bau: like a good way to test your hypothesis would be like in a lot of the toxic. In the case of toxicity. We saw a lot of
David Bau: value. Vectors that were similar to the toxic direction, but not necessarily
David Bau: near, like cuss words in the unembedding space. Okay? So maybe if we only intervened on these non-toxic sorry non, like cuss word value vectors, and still saw a drop in like
David Bau: toxic outputs. Maybe that's 1 way to control for our customers. Or maybe another thing could be that if you ablate those toxic vectors altogether and
David Bau: try to make the model say something toxic.
David Bau: Like, somehow, then the model can still find workarounds to say the toxic words
David Bau: like. By by ablating, I mean, 0 ablated like, not like some negative. Yeah.
David Bau: Yeah. I think ultimately, what
David Bau: is a good question here is like, what is the relationship between like hidden State representations and the
David Bau: embeddings that happen to be near them.
David Bau: Like, I think there's some interesting relationship there like, why is it that if you
David Bau: unembed these concept representations you see tokens related to that concept show up. I think that's
David Bau: ultimately like, what's what might be the most interesting question here.
David Bau: because that is how you located them. It is like,
David Bau: yeah, it is. It happened that like often, it works when you find these patterns. Yeah,
David Bau: so like empirically, we're just. I'm just like reusing this method over and over again. But there's still an unknown there like, why? Well, for the toxic case. Yeah, that's how you found that right. But for for this reasoning cases
David Bau: is, it's not how you found it right, is it?
David Bau: Your your probe is not
David Bau: being. You're not probing at the point when it's saying got it, or didn't get it, or something like that. Right? You're probing at a different point. We we are probing
David Bau: like right when a attempt is made. So at at. In this case we're probing like
David Bau: at these squared typesteps. Right? So you're not probing at the place where it would say.
David Bau: You know, yes or no, you got probing at a different place, right? So this is, I mean, it is a little bit. It's surprising. Right? Yeah, yeah, the the vocabulary.
David Bau: It it kind of is about to say yes or no, but it says it in a special way
David Bau: like this. Say not, you know, this works, but it's like, Oh.
David Bau: in case there's any confusion like the parentheses, not in the text. It is. This is what the language model is saying, oh, it is. Yeah. Yeah. Oh, I see. But what I want to say is like.
David Bau: right after it says, for instance, negative 10. It doesn't predict.
David Bau: not. It has predict open parentheses, and then oh, but it is, it's pretty close. It's getting ready to say no, it's getting ready to say no, but it's not literally saying the token. It has already decided that what it wants to say. But it is saying that in a specific format, yeah, yeah.
David Bau: Okay, okay, cool. Yeah. Thank you. I mean, I could probably let you finish it. Yeah, I thought about like what you've asked a lot. And like in the last paper, I think we'll revisit this idea of how concepts and token the relationship between like concept representations and their token. There's tokens.
David Bau: Okay? Okay, so we've identified these glu value vectors that we think are relevant for
David Bau: self verification. We can do the same thing. We can turn these off during right before it tries to do self verification.
David Bau: And there's like 3 possible outcomes that we see here. One is, the intervention fails, and so the model will still correctly verify its output and output. This works and output the answer.
David Bau: In the case. When the intervention worked, we can make the model fail at self verification and say, Oh, this is not equal to 28. So let's try some other themes, and it will indefinitely continue its chain of thought forever
David Bau: interesting. We see sometimes a case of partial success where we initially are able to disable self verification.
David Bau: But then the model self corrects itself and says, Wait, what we just tried is actually cracked. And so this is the answer.
David Bau: I guess, just to
David Bau: when we intervene, we are only intervening at a specific at that specific typeset, and we don't intervene at any other time step
David Bau: when it corrects itself? Does it always correct itself right away? Or does it ever correct itself? Some number of steps later.
David Bau: usually almost right after, I think by the time it's like moves on, it's it never comes back together.
David Bau: So you end up seeing actually a very low success rate here.
David Bau: What's kind of funny here is
David Bau: here. We're just grabbing the activation, grabbing the value vectors that are most similar to this w valid probe direction.
David Bau: but you can also grab a few value vectors that
David Bau: are most similar to your like w invalid direction, and also turn those off as well.
David Bau: and like intuitively, you might think the like invalid direction shouldn't matter when it's trying to do some verification. And so why, you see, this improvement is again like you have to think of the geometry of these vectors.
David Bau: where, if you take some of your value, vectors for valid, and flip them, and then unembed them. You see.
David Bau: the opposite semantics show up. So you see words like wrong invalid, bad failure, and so on.
David Bau: Same thing with the invalid value directors. They originally
David Bau: encode things like failed not enough, but if you flip their directions and unembed it, you see words like successful, smooth, finally equal, sufficient so on and so on.
David Bau: So like the same idea here.
David Bau: even though this, like w invalid direction, seems to encode
David Bau: like invalid tokens like during the forward pass.
David Bau: because they're inactive, they get multiplied by some negative value, end up pointing in the valid direction, and then contribute for self verification.
David Bau: So this is kind of why we see this bump in intervention, success rate.
David Bau: But there's still this other like 50% to explain. For.
Chris Wendler: Andrew.
David Bau: Yeah.
Chris Wendler: Does it ever happen that? You. So? Because in the get glue architecture you could also multiply them with a negative number, right?
Chris Wendler: So you put them through your activation function and they would be on the positive side. But, like the second branch
Chris Wendler: theoretically, could output a a negative number right?
David Bau: Second branch being.
Chris Wendler: So like you have 2 key.
David Bau: Okay?
David Bau: Oh, yeah. That's a good point.
Chris Wendler: Has this ever happened?
Chris Wendler: Yeah, I'll do
Chris Wendler: imagine that this may give rise to the value, vectors being meaningful into both directions as well.
David Bau: Yeah, that's a good point. I should check the actual activations and see that they are negative.
David Bau: I haven't looked at that
David Bau: good point.
David Bau: So that's kind of the top down analysis. We started by probing at the last layer, looking for weights that are similar to it. Now I'll show you a bottom up analysis.
David Bau: The thing about this task is that the task actually specifies what the correct answer is in the context.
David Bau: And so a reasonable approach for the transformer model. To check its solution is to simply attend to this time, step and compare like.
David Bau: would it just output it against
David Bau: like the answer provided in the context.
David Bau: And so we basically look for attention heads that spend at least 10% of its attention to the answer token.
David Bau: And I'll just start referring to these as previous token heads.
David Bau: And this isn't like new people have found previous token heads before they seem to play a role in induction heads, and so on.
David Bau: And so in this, in the case of quent a. 3 B.
David Bau: There are, I suppose, 33 previous token heads.
David Bau: and it turned out that if you turn off these 33 attention heads, you get perfect intervention, success rates.
David Bau: And so then the model will always fail at self verification and generate its chain of thought forever.
David Bau: One thing that you might wonder is whether there's a relationship between these value vectors that we found before, and these attention heads these previous token heads.
David Bau: And my claim is that these 2 sets of weights share the same, quote unquote verification, subspace.
David Bau: So what I what do I mean by this?
David Bau: So this is where we'll revisit, like the Qk. And ov circuits where ov determines what information to write. Qk. Determines what information to read
David Bau: and people have suggested what they call like
David Bau: interlayer communication channels where they come up with this composition score to decide which attention heads are talking to each other.
David Bau: So
David Bau: here w. 1 might refer to an ov circuit of one head, where W. 2 might be the Qk circuit of another head, and the more these 2 values are aligned.
David Bau: the more it might be that these 2 attention heads are talking to each other.
David Bau: So what whatever like the Ob circuit is writing might be meant for the Qk circuit to read
David Bau: so analogous, we we do something similar, except instead of the Qk circuit of another attention head, we're going to replace that with the key. Vectors of Glus.
David Bau: So like this is our quote, unquote composition score that we come up with.
David Bau: and I'll walk us through what we're looking at so previously. If you look at how the value vectors are activated, it relies on 2 vectors like K Gate and K up
David Bau: And here, instead of the input vector X, we're just going to
David Bau: all right. So we're like looking at a specific pair of key and value vectors. But here, instead of the some input vector X which is going to be the model's hidden state. We're simply going to replace that with the ov circuit of your attention heads.
David Bau: And then.
David Bau: like this is this score here is for just a single pair of key value vectors. But given a set of glu valid vectors, we're just going to take the average of score of those value vectors.
David Bau: does this make sense.
David Bau: We're basically checking how much an Ov circuit aligns with the key vectors of a Gov vector
David Bau: and an analogous way to think of this is like again, I've talked about how these
David Bau: value vectors have an activation region. And so again, we can think of
David Bau: these key vectors shaping some polytope. That is an activation region for a value. Vector, and if you think of just the intersection of all these activation regions for our glu valid weights.
David Bau: we're basically checking how much the ob circuit of an attention head a lot quote unquote aligns with this
David Bau: set of activation regions that I'm calling a verification subspace up, pop the formula back up again for the score, I see.
David Bau: I see. So you're you're aligning the
David Bau: the the parameters, the gate parameters of your glu with the parameters. Yep, this will tell us which
David Bau: attention heads activate the glu guide vectors the most.
David Bau: and when we do this we like sort our 33 previous token heads based on this score. And we, like a plate incrementally update one at a time.
David Bau: And it turns out when you do this by turning off as few as 3 attention heads.
David Bau: you can disable self verification.
David Bau: Does that make sense any questions so far
David Bau: sorry? What was the last thing that you just said
David Bau: right? So like, given the 3 3 previous token heads.
David Bau: We're going to try
David Bau: instead of like turning off all 33. We're going to try turning off just one
David Bau: see how well that works. Try turning off 2 try to turn off 3, and so on. And when you do this, it turns out that
David Bau: after you turn up 3 heads you can disable self verification. Okay.
David Bau: and
David Bau: so now here, I'm showing you the activation of the Gl view. Valid vectors. So each corresponding like neighboring pair of blue and orange bars, indicate the activation of
David Bau: a specific glu value vector before and after turning off the 3 attention hits.
David Bau: And so you can see that by turning off the 3 heads you can basically
David Bau: deactivate all of your glu value glu valid value vectors.
David Bau: So my claim is like in this case, the
David Bau: 3 attention heads takes the model into this quote, unquote verification subspace because it thinks it's found solution. And the same quote, unquote verification subspace
David Bau: activates your Glv weights which will promote the likelihood of tokens like success or correct from being generated.
David Bau: Last part for this is we then look at what these components do in the base model.
David Bau: So just client, 2.5 3 b before Rrl. Itery.
David Bau: But of course the base model isn't trained to do countdown
David Bau: but what we can do is we can provide the countdown task in the form of in context, learning
David Bau: with the same like structure that we like.
David Bau: And so if you give.
David Bau: I think we use like 5 demonstrations, then it turns out. The model can solve countdown using the same exact chain of thought sequence.
David Bau: So this like adds evidence to like
David Bau: parallel work that, like it, seems like Rl. Is just enhancing existing mechanisms. I know. I think Nikhil has a paper on that
David Bau: of
David Bau: right. So in this Icl version of the setup and the base model. We try like turning off these
David Bau: same components.
David Bau: and it turns out like, if you turn off the attention or previous token heads, all 33 of them, you still see very high success rates
David Bau: out of range here means like the model generates something that doesn't Co, that doesn't adhere to the chain of thought sequence that we want it to.
David Bau: Interestingly, if you turn off only the 3 verification heads, we see a much lower success rate.
David Bau: and we haven't like explored this further. But my hypothesis is
David Bau: with the Rl. Step. It's really enhancing these 3 heads to do the self verification task.
David Bau: which made it easier for us to localize.
David Bau: Are these the same heads? Yeah. Okay, same exact heads.
David Bau: Awesome.
David Bau: now. So we looked at plan 3 B pre trained and rl, fine tuned. Now we're looking at a general reasoning model. Deepseq, R. 1 14 B.
David Bau: And we do the same exact procedure like training. We actually don't train the pro. But
David Bau: I'll talk about how we can
David Bau: repeat this procedure. But the point is, we can apply the same procedure to find glu eights, attention heads, and so on, and turn those off.
David Bau: So in this case I think there are.
David Bau: I want to say, 2,000 attention heads.
David Bau: But if you turn off
David Bau: the either previous token heads or verification heads based on, like the same criteria that I described before, you still see very high success rates.
David Bau: But interestingly now.
David Bau: in the majority of cases we see, like a partial success where the model will initially say, Oh, yeah, 26 is not equal to 26. But then, right after, say, wait a second. It actually is 26 and self correct itself. So it seems like it has a much more
David Bau: press, like additional mechanisms for self correction, in addition to like self verification.
David Bau: So by simplifying, by by studying a simplified setting for reasoning. We do find this subspace that we think is relevant for self-verification, which gives us hope that perhaps we can monitor the latent space of these models to understand what they're doing for these reasoning tasks
David Bau: any questions before I try to quickly go through the last project.
David Bau: So if if Rl. Is just enhancing capabilities that are already, or ex encouraging the model to express capabilities that are laden in the base model.
David Bau: But then our one is seemingly significantly better at recovering after it's
David Bau: after we intervene on it. Why, why do you think that is? Is it just just more thorough? Rl to express more latent base model capabilities. Or do you think our one learned something
David Bau: from? I think it's just like it has a.
David Bau: It has like more parameters like we're comparing like 3 B models versus 14 B, and I think that's just where that additional capability comes from.
David Bau: That's my guess. Sure.
David Bau: Okay,
David Bau: so hopefully, I've convinced you that like, if we understand the geometry of these model weights and activations, we can understand their inner computations.
David Bau: And then so in this left
David Bau: project I'll be talking about like similar geometries we see across language models.
David Bau: So this is work with some amazing mentors. This is under review for Cullen right now.
David Bau: And in this work we're only going to be talking about the embedding space
David Bau: of language models, the embedding and unembedding weights.
David Bau: And what we're going to study is different geometric properties of these embedding spaces and compare them across
David Bau: language models in the same family.
David Bau: The reason why we do. This comparison within the same family is because we need the embedding points to match across embedding spaces.
David Bau: which basically means, like, we need to compare models that share the same tokenizer. And so we're just looking at models that share the same tokenizer.
David Bau: And we compare the quote unquote, global and local geometry of these embeddings.
David Bau: Starting with global, what we can do is given a set of points in the embedding space, we can compute a distance matrix. So given this some point, for, say the word he like how similar is it to the word she chair, desk, and so on.
David Bau: using cosine similarity scores. And so you construct this distance. Matrix d, that's like
David Bau: V, by V, the number of vocabs by number of vocabs. And then all you have to do is compare how similar the distance. So I suppose I think of this D as capturing the relative orientations of words amongst themselves.
David Bau: and so you can simply compare how similar these distance matrices are across language models using Pearson correlation scores.
David Bau: And so this is what it looks like for Gpt. 2. Where you, when you compare Gpt. 2 small medium, large and extra extra large.
David Bau: You see pretty high correlation scores indicating that the relative orientations of these token embeddings are actually quite similar across language models.
David Bau: Does this make sense so far?
David Bau: And if you look at the other 3
David Bau: other 2 model families, you actually see a very different picture for Llama 3.
David Bau: Any guesses as to why this is
David Bau: so? It turns out like in the case of Gpt. 2. And Gemma. It quote unquote ties its embeddings, which means it shares the same weights for both embedding and unembedding, whereas the case of Lama 3. It it doesn't tie.
David Bau: It's embedding weights. So it learns separate weights for the unembedding part.
David Bau: except for like one B and 3 b, 1 b and 3 B share, one B and 3 B tie there
David Bau: embedding on embedding weights.
David Bau: So in the case of Lama 3, if you do the same thing for the unembedding vectors, you see very high correlation, very high light edge.
David Bau: So this tells us that the kind of global, like relative orientation of token embeddings across language models are similar.
David Bau: Now we can. We can look at the local geometry.
David Bau: and we can characterize local geometry 2 ways. One is by using locally linear embeddings, loes, and the other is by defining and looking at intrinsic dimensions.
David Bau: So what are Lls linearly local embeddings? The goal of Ll. Is to take some point.
David Bau: take K of its nearest neighbors, and to express this point as though, as some linear combination of its K nearest neighbors.
David Bau: So here, if we pretend like the, if, if we
David Bau: care to reconstruct this blue. Vector. We can pretend like, let's pretend, like the 3 other vectors are its 3 nearest neighbors. We want to find
David Bau: a set of weights.
David Bau: That can be applied to these 3 neighboring embeddings to reconstruct this blue. Vector so here,
David Bau: so in particular, this W vector here is going to be of shape. R, by B,
David Bau: except this, this, vector, is going to be very sparse, because, except for the K nearest neighbors, all other entries are going to be zeros.
David Bau: And you can imagine trying to learn this W vector, for all of our points.
David Bau: I try to just re-express every point based on
David Bau: some waited some of his nearest neighbors.
David Bau: So that's the goal to learn this. W matrix.
David Bau: Does that make sense?
David Bau: And so you can think of this as trying to optimize this objective where
David Bau: we're trying to reconstruct every point based on its nearest neighbors, where W.
David Bau: Each entry of W is going to be 0. If it is not the nearest neighbor of the point we're trying to reconstruct, and, like the rows, have to sum up to one.
David Bau: and it turns out. There's a close form solution for this which makes our lives easier.
David Bau: But once we have this W matrix, then we can compare how similar these W. Eights are across language models.
David Bau: And so if you look at llama
David Bau: like one B, 3 B, 8 B, 70 B, and compute these Ll rates and compare them. You end up with very high cosine similarity scores
David Bau: for most of the case.
David Bau: What's interesting here is this big spike in cosine similarities where it's 0
David Bau: any guesses as to what that might be.
David Bau: So we think these are.
David Bau: quote unquote under trained tokens. Some people call these glitch tokens, which are tokens that the tokenizer has
David Bau: like dedicated. It's created a token for these random tokens, but they don't show up in your training data. So then, like they end up the embeddings for these, which tokens being just like randomly initialized vectors, which
David Bau: in like high dimensional space, would have equality score of 0.
David Bau: And so I don't think I have.
David Bau: Like we make this claim. But one of the reviewers asked for evidence of this, and so we look for overlap and like
David Bau: which tokens that some prior work has found and I forget the number. But there is some overlap in
David Bau: our tokens and their tokens.
David Bau: Okay, the other way to characterize local geometry is by defining this intrinsic dimension measure. Where what we do is, we take a
David Bau: token embedding. We look at its k nearest neighbors, and we apply Pci on it. We just count up the number of principal components needed to explain 95% of the variance in that like local
David Bau: neighborhood, and so that the number of principal components there would be our intrinsic dimension.
David Bau: And, interestingly, if you take tokens, compute their intrinsic dimensions like the tokens with
David Bau: very low intrinsic dimension, often form very coherent semantic clusters, whereas tokens with higher intrinsic dimension.
David Bau: You don't see that semantic clustering
David Bau: so like for instance, if you look at the token for, like I don't know, Nissan, it has a intrinsic dimension of 5, 99
David Bau: out of, I think this is Gpt. 2. So like out of 1024 mentions
David Bau: but if you look at its nearest neighbors, you see words like Mazda Hyundai, and so on. Same thing with like East numbers, police, and so on. But as you go up higher in tokens, with
David Bau: higher intrinsic dimensions, you see, words like
David Bau: this is just from like eyeballing some of the data. But
David Bau: you see words that seem to be either like suffixes or prefixes.
David Bau: or different lexical patterns, like words that start start with g
David Bau: like here. I think these are suffixes that start with 8, and so on.
David Bau: Right? So if if you like, computed the intrinsic dimension of all of your embeddings.
David Bau: and then compare them across language models with peerson correlation. You again see some correlation.
David Bau: Sorry you are. You're calculating the intrinsic dimension of a token. Can you please go back to the previous slide and like, just recap, how are you doing that? So we take.
David Bau: Say, we want to measure the inference dimension of the word heat. Yeah. We grab the embedding look at. We grab K of its nearest neighbors for K equals.
David Bau: I think we played with like 100 300.
David Bau: And then with those 300 points, you apply Pca.
David Bau: and then you just count up the number of principal components needed to explain like 95%, or say 99%, some threshold value of the variance.
David Bau: Okay? So and it turns out like without needing all like, say, 1,024
David Bau: dimensions after like, say, like 500 ish dimensions, you can explain 95% of the variance. So we're calling that the intrinsic dimension 200 nearest neighbors. Right?
David Bau: So technically, if I just do QR. Factorization, I should not get more than 300 basis directions.
David Bau: how can I get 500?
David Bau: Does that make sense?
David Bau: Yeah, yeah, yeah, so, okay, so like the numbers here. Sorry. So the number the Ids here should shown we grabbed
David Bau: like a little over a thousand 24 neighbors because of that reason.
David Bau: Okay, like the actual, like, absolute value of these ids don't matter much, but if we instead did something like 300 or a much smaller number
David Bau: like these Ids change accordingly. To
David Bau: some smaller value like, 200. Ish, say, or.
David Bau: Yeah, okay, yeah, that was a good question. Okay, yeah. So K, is actually 1,000. Sorry in this table. K is a thousand.
David Bau: We are asked like, why? A 1,000. And so we tried 300 500. So.
David Bau: but we started with a thousand because of that reason, because otherwise. Yeah, yeah. So there's always topology. People who, you know, have ways of measuring, you know, intrinsic dimension. Have you compared your method?
David Bau: We haven't, I know, like Qnn is another way.
David Bau: we haven't. But I suspect we would see similar results.
David Bau: Yeah, I'm starting to do some reading in that space to learn of cool these things. Yeah.
David Bau: okay, how much time do I have? Last bit.
David Bau: So one thing, that kind of beat
David Bau: given that these language models share both global and local geometry is that you can simply transfer your steering vectors from one model to another.
David Bau: So imagine you have 2 language models.
David Bau: Imagine you have some steering. Vector. B, what you can do is
David Bau: you can learn a simple linear mapping across your unembedding spaces.
David Bau: Call that T, and you can use like even something as simple as like minimizing these squares.
David Bau: And you can actually apply this same linear transform on your staying vector V and apply it on your other language model, and it will still work.
David Bau: So here are some results.
David Bau: So here, what we're looking at is some refusal steering. Vector
David Bau: and the dashed blue line is if you steered Lama 3, 8 B
David Bau: with a steering vector derived from Lama 3, 8 B from the same language model.
David Bau: But if you take that steering vector from Lama 3, 8 B, and just simply transfer it to Lama 3 B or one B, and do the steering, you see the same effect.
David Bau: So like this multiplier is the alpha value. For like, how strongly you wanna scale your string vectors
David Bau: and then like more results, for like different steering vectors.
David Bau: and if we were to think about like, why this works, this is where, like your point about the logic space.
David Bau: it's relevant where, like, really, ultimately, you can think of the next token prediction task as this bilinear model where you have the last layer hidden state and you're unembedding space.
David Bau: You're, I think of it as just taking the dot product of the last layer with all your committees. I'm looking for the one that I start products. But
David Bau: right? So like you said, if you simply add a direction for the word he then, of course, the likelihood of he goes up
David Bau: and I think this is why this steering transfer of steering vector works because it turns out like, if you look at some of the steering vectors.
David Bau: So here, if you looked at the 1st 3 whoops.
David Bau: so the 1st line is a steering vector for the for the 3 B llama model. And if you unembed it. You see words like straightforward simple show up same thing with the 8 B model. You see words related to my pick show up, and after this transfer from 8 B model to 3 B. Because the unembedding space is very similar, you still see very similar words show up
David Bau: same thing with like the opposite direction.
David Bau: And like for many of these concepts, you see this pattern.
David Bau: The point so like
David Bau: this says, like the unembedding spaces, are similar, but I think this also means, like the hidden States, are similar to like. There's a very close kind of symmetric relationship between your
David Bau: last few layers at the on embedding space.
David Bau: So I don't have like evidence of this, but I think the hidden States
David Bau: end up becoming very similar as well.
David Bau: And so there's a
David Bau: this is like my last slide, like, why do these language models learn very similar representations? This probably has to do with the data. I wonder if there's a relationship with neurocollapse?
David Bau: Where?
David Bau: I I don't know if people are familiar with this concept. But like people have found for things like image classification
David Bau: like the representations for different classes, collapses into like single points.
David Bau: Which I kind of think of as just like the model having found an optimal representation to do this classification task. For, like optimal like, say, separability or whatever.
David Bau: And so, perhaps, like, you can think of language models as also doing classification context. And perhaps
David Bau: this is what your costs looks like. Where, like, they all converge to the optimal representation. To do this
David Bau: context classification task.
David Bau: So then, the last hidden states and this unembedding space end up looking very similar. That's 1 hypothesis.
David Bau: Yeah. So that's that's all I have.
David Bau: I'll just conclude with like, even though these in these intercomputations seem scary. Like, if we understand the geometry of these embeddings and activations, we can reverse engineer even complex models.
David Bau: Yeah, that's all I have. Thank you for your.
David Bau: So I'm kind of curious if you've like done some more like qualitative analysis of like what lies in different regions of that like similarity histogram that you showed the beginning of part 3,
David Bau: like, for example, like do function, words like end up in like very different parts than like the I guess different kinds of content words this
David Bau: more like the yes.
David Bau: So
David Bau: I think the spec at 0 makes sense here. But I wonder if, like this program basically just tracks like word frequency in the models trading set.
David Bau: Yeah.
David Bau: I think another way to put, it is like whether this tracks, how well a token has been treated. Perhaps.
David Bau: I do think my hope is.
David Bau: if we know what an ideal gentlemanly is, we can kind of use that as a metric, for how well
David Bau: model is trained at a token level.
David Bau: So so like, if we know that like.
David Bau: for instance. So okay, in a different example, like, if we know what governs the intrinsic dimension of the tokens, maybe it's a property of the data.
David Bau: Then we know what like during training, we can observe how these intrinsic dimensions evolve for these tokens. And we can hopefully identify okay like this token
David Bau: has seen enough context, and it's like fully trained. But some of these tokens need more data.
David Bau: Oh, because it hasn't reached its ideal transit division yet.
David Bau: So I'm I'm I've been thinking about how we can use
David Bau: geometry to either, for like tracking, either for evaluation. I know we've recently talked about that yesterday or for like training, monitoring, training.
David Bau: I don't know if I answer your question.
David Bau: Your question was whether this just reflects frequency, yeah, or if there's like, maybe some like structured pattern like structured. But that's also a really cool idea. I guess that seems sort of equivalent in a way to like just sampling training examples that are kind of adversarial, or that like have very high perplexity to the model, I would think, but at more like the token level as opposed to like a full example. Level.
David Bau: So sorry. What was the last part like your idea about like I guess. Monitoring like the intrinsic dimensionality of certain tokens to look for things that are maybe under trained. It feels a lot like just sampling adversarial examples for training in a way.
David Bau: adversaryo meaning like, like, they receive very high perplexity. Looking at the language model.
David Bau: yeah, potentially.
David Bau: I was thinking. And yeah, yeah, I was thinking in terms of like identifying which tokens need more data.
David Bau: In terms of identifying like
David Bau: different patterns.
David Bau: And like, whether this tracks like frequency data,
David Bau: not sure. But I do think there's some meaningful like features represented just in the bending space itself. So, for instance, like here.
David Bau: like, if you run saes, people be like. Oh, we found a feature for like words concept that triggers for like words that start with G or whatever like. I wonder if that's just reading
David Bau: directly from this a feature that's already in the embedding space.
David Bau: I guess this is to say, like I'm interested in like better understanding.
David Bau: like different patterns or structures that show up here.
David Bau: Maybe it has to do with word frequency. But I haven't looked at it yet. If you want to explore this space? Like, yeah, relatedly, have you considered or given much thought to more principled ways of selecting? KI would expect that different subspaces want different K, values. Yeah, yeah, that's a good point.
David Bau: so I think in the case of 2 nn, which is a different way of measuring intrinsic dimension, you only need to.
David Bau: It only relies on 2 of its nearest neighbors. I don't totally understand what it does from there, but like that might be a better way to capture intrinsic dimension.
David Bau: Cool.
David Bau: Yeah.
David Bau: I'm very new to actually like I've been claiming like, Oh, we have to understand the geometry of these models. But I'm very new to it.
David Bau: Geometry and topology. Actually. So, yeah, Prince has a question.
Chris Wendler: And like I think, like this work is amazing.
Chris Wendler: And I I wouldn't, you know.
Chris Wendler: short, sell yourself, I mean, your geometry is just very applied, and, you know, attached to the measurements which for me makes it great.
Chris Wendler: because I can actually understand what you're doing.
Chris Wendler: And
Chris Wendler: so I was wondering whether you have any thoughts on extending this to maybe analyzing the geometry of words, or more like anything that goes beyond single tokens.
Chris Wendler: because that would be.
David Bau: Yeah, yeah, if we can analyze like phrases
David Bau: and the hidden states, I think that'd be billing.
David Bau: I just don't know how to control. For like context, like like, what kind of
David Bau: context we should give it, whether it should just be like a whole bunch of random Wikipedia sentences. Or
David Bau: wow, maybe we can control. For like
David Bau: the grammar of the context using like our street or whatever
David Bau: yeah, I I think there's work to be done there. I haven't thought too much about what that might look like.
Chris Wendler: Alright. Yeah, if you want to chat about this, I would be happy to chat.
David Bau: So a very minor question.
David Bau: So so throughout his presentation, you kept calling the Activation region Polytope. Why, Polytope like? Why not a sphere.
David Bau: In the case of glu's like, I don't. I actually don't know what
David Bau: it looks like. So if we look if we like, pull up the
David Bau: yeah, I I don't know what
David Bau: shapes this results in. Like what the the subspace that satisfies this constraint.
David Bau: oh, because you're cutting it because of Jellu. Then it becomes a polytope. Okay, there's a jelly. There's also like 2 key key vectors here. I don't know how to characterize what kind of shape results in that. Oh, but 2 key vectors. But the activating region is ultimately one
David Bau: dot product with the value. Vectors. Sorry? No, no. No. Oh, okay. Okay. Okay, okay, no one dot product with the key vectors.
David Bau: there are 2 dot products here, one with K up, one with K. Gate, and there is a like a channel. Wise multiplication.
David Bau: It's hard to think about.
David Bau: And what happens if you just like, multiply 2 things in different representations.
David Bau: Yeah, okay, yeah, maybe there's a name for this shape. I don't know what it is.
David Bau: Cool. Thank you.
David Bau: Thank you so much. Yeah, thank you. Thank you. Thank you.
David Bau: So let's let's take a little break.
David Bau: Oh, did we not record? Did the recording stop?
David Bau: No, it's recording.
David Bau: Oh, how do I stop it here
David Bau: like a stop so backpack.