One of the better visualisations I've seen with the exception of Q/K/V weights and how they are presented. I believe they should be put more upfront since they are the core learnable parameter of attention. IMO, they should also be part of the "Head N of M" block since each head has its own weights (although they all can be collapsed into one huge matrix computation).