GQA: обучение обобщённых трансформерных моделей с grouped-query attention на основе чекпоинтов с многоголовым вниманием

Multi-query attention (MQA), использующее одну пару ключ–значение, существенно ускоряет вывод декодера. Однако MQA может приводить к снижению качества, и может быть нежелательно обучать отдельную модель исключительно для более быстрого вывода. Мы (1) предлагаем процедуру дообучения (uptraining) существующих чекпоинтов языковых моделей с многоголовым вниманием в модели с MQA, используя 5% от изначальных вычислительных ресурсов предобучения, и (2) вводим grouped-query attention (GQA) — обобщение multi-query attention, которое использует промежуточное число пар ключ–значение (больше одной, но меньше числа голов запросов). Мы показываем, что дообученная GQA достигает качества, близкого к многоголовому вниманию, при скорости, сопоставимой с MQA.