KVComm: Enabling Efficient LLM Communication through Selective KV Sharing

1. Problemstellung

Große Sprachmodelle (LLMs) werden zunehmend in Multi-Agenten-Systemen eingesetzt, in denen eine effektive Kommunikation zwischen den Modellen entscheidend ist. Bestehende Kommunikationsprotokolle weisen jedoch erhebliche Nachteile auf:

• Natürlichsprachliche Kommunikation: Methoden wie „Natural Language Debate" (NLD) oder der Austausch von Embeddings (CIPHER) erfordern mehrere Dekodierschritte, was hohe Inferenzkosten verursacht. Zudem geht durch den Sampling-Prozess beim Erstellen neuer Tokens wertvolle Information verloren.
• Versteckte Zustände (Hidden States): Ansätze wie „AC" (Activation Communication), die den versteckten Zustand des letzten Tokens übertragen, leiden unter einem Informationskonzentrations-Bias. Der letzte Token enthält zwar viel Information für die Ausgabe, aber das Übertragen aller Tokens versteckter Zustände führt zu einem Dilemma:
  • Werden Zustände aus frühen Schichten übertragen, ist der Rechenvorteil gering (ähnlich wie das einfache Aneinanderhängen der Eingaben).
  • Werden sie in spätere Schichten eingefügt, bricht die Leistung stark ein.
• Ziel: Es fehlt ein Protokoll, das sowohl effektiv (hohe Informationsdichte) als auch effizient (geringe Rechen- und Übertragungskosten) ist und auf verschiedene Aufgaben und Architekturen anwendbar ist.

2. Methodik: KVComm

Die Autoren schlagen KVComm vor, ein Framework, das die Kommunikation zwischen LLMs durch selektives Teilen von Key-Value (KV) Paaren ermöglicht.

Kernkonzept

Anstatt ganze versteckte Zustände oder Text zu übertragen, nutzt KVComm die KV-Paare, die im Attention-Mechanismus jeder Schicht des Senders ($M_s$) generiert werden. Diese Paare enthalten reichhaltige semantische Informationen. Der Empfänger ($M_r$) integriert diese KV-Paare direkt in seinen Attention-Mechanismus, ohne dass $M_s$ dekodierte Tokens erzeugen muss.

Selektionsstrategie

Da das Übertragen aller KV-Paare aller Schichten zu teuer wäre, wird eine selektive Strategie entwickelt, um nur die informativsten Schichten auszuwählen. Diese basiert auf zwei Hypothesen:

1. H1 (Semantische Transferierbarkeit): KV-Paare aus mittleren Schichten enthalten die übertragbarsten semantischen Kenntnisse (frühe Schichten erfassen Oberflächenmuster, späte Schichten sind taskspezifisch).
2. H2 (Aufmerksamkeitsverteilung): Schichten mit einer stärkeren Aufmerksamkeitsverteilung (d.h. hohe Gewichtung auf Kontext-Token) sind effektiver für die Kommunikation.

Algorithmus zur Auswahl:

1. Attention Importance Score ($S_a$): Für jede Schicht wird ein Score berechnet, der den durchschnittlichen Aufmerksamkeitsgewichten aller Heads auf die Kontext-Token während der Prefill-Phase entspricht.
2. Gaussian Prior: Um Hypothese H1 zu unterstützen, wird eine Gauß-Verteilung als Prior verwendet, die Schichten um eine bestimmte Tiefe ($\mu$) bevorzugt.
3. Bewertung: Der endgültige Auswahl-Score $S_l$ für eine Schicht $l$ ist eine gewichtete Kombination aus dem Attention-Score und dem Gaussian-Prior:
   $$S_l = \alpha S_a^l + (1-\alpha) P^l$$
   Dabei sind die Top-$M$ Schichten mit den höchsten Scores für die Kommunikation ausgewählt.
4. Kalibrierung: Die Auswahl kann bereits mit einem einzigen Kalibrierungs-Beispiel (Sample) erfolgen und generalisiert gut auf den gesamten Testdatensatz.

Integration

Der Empfänger $M_r$ fügt die ausgewählten KV-Paare des Senders in seine eigenen KV-Puffer ein (Konkatenation), sodass $M_r$ sowohl auf den eigenen Kontext als auch auf die vom Sender bereitgestellten Informationen achten kann.

