Multi-Query attention from PaLM

[ananyahjha93]

🚀 The feature, motivation and pitch

Iz and I were talking about this today, can be under nice to have

This is from the PaLM paper:

"The standard Transformer formulation uses k attention heads, where the input vector for each timestep is linearly projected into “query”, “key”, and “value” tensors of shape [k, h], where h is the attention head size. Here, the key/value projections are shared for each head, i.e. “key” and “value” are projected to [1, h], but “query” is still projected to shape [k, h]. We have found that this has a neutral effect on model quality and training speed (Shazeer, 2019), but results in a significant cost savings at autoregressive decoding time. This is because standard multi-headed attention has low efficiency on accelerator hardware during auto-regressive decoding, because the key/value tensors are not shared between examples, and only a single token is decoded at a time."

Alternatives

No response

Additional context

No response

[epwalsh]

Duplicate of #30 😉