[0:00:01] In our last video, we found that we could describe two-layer attention-only transformers with a kind of handy equation. And the equation is three kinds of terms, the direct path terms, the individual attention head terms, and the virtual attention head terms. And something that might be useful to know before we spend a lot of time trying to go and really understand the behavior of the model is how important
[0:00:30] Are those different terms? And it turns out we can actually just explicitly measure that. And so that's what we'll do. So for the direct path term, which is just going straight up down the middle, if you haven't, if you're measuring relative to just making random predictions, you reduce the loss by 1.8 maths when you add that term to the model. So that's a pretty nice draw, especially just for a single, very simple term.
[0:01:04] The individual attention heads then, if we go and we add those in on top of the direct path term, so we have the individual attention heads, and they go and give us a reduction of 5.2 knots. Now, since there's 24 terms, it's actually quite a bit less than the direct path term per head. It's only 0.2 knots. But, yeah, they're still doing a lot of work.
[0:01:36] And then the virtual attention heads, if you already have the individual attention heads, they only give us a marginal additional 0.3 nats. And because there's quadratically many of these, right? Because for every first layer attention head and every second layer attention head, there's one of these pairs that are the virtual attention heads.
[0:02:02] The amount of loss reduction we get per term is very small, 0.002. So probably if we're trying to prioritize what we want to understand, our top priority should be, of course, the WWE, which is going to be very easy to understand. It's sort of representing bigrammish statistics and gives us a lot of bang for the effort. But then we want to study the individual attention head terms, probably.
[0:02:34] Now, you might wonder, how is it that we calculated this? And there's kind of a nice algorithm you can use to go and calculate how much these different terms contribute to the loss. And you can use it for this model, but you can also use it for larger models. It scales to potentially very large models. And the trick is, first you just run the model, and you save all your attention patterns, because we want to hold the attention patterns fixed. So you do that.
[0:03:04] Then you run the model and you're going to force all the attention patterns to be the attention pattern that you recorded. But instead of adding the attention head outputs to the residual stream, you'll just save them and replace them with zero. So you'll go and record what the attention heads would have added to the residual stream, but then you'll add zero instead and you record the loss. So that's just calculating the direct path contribution.
[0:03:35] Then you're going to go and iteratively run the model. And every time you run it, you're going to force the attention patterns to be the original attention patterns. But instead of going and adding the true output of your attention heads to the residual stream, you'll add whatever you would have added at the previous step that you replaced. And then you'll go and you'll save the output so you can go and add it on the next step.
[0:04:05] The nice thing about this is that it means that all of the attention heads only saw at most n-1 attention heads, or their values were only affected by n-1 attention heads, at least through the OV circuit, not through the effects of the attention patterns. But if you hold the attention patterns fixed, then that's true. And so that's a way that you can go and understand that.
[0:04:32] Of course, you could come up with all sorts of variants on this algorithm. So if you wanted to understand how important things were also going through the effects on attention patterns, you could also do that. Okay, so we think that probably we can ignore the virtual attention heads, or at least they aren't our top priority to understand for this model. Now, another question you could ask is we have really two, you know, for the attention head terms, we sort of have two different things. We have the
[0:05:03] First layer attention heads and the second layer attention heads. And we can ask in terms of what is their direct effect on the loss? How much do they affect the loss? And there's sort of two ways you could measure this. One is you could measure relative to there not being any attention heads, or you could measure how much do they add marginally on top of the other.
[0:05:36] Other heads already being present. So if you want to measure the layer one heads, you could either measure relative to the case where there's just no other attention heads or relative to the case where you've already added the layer two attention heads. And to make a long story short, the layer two heads have a much larger direct effect than the layer one heads. If you measure it in terms of the relative to only having the direct path, it's 5.2 NATs versus 1.3 NATs.
[0:06:06] If you measure it relative to the other one being present, it's 4 NATs versus 0.05 NATs. And if you then go and divide by the number of heads, those are obviously getting to be very small numbers for the layer 1. Now, that doesn't mean that the layer 1 attention heads aren't important. They can affect the attention patterns of the layer 2 attention heads, and we'll see that they're doing a really critical job in that. But it means that in terms of this overall equation we have, we don't need to go and worry as much about the layer 2 attention heads.
[0:06:38] So we can mostly understand the model by understanding the behavior of 12 attention heads. And those 12 attention heads, their attention patterns are influenced by earlier attention heads. But if we understand those 12 attention heads, for instance, you could just empirically understand their attention behavior as a starting point and their OV circuits, you'd understand the model's behavior. But yeah, it's really important to keep in mind that whenever you see
[0:07:11] The attention pattern here, particularly, not only is it nonlinear, but for later layers, that also contains potentially lots of effects of previous attention. And in fact, we can make that explicit if we want. So when we looked at the first layer, the equation for our attention pattern was this. And so it's a little bit complicated, but this is the thing that we called our QK circuit matrix.
[0:07:43] And we're multiplying by tokens on both sides. But when we go and then move to the two-layer model, all of a sudden it becomes a lot more complicated. And what's going on is that rather than just having the tokens get embedded and then affect things, they first go through all the attention heads in the previous layer, and then they hit the QK matrix.
[0:08:16] Then the keys also were computed using all of the attention heads after they went through the token. So on both sides, we have attention heads. And you can expand this all out, of course. And if you want to really analyze this, then you'll have a bunch of terms that involve first layer attention heads going and contributing to the attention pattern. You'll still have something like a QK circuit, but it'll be a bit more complex.
[0:08:46] In any case, now that we know what the right terms to focus on are, we're ready to go and explore the model and understand its behavior in our next video.