3. Wichtige Beiträge

1. Analyse von Aktivierungsdaten: Die Arbeit identifiziert die Grenzen der Verwendung von versteckten Zuständen (Hidden States) als Kommunikationsmedium, insbesondere den Informationskonzentrations-Bias des letzten Tokens und das Dilemma bei der Nutzung aller Tokens.
2. KVComm Framework: Einführung eines neuen Protokolls, das KV-Paare als effizientes Medium nutzt. Es ist der erste Ansatz, der die Auswahl nicht-kontinuierlicher Schichten (non-contiguous layers) ermöglicht, was flexibler und effektiver ist als das Übertragen eines einzigen zusammenhängenden Blocks.
3. Selektionsstrategie: Entwicklung einer Strategie basierend auf Attention-Importance-Scores und einem Gaussian-Prior, die sich als robust und mit minimalem Kalibrierungsaufwand (ein Sample) erweist.
4. Umfassende Evaluation: Experimente mit neun verschiedenen Modellpaaren (inkl. LLaMA, Qwen, Falcon) und diversen Aufgaben (Reasoning, Mathematik, QA).

4. Ergebnisse

Die Evaluation zeigt, dass KVComm bestehende Methoden deutlich übertrifft:

• Leistung: KVComm erreicht eine Leistung, die mit der Skyline-Obergrenze (direktes Aneinanderhängen von Kontext und Query ohne Kommunikation) vergleichbar ist, obwohl nur ein Bruchteil der Daten übertragen wird. In einigen Fällen übertrifft KVComm sogar die Skyline-Methode, vermutlich durch Regularisierungseffekte und komplementäre Fähigkeiten der Modelle.
• Effizienz:
  • KVComm überträgt nur 30% bis 70% der KV-Paare (Schichten), erreicht aber oft bessere Ergebnisse als Methoden, die 100% der Daten (wie Skyline) oder ganze Textsequenzen übertragen.
  • Im Vergleich zur Skyline-Methode reduziert KVComm die Rechenkosten um den Faktor 2,5x bis 6x.
  • Im Vergleich zu Methoden, die den gesamten KV-Satz übertragen, reduziert sich die Kommunikationsmenge um bis zu 3x.
  • Der Speicherverbrauch sinkt ebenfalls signifikant (bis zu 73% weniger auf dem Datensatz Tipsheets).
• Vergleich mit Baselines: KVComm übertrifft konsistent Methoden wie NLD, CIPHER und AC (Hidden States) auf komplexen Aufgaben mit langem Kontext. NLD und CIPHER funktionieren nur bei sehr einfachen Aufgaben, bei denen wenig Information übertragen werden muss.
• Selektion vs. Zufall: Die gezielte Auswahl von Schichten (basierend auf Attention-Scores) übertrifft eine zufällige Auswahl deutlich, insbesondere bei niedrigen Übertragungsraten (30%).

5. Bedeutung und Ausblick

• Paradigmenwechsel: Die Arbeit etabliert KV-Paare als effektives Medium für die Inter-LLM-Kommunikation, das den Informationsverlust natürlicher Sprache und die Ineffizienz versteckter Zustände umgeht.
• Skalierbarkeit: Durch die Reduktion der Kommunikationskosten und der Rechenlast ermöglicht KVComm skalierbare und effiziente Multi-Agenten-Systeme.
• Zukunftspotenzial: Das Framework ist erweiterbar auf mehrere Sender/Empfänger und könnte durch adaptive, online-kalibrierende Mechanismen weiter optimiert werden. Eine aktuelle Einschränkung ist die Abhängigkeit von ähnlichen Basis-Architekturen (gleiche oder fine-tuned Modelle), da die Struktur der KV-Paare zwischen verschiedenen Modellfamilien variiert. Zukünftige Arbeiten könnten Adapter oder Projektionen für heterogene Architekturen entwickeln.

Zusammenfassend bietet KVComm einen robusten, effizienten und leistungsfähigen Ansatz für die Kommunikation zwischen LLMs, der die Grundlage für die nächste Generation von kollaborativen KI-Systemen legt.