MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding
Zuhri et al
arxiv.org/abs/2406.09297

Помните multi-query attention где вместо того чтобы у вас были разные key, query, value на каждую голову вы используете одни и те же kv для всех и разными бывают только query? Это очень классная идея которая сильно сокращает размер kv cache позволяя упихивать более длинные тексты или большие батчи в тот же объем памяти с минимальной потерей качества.

Авторы multi-layer kv heads предлагают сделать еще один шаг в эту сторону и предлагают шарить kv между соседними слоями тоже. То есть мы разбиваем наши слои на блоки, например по 4 слоя. И в каждой группе из 4 слоёв только лишь первый слой считает kv, остальные используют только query.

Потеря качества хоть и небольшая, но достаточно заметная, но иногда приходится упихивать неупихиваемое в GPU и хорошо иметь больше способов делать tradeoffs.