Introduction to Multi-Head Attention (MHA) | The attention family series - Lecture 8

This is Lecture 8 of The Attention Family. In this lecture we run several attention heads in parallel, see why we need them, build them first with a simple loop, and then build them the way GPT actually computes them: one large matrix multiplication for all heads together. In this lecture I cover: - Why one head is not enough: in "The dog chased the ball, but it could not catch it", the last "it" needs its noun (ball) and its verb (catch), and one row of softmax weights has to answer both questions at once - The main idea: each head has its own Wq, Wk, Wv. Two heads of 6 by 3 are concatenated into a 6 by 6 context matrix, the same width as the input - Why we multiply by W out after concatenating, so that every output feature uses information from every head - The dimension rule: number of heads times head dimension equals the width. GPT-2 small: 768 = 12 × 64. GPT-3: 12288 = 96 × 128 - Looping over heads: copies of the causal head from lecture 7, run one after another. 2 heads take about a millisecond, 4000 heads take about 0.9 seconds - The weight split: one large Wq, Wk, Wv for all heads, then view to split out the heads axis - Why we transpose(1, 2) before Q @ K.T, so that we get one 6 by 6 score matrix per head instead of comparing heads with heads - Every tensor shape from input to output: (2, 6, 6) to (2, 6, 2, 3) to (2, 2, 6, 3), scores at (2, 2, 6, 6), and back to (2, 6, 6) - The five steps used in practice: project, split, attend, merge, multiply by W out - A question for the next lecture: "All dogs are animals" and "All animals are dogs" give the same attention outputs, because attention does not know word order Slides for this lecture: https://lecture8-attention-mechanisms.vercel.app Colab notebook (the loop version with timing, the weight split step by step with an equality check against the loop, and the word order experiment): https://colab.research.google.com/drive/1WVYWQh4YcjQq-kPDi1Hug-KYDud-bDkK?usp=sharing Next lecture: how position information gets into attention. Dr. Sreedath Panat MIT PhD, IIT Madras BTech, Co-founder, Vizuara AI Labs