Shazeer (2019), "Fast Transformer Decoding: One Write-Head is All You Need," introduces multi-query attention as a modification to standard multi-head attention aimed at speeding up autoregressive decoding. What does multi-query attention change relative to standard multi-head attention, and why does this help?
- It reduces the number of query heads down to a single shared query head, while each head keeps its own separate key and value projections, cutting the number of query computations performed at each step
- It removes the key and value projections entirely, having every attention head operate directly on the raw token embeddings instead of projected keys and values
- It increases the number of key and value heads beyond the number of query heads, giving each query head access to several redundant copies of the same key and value vectors
- It keeps multiple separate query heads but has all of them share a single set of key and value projections, shrinking the size of the cached keys and values that must be stored and read back at every decoding step, which reduces the memory-bandwidth cost that dominates incremental decoding
Why D? And why not the others?
Correct answer: D. It keeps multiple separate query heads but has all of them share a single set of key and value projections, shrinking the size of the cached keys and values that must be stored and read back at every decoding step, which reduces the memory-bandwidth cost that dominates incremental decoding
Shazeer's multi-query attention keeps the multiple, separately learned query heads of standard multi-head attention, but has all of those heads share one single set of key and value projections instead of each head computing and caching its own; because the memory bandwidth needed to read the cached keys and values back from memory at every incremental decoding step is the dominant cost of autoregressive generation, shrinking the cached keys and values down to a single shared copy per layer substantially reduces that bandwidth cost, at a modest expense in model quality. The option describing a single shared query head with separate keys and values per head is wrong because it inverts the paper's actual change: multiple query heads are kept, and it is the keys and values that are shared, not the queries. The option describing removal of key and value projections entirely is wrong because keys and values are still computed and cached, just from one shared projection rather than one per head. The option describing more key/value heads than query heads is wrong because multi-query attention reduces the key/value heads down to exactly one, rather than increasing them beyond the query head count.
Source: Shazeer, "Fast Transformer Decoding: One Write-Head is All You Need" (2019), arXiv:1911.02150