Multi-head attention, GQA, and MQA

Multi-head attention applies several sets of learned projections to the same input, then combines their outputs. This lesson follows the tensor shapes and explains why sharing key and value heads can reduce the memory needed during generation.

The previous lesson computed a single masked attention operation. Each position combined information from itself and earlier positions using one set of query, key, and value projections.

A single head provides one learned way of comparing and combining token representations. Multiple heads give the model additional projection spaces in which it can represent different relationships.

This is the purpose of multi-head attention.

Multi-head attention

Instead of using one set of W_Q, W_K, and W_V projections, multi-head attention uses several sets and calculates their attention outputs in parallel.