Reduzierte KV-Cache-Nutzung durch neuen Aufmerksamkeitsmechanismus

Einführung in die Herausforderungen der LLM-Inferenz

Die zunehmende Verwendung von großen Sprachmodellen (LLMs) und das Aufkommen neuer Inferenzparadigmen haben die Herausforderung der effizienten groß angelegten Inferenz in den Vordergrund gerückt. Ein erheblicher Engpass ist der Key-Value (KV)-Cache in traditionellen Aufmerksamkeitsmechanismen, der mit der Batchgröße und Sequenzlänge linear ansteigt und zu einem "Speicherfresser" wird, der die Skalierung und Erweiterung von LLMs behindert.

Obwohl Varianten wie MQA, GQA und MLA entstanden sind, um dieses Problem zu beheben, haben sie oft Schwierigkeiten, die Leistung unter strengen Speicherbeschränkungen aufrechtzuerhalten oder Komplexitäten einzuführen, die technische Herausforderungen und Kompatibilitätsprobleme mit sich bringen.

Multi-Matrix Factorization Attention (MFA)

Ein kürzlich von Forschern von Stepes, der Tsinghua Universität und anderen Institutionen veröffentlichtes Paper stellt eine neuartige Aufmerksamkeitsmechanismusarchitektur vor: Multi-Matrix Factorization Attention (MFA) und ihre Variante MFA-Key-Reuse (MFA-KR). Dieser Mechanismus reduziert die Kosten der Sprachmodellinferenz erheblich und verbessert gleichzeitig die Leistung.

MFA und MFA-KR übertreffen nicht nur MLA in der Leistung, sondern erreichen auch die Leistung des traditionellen MHA, während die KV-Cache-Nutzung um bis zu 93,7 % reduziert wird.

MFA ist auf Einfachheit, einfache Reproduzierbarkeit, geringe Empfindlichkeit gegenüber Hyperparametern und Kompatibilität mit verschiedenen Pos-Embedding-Methoden ausgelegt.

MFA Ansatz und Analyse

Das Forschungsteam analysierte das allgemeine Design und die Kapazität von Aufmerksamkeitsmechanismen und identifizierte zwei kritische Dimensionen, die mit der Kapazität zusammenhängen. Diese Analyse führte zur Entwicklung neuer Analysemethoden und Designprinzipien.

Sie führten das Konzept der Generalized Multi-Head Attention (GMHA) als vereinheitlichenden Rahmen für das Verständnis verschiedener MHA-Varianten ein.

Das Team untersuchte auch die Berechnung und Speicherung von Key-Values aus einer Inferenzperspektive und untersuchte die Modellkapazität aus einer Dekompositionsperspektive.

Fully Parameterized Bilinear Attention (FPBA) wurde als theoretische Obergrenze der Leistung etabliert. Sie fanden heraus, dass MHA und seine Varianten Low-Rank-Zerlegungen von FPBA sind.

Vergleich mit MQA und MLA

Die Analyse konzentrierte sich auf zwei repräsentative Verbesserungsschemata: Multi-Query Attention (MQA) und Multi-Head Latent Attention (MLA).

MQA verwendet eine aggressivere Parameter-Sharing-Strategie, bei der sich alle Aufmerksamkeitsköpfe den gleichen Satz von Key-Value-Parametern teilen. Dies reduziert den Speicherverbrauch, kann aber die Ausdruckskraft des Modells beeinträchtigen.

MLA führt einen gemeinsamen latenten Raum für die Parameterkompression ein, aber die tatsächliche Ausdruckskraft ist durch die kleinste Dimension begrenzt, was bedeutet, dass die Erhöhung der Zwischendimensionen die Leistung nicht wesentlich verbessert.

MFA-Schlüsselinnovationen

Die Entwicklung von MFA wurde durch das Ziel angetrieben, einen Aufmerksamkeitsmechanismus zu schaffen, der den Ressourcenverbrauch minimiert und gleichzeitig die theoretischen Leistungsgrenzen erreicht. Das Design von MFA umfasst drei wichtige Innovationen:

• Signifikante Erhöhung der Anzahl und Dimension der Aufmerksamkeitsköpfe, um die Modellkapazität zu maximieren.
• Einsatz einer aggressiven Low-Rank-Zerlegungsstrategie, um die Parametereffizienz aufrechtzuerhalten, während die Anzahl und Dimension der Aufmerksamkeitsköpfe erweitert wird.
• Verwendung eines einzelnen Key-Value-Head-Designs, um den Speicherverbrauch auch bei erhöhter Modellkomplexität minimal zu halten.

Kapazitätsmessung und Vergleich

Um MFA und andere Aufmerksamkeitsmechanismen weiter zu analysieren, führte das Team zwei wichtige Metriken ein:

• Total Effective Rank (TER): Das Produkt aus der Anzahl der Aufmerksamkeitsköpfe und dem Factorization Rank per Head (FRH).
• Shared Latent Subspace Dimension (SLSD): Die Dimension des versteckten Raums, der von allen Aufmerksamkeitsköpfen gemeinsam genutzt wird.

MFA erreicht im Vergleich zu MQA einen höheren SLSD und TER.

Im Vergleich zu MLA erreicht MFA eine kleinere KV-Cache-Größe und einen höheren TER bei ähnlichen Parameterbudgets, während ein vergleichbarer SLSD beibehalten wird.

Im Vergleich zum traditionellen MHA hat MFA einen höheren TER, obwohl sein SLSD kleiner ist.

Experimentelle Ergebnisse

Umfangreiche Experimente wurden durchgeführt, um die Leistung der neuen Architektur in größeren Maßstäben zu bewerten, wobei Modelle von 1B bis 7B Parametern und Trainingsdaten von 10B bis 1T getestet wurden.

MFA demonstrierte Skalierungsfähigkeiten, die mit dem traditionellen MHA vergleichbar sind, und behielt auch in größeren Maßstäben eine ausgezeichnete Leistung bei.

Während MFA-KR eine etwas geringere Leistung zeigte, entsprach sein Skalierungstrend dem von MHA. Die Speichervorteile von MFA und MFA-KR nahmen mit der Modellgröße weiter zu, wobei MFA 87,5 % Speichereinsparungen erreichte und MFA-KR die Speichernutzung in der größten Größenordnung auf 6,25 % reduzierte.

Ablationsstudien

Ablationsstudien validierten die Wirksamkeit von MFA und MFA-KR. Ihre Leistungsvorteile wurden auch über verschiedene Mainstream-Positionskodierungsmethoden hinweg bestätigt.

Ausblick

MFA bietet signifikante Verbesserungen mit einem einfachen Design und behebt effektiv den Speicherengpass bei der LLM-Inferenz, ohne zusätzliche technische Komplexität hinzuzufügen. Es integriert sich nahtlos in das bestehende Transformer-Ökosystem und beschleunigt die Anwendung von LLMs in verschiedenen Szenarien.