Tensor by Tensor

02.20 · UNIT 06 · Attention and position · Lab

Lab: implement multi-head attention

Multi-head attention runs several attention mixes with separate learned projections.

PLAIN-LANGUAGE INTRODUCTION

What is this?

Multi-head attention runs several attention mixes with separate learned projections.

One simple example

Head 1 uses weights [0.7,0.2,0.1]. Head 2 uses [0.1,0.3,0.6]. Each row sums to 1.

What goes in?

A token tensor plus one set of query, key, and value projections per head.

What comes out?

One context per head, joined back into the model width.

Why does it matter?

Different heads can learn different weighted relationships.

What is it not?

More heads do not automatically improve a model.

WORK THROUGH THE IDEA

See the idea in more detail

  1. Split the model width across heads. Each head makes its own query, key, and value tensors.
  2. For three visible tokens, head 1 has illustrative weights [0.7,0.2,0.1].
  3. Head 2 has [0.1,0.3,0.6]. Each row sums to 1 after softmax.
  4. Each head makes a weighted value sum. Join the results, then project to the original model width.
  5. Common mistake: model width must divide cleanly by the number of equal-width heads.
Open the detailed notes ↗