GQA: Training von generalisierten Multi-Query Transformer-Modellen aus Multi-Head-Checkpoints

Multi-Query-Attention (MQA), die nur einen einzelnen Schlüssel-Wert-Kopf verwendet, beschleunigt die Dekodierer-Inferenz drastisch. MQA kann jedoch zu Qualitätsverlusten führen, und außerdem kann es unerwünscht sein, ein separates Modell nur für eine schnellere Inferenz zu trainieren. Wir (1) schlagen ein Rezept vor, um bestehende Multi-Head-Sprachmodell-Checkpts in Modelle mit MQA unter Verwendung von 5 % der ursprünglichen Vortraining-Rechenressourcen zu uptrainieren und (2) führen die Grouped-Query-Attention (GQA) ein, eine Verallgemeinerung der Multi-Query-Attention, die eine intermediäre Anzahl von Schlüssel-Wert-Köpfen (mehr als einen, weniger als die Anzahl der Abfrageköpfe) verwendet. Wir zeigen, dass uptrainierte GQA eine Qualität erreicht, die der Multi-Head-Attention nahekommt, mit vergleichbarer Geschwindigkeit zu MQA.