ch03/01_main-chapter-code/ch03.ipynb - Masked MHA

[d-kleine]

I just had a discussion with Grant from 3Blue1Brown, I want to point out an important aspect that's missing in the visualizations, in my opinion, and might be important for the understanding from other readers as well:

3.4.2 Implementing a compact SelfAttention class

3.5 Hiding future words with causal attention

(P.S.: The normalized attention weights in the last graphic also don't add up row-wise to 1.0 perfectly for most cases there, maybe this could be fixed too.)

From the code, you can understand that the queries are multiplied by the keys:

attn_scores = queries @ keys.T

I think it would be good in the last graphic to indicate with axes labels that the queries are the rows, whereas the keys are the columns - because sometimes you see visualizations where the attention weight table is transposed, like here.

Quick draft:

[rasbt]

Thanks for suggesting, that's an interesting one! So what you have in mind is something like the following?

The only thing I am maybe a bit worried about is that this figure is a bit crowded then, and this annotation might be not very clear. We do compute the attention scores as attn_scores = queries @ keys.T so from from that, maybe it's already sufficient to see that the rows come from the queries?

[d-kleine]

Thanks for suggesting, that's an interesting one! So what you have in mind is something like the following?

My idea would be to clearly illustrate that the queries will be the rows and keys the columns in the attention matrix, for instance like that:

For me, these changes are not that important. More important for me are the changes in the following visualization, indicating the rows represent the query indices and the columns the key indices.

For the attention matrix, I would indeed add Query/Key sequence position labels, for instance like this:

Also, to enhance its accuracy of the visualization of the normalized attention weights matrix, I recommend ensuring that the normalized attention weights in each row sum precisely to 1.0 (100%). Presently, this condition is not met for most rows, which might lead to confusion among readers. Achieving this consistency would greatly improve the clarity and reliability of the data, making the matrix more intuitive and reflective of the intended proportional relationships among the weights.

The only thing I am maybe a bit worried about is that this figure is a bit crowded then, and this annotation might be not very clear.

I see what you mean; there's a lot going on in this visualization! It seems like the main issue is the amount of text within the workflow diagram itself. Personally, I find that including too much text inside a workflow can make it a bit cluttered. As an alternative, maybe you could try moving the textual descriptions to the caption or text instead? That might help make the diagram cleaner and easier to follow.

We do compute the attention scores as attn_scores = queries @ keys.T so from from that, maybe it's already sufficient to see that the rows come from the queries?

From my point of view, this is only clear when examining the code. The diagram alone does not make it obvious what the rows and columns in the attention weights matrix represent. Furthermore, this understanding can be influenced by the vector representation convention (such as whether column or row vectors are used). This distinction is crucial, especially for comprehending how masking is applied later on.

[d-kleine]

As an idea for the supplementary materials, I think it would be cool to see Grouped Query Attention (GQA) and Multi-Query Attention (MQA), especially for understanding Llama or Mistra, implemented "from scratch" in this chapter too. I think this would a great deep-dive after understanding Multi-Head Attention (MHA).

[rasbt]

100%. I plan to convert GPT to Llama 3 some time, and as part of this effort, I want to explain/implement GQA. (And maybe something similar for Mistral). In the meantime, you might like our CausalAttention class in LitGPT, which already supports both GQA and MQA: https://github.com/Lightning-AI/litgpt/blob/0bb34ab3aad9fba0716fe1c37e2924994eb0d400/litgpt/model.py#L191

[d-kleine]

Great, cool! 👍🏻 Looking forward to this! 🙂

[rasbt]

I just realize that I never responded to this 😅! Thanks so much for taking the time to actually draw and outline this!

For me, these changes are not that important. More important for me are the changes in the following visualization, indicating the rows represent the query indices and the columns the key indices. For the attention matrix, I would indeed add Query/Key sequence position labels, for instance like this:

I would say though the 2nd figure shown here shows the same attention matrix as the one above. So if it is clear from the queries @ keys.T calculation that the row values are derived from the query and the key values are derived in the columns, then it should also be clear in the 2nd figure, because it's the same figure?

Of course, you are free to multiply the matrices as keys.T @ queries so the results are transposed, but I don't know anyone who does that. I.e. all common implementations and figures assume queries @ keys.T afaik.

[d-kleine]

So if it is clear from the queries @ keys.T calculation that the row values are derived from the query and the key values are derived in the columns, then it should also be clear in the 2nd figure, because it's the same figure?

Of course, you are free to multiply the matrices as keys.T @ queries so the results are transposed, but I don't know anyone who does that. I.e. all common implementations and figures assume queries @ keys.T afaik.

Therefore I had the discussion with Grant (3Blue1Brown). He handled the vectors as column vectors, and therefore computed the matrices as keys.T @ queries.

I see this from the eyes of a learner: Especially for readers are just learning the concept of the attention mechanism with this book, this might not be clear. That's the crucial point for me. There is a attention weights table, but I can't clearly see if you either multiply keys.T @ queries or queries @ keys.T.

And I also think it's good practice to label any kind of visualization (plots, tables). Again, seeing this from the eyes of a reader having no experience with the attention mechanism:

• Attention weights matrix (without column and row labels)
  

• Attention weights matrix (with column and row labels)
  

What do you think is more concise and understandable for someone unfamiliar with the attention mechanism?

[rasbt]

I see what you mean. I am not sure though if a "Query/key position" label would add confusion though because then you'd have to explain that you can also do the multiplication the other way around, which would be going off on a tangent on alternative implementations.

Therefore I had the discussion with Grant (3Blue1Brown). He handled the vectors as column vectors, and therefore computed the matrices as keys.T @ queries.

I see. I think this is highly unusual though, in machine learning and deep learning, you usually always consider training examples as rows and features as columns. The only exceptions I can think of are older statistics textbooks.

[d-kleine]

I see what you mean. I am not sure though if a "Query/key position" label would add confusion though because then you'd have to explain that you can also do the multiplication the other way around, which would be going off on a tangent on alternative implementations.

I see. I think this is highly unusual though, in machine learning and deep learning, you usually always consider training examples as rows and features as columns. The only exceptions I can think of are older statistics textbooks.

Yeah, I agree. But imagine you are for example a young student (high school or university). In school, you are typically taught to represent vectors as columns. In machine learning, you do that differently, as you said. Maybe you should add this information as a short note in the text, to make this clear. It think this info can be really helpful for some reader understanding why queries are put in the rows and keys in the columns in the attention weight matrix, and not vice versa.

[rasbt]

Personally, I find that including too much text inside a workflow can make it a bit cluttered. As an alternative, maybe you could try moving the textual descriptions to the caption or text instead? That might help make the diagram cleaner and easier to follow.

I agree, but this is a publisher convention. Originally, I had more explanations in the text, but I was asked to use figure annotations instead. Not much I can do here, unfortunately.

Also, to enhance its accuracy of the visualization of the normalized attention weights matrix, I recommend ensuring that the normalized attention weights in each row sum precisely to 1.0 (100%).

That's a good point, too! I originally rounded the values everywhere but the editor suggested that numbers in the figures have to be the same as in the code. So I then truncated the numbers everywhere instead of rounding them. But then yeah, it gets tricky in cases like this. It's mentioned in the first figure where the truncation is applied. I.e., it says " (Note that the numbers in the figure are truncated to one digit after the decimal point to reduce visual clutter.)" in the figure caption. Maybe I should copy & paste that phrase to the other figures as well.

[d-kleine]

Ah, I got it - that makes sense! Thanks!

If the book doesn't include this background information yet, I think it would be helpful to also add it as a side note in an appropriate section, e.g. the section about normalization or when building GPT-2 from scratch.

[d-kleine]

I just have checked in the ipynbs in the chapters about the the MHA. There is one little thing I still miss in the book:

• an all-in-one overview figure about the transformer architecture and the attention mechanism

I just drew a little illustration for myself - even though it might be a little too late, what do you think about implementing a figure like this?

[rasbt]

Thanks for the suggestion. This looks really nice and clear. I probably can't show the left part (transformer block) in chapter 3 because it's not been introduced, yet. But the MHA part looks really nice. It could maybe be shown as a chapter summary after Figure 3.22.

Unfortunately, chapter 1-6 have already been layouted, and I am not allowed to make any big changes that are not typo fixes, because they would change the page numbers. However, I will keep this idea in mind for a potential 2nd edition in 2-3 years! I have a whole document of ideas already :).

[d-kleine]

Ok cool, thanks!

From my point of view, we have covered everything for now about the MHA in the book. So if you like, you can close this discussion. 🙂

[rasbt]

Sounds good, and many thanks for the good and detailed feedback!