Lu Baotong von Microsoft Research: Neuformung der Modell-Attention mit Vektor-Suche – Attention

In großen Sprachmodellen ist die Fähigkeit zur Inferenz über sehr lange Kontexte einer der Hauptengpässe, die ihre Leistung beeinflussen.

Dies liegt an der quadratischen Komplexität der Self-attention und der Tatsache, dass der KV-Cache den Grafikspeicher frühzeitig beansprucht und diese Beanspruchung mit zunehmender Länge steigt. Beispielsweise kann bei einem 8B-Modell der KV-Cache beim Inferieren eines Kontexts von 1M Token leicht über 100 GB Grafikspeicher beanspruchen. Ein normaler GPU kann in einem solchen Fall einfach nicht funktionieren.

Angesichts dieses Problems konzentriert sich ein neues Interviewartikel aus der Attention-Serie auf einen neuen Mechanismus, der in der Studie  Retrieval Attention: Accelerating Long-context LLM Inference via Vector Retrieval  vorgeschlagen wird: ein trainingsfreies, dynamisches und sparses Attention-Schema für die Inferenz über sehr lange Kontexte.

Im Folgenden finden Sie eine Zusammenfassung eines Interviews zwischen Oasis und Dr. Baotong Lu, einem der Kernautoren dieser Studie und einem erfahrenen Forscher am Microsoft Research Asia. Das Lesen des gesamten Artikels dauert etwa 20 Minuten.

Enjoy

Der Kerngedanke von Retrieval Attention besteht darin, dass jeder Query in der Tat nur mit einem kleinen Teil der Keys stark interagieren muss, während der Rest der Attention redundant ist. Darüber hinaus ist Attention von Natur aus sparse.

Deshalb besteht die Kernmaßnahme des Forschungsteams darin: Die meisten KV-Vektoren vom GPU auf die CPU zu verlagern. Jedes Mal, wenn das Modell eine Inferenz durchführen muss, wird die Annähernde-Nächste-Nachbarschaft-Suche (ANN Search) verwendet, um die wenigen Keys zu finden, die für den aktuellen Query am relevantesten sind. Beispielsweise werden nur die ersten 1 % ausgewählt und dann zusammen mit den wenigen "vorhersagbaren" KV auf dem GPU parallel die Attention berechnet und die Ergebnisse zusammengeführt.

In der Praxis kann ein 8B-Modell auf einem RTX4090 (24 GB) bei einem Kontext von 128K stabil generieren, wobei die Zeit pro Token etwa 0,188 s beträgt und die Genauigkeit fast identisch mit der der Vollattention ist (∞-Bench, RULER und andere Benchmarks).

In der Folgearbeit RetroInfer wurde der Leistungsvorteil weiter ausgebaut. Im Vergleich zur Vollattention (Full attention) auf einem A100 GPU wurde eine 4,5-fache Dekodierdurchsatzrate erreicht. Bei einem Kontext von 1M Token wurde im Vergleich zu anderen GPU-CPU-sparsen Attention-Systemen eine 10,5-fache Durchsatzrate erzielt.

Die Kerninnovation von Retrieval Attention besteht darin, dass es den Vektorretrieval-Mechanismus in den Attention-Berechnungspfad einführt und auf der Berechnungsebene eine echte dynamische Sparsifizierung realisiert.

Konkret hat der Autor einen attention-aware Vektorindex (attention-aware retrieval index) aufgebaut: Indem er während der Vorauffüllphase die genaue Abbildung zwischen Query und Key lernt und diese Beziehung auf die Key–Key-Graphstruktur projiziert, kann der Retrievalprozess mit einer sehr geringen Scanrate (etwa 1–3 %) eine hohe Recallrate erreichen. Diese Gestaltung ermöglicht es jedem Query, die wichtigsten Keys mit minimalem Rechenaufwand dynamisch zu lokalisieren, wodurch das Problem der Ineffizienz bei der Anwendung der traditionellen ANN-Suche auf die Attention aufgrund der unterschiedlichen Verteilungen von Query und Key überwunden wird.

In der Systemarchitektur schlägt Retrieval Attention einen weitergehenden Zweiwege-Attention-Mechanismus mit CPU–GPU-Kooperation vor: Der GPU ist für die Aufrechterhaltung eines kleinen Teils des "vorhersagbaren" lokalen KV-Caches zuständig, während die CPU die großen KV-Speicher dynamisch über die Retrievalmethode aufruft. Die beiden Berechnungen laufen unabhängig und parallel ab, und schließlich werden die Ergebnisse durch eine numerisch stabile Re-Skalierungsformel fusioniert, was sowohl den Grafikspeicherbedarf als auch die Inferenzverzögerung verringert.

Es ist erwähnenswert, dass der gesamte Mechanismus keine erneute Training des Modells erfordert. Retrieval Attention kann als ein austauschbarer Modul in das bestehende Transformer eingefügt werden. Nur indem die Vorwärtslogik der Attention-Schicht geändert wird, kann die Inferenz über lange Kontexte ohne Genauigkeitsverlust erheblich beschleunigt werden. Dieser Gedanke des "trainingsfreien dynamischen sparsen Attention" bietet einen praktikablen technischen Weg für die Skalierbarkeit von Sprachmodellen mit langen Kontexten.

Im Folgenden finden Sie einen Teil des Dialogs zwischen Oasis und Dr. Lu:

Oasis: Bitte geben Sie zunächst einen Überblick über Ihren Forschungsbackground, insbesondere über Ihre Forschung im Bereich Datenbankmanagementsysteme (Database Management System), und erläutern Sie, welche Überlegungen Ihnen zu diesem Entwurf geführt haben.

Dr. Lu: Mein Hauptforschungsgebiet sind Datenbankmanagementsysteme, d. h. wie man riesige Datenmengen effizient speichert, abruft und verwaltet. Aus wissenschaftlicher Sicht scheinen Datenbanksysteme und maschinelles Lernen zwei relativ unabhängige Bereiche zu sein, wobei der eine eher auf Systeme und Strukturen, der andere eher auf Modelle und Optimierungen ausgerichtet ist. Aber auf der untersten Ebene gibt es einige gemeinsame Probleme, wie beispielsweise, wie man Informationen effizienter organisiert, wenn die Ressourcen begrenzt sind.

Während meines Promotionsstudiums habe ich hauptsächlich die Auswirkungen neuer Hardwarearchitekturen, wie etwa persistenten Arbeitsspeicher und getrennten Arbeitsspeicher, auf das Datenbankdesign untersucht. Dies umfasste insbesondere die Entwicklung neuer Datenstrukturen und Datenbankindizes für diese Hardware sowie die Optimierung der Abfrageleistung von Speichersystemen. Man könnte sagen, dass der Kern dieses Forschungsgebiets die Verbindung zwischen den Veränderungen der untersten Hardwareebene und den Leistungsproblemen der Datenverwaltung ist.

Oasis: Sie haben gerade erwähnt, dass der Kern Ihrer Forschung die effizientere Datenverwaltung in Systemen unter verschiedenen Hardwarearchitekturen ist. Wie ist es Ihnen gelungen, von den traditionellen Datenbankoptimierungsproblemen zu Ihrem heutigen Forschungsgebiet in Bezug auf die Modellinferenz zu gelangen?

Dr. Lu: Die Entwicklung traditioneller Datenbanksysteme besteht eigentlich aus einer Reihe von Technologien, die ständig verbessert werden, um eine effizientere Datenabfrage und -speicherung zu ermöglichen. Das Kernziel ist einfach: wie man Abfragen schneller und stabiler macht. Beispielsweise dienen die Optimierung von Indexstrukturen und die Verbesserung von Caching-Mechanismen dazu, dass das System auch bei zunehmenden Datenmengen die Abfrageleistung aufrechterhalten kann.

Aber mit den Veränderungen der Hardwarearchitektur, wie etwa persistenten Arbeitsspeicher, getrennten Arbeitsspeicher oder hierarchischen Speichersystemen, gibt es viele Inkompatibilitäten zwischen diesen neuen Hardwarekomponenten und den traditionellen Modellen. Deshalb müssen wir uns erneut fragen: Wie sollte man das Datenbanksystem in dieser neuen Architektur gestalten, um das Potenzial der Hardware optimal auszunutzen?

Meine Forschung während des Promotionsstudiums hat sich hauptsächlich um diese Frage herumgedreht.

Einerseits habe ich mich auf die Leistungsoffizierung des Systems unter neuen Hardwarebedingungen konzentriert, wie etwa die Reduzierung der Zugriffsverzögerung und die Verbesserung der Indexeffizienz. Andererseits habe ich auch untersucht, wie man das System intelligenter gestalten kann, wenn es mit großen Aufgaben konfrontiert wird. Diese Erfahrungen haben auch meinen heutigen Forschungsrichtung beeinflusst.

Heute konzentriert sich unser Team hauptsächlich darauf, wie man die "Retrieval"-Logik aus traditionellen Datenbanksystemen auf die Modell-Ebene übertragen kann. Beispielsweise hat das Modell beim Inferieren über lange Kontexte ähnliche Engpässe wie eine Datenbank, nämlich wie man in einem großen Speicher schnell die wichtigen Informationen findet. Dies hat uns dazu gebracht, zu denken, dass man vielleicht die Retrieval-Idee aus der Datenbank mit dem Attention-Mechanismus des Modells verbinden kann.

Unser Ansatz besteht darin, bewährte Vektorretrievalmethoden aus der Datenbank in den Inferenzprozess von Sprachmodellen zu übertragen, damit das Modell nur auf die "relevantesten" Informationen zugreift, wenn es generiert. Mit anderen Worten, anstatt das Modell neu zu trainieren, möchten wir durch die Systemdesign das Modell effizienter nutzen lassen, um die vorhandenen Informationen zu nutzen. Dies ist eigentlich der Ausgangspunkt unserer heutigen Studie.

Oasis: Können Sie uns bitte näher erläutern, wie dieser Forschungsansatz, der die Hardware berücksichtigt, entstanden ist? Wie haben Sie damals auf der Hardware- oder Systemebene über diese Verbindung nachgedacht?

Dr. Lu: Dieser Gedanke ist eigentlich aus einem sehr konkreten Problem entstanden.

Wir haben festgestellt, dass der Attention-Mechanismus im Wesentlichen ein Vektorretrievalprozess ist. Wenn das Modell generiert, muss es ständig in den vorhandenen Kontexten nach relevanten Informationen "suchen". Dabei ist der Memory (d. h. der Grafikspeicher) die größte Einschränkung.

Der Memory heutiger GPU ist sehr begrenzt. Wir haben festgestellt, dass das Modell beim Inferieren über lange Kontexte ständig eine große Menge an KV-Caches anhäuft, und die Verwaltung dieser Caches ist eigentlich sehr ähnlich der Logik von "Index + Cache" in traditionellen Datenbanksystemen. Deshalb haben wir uns gefragt, ob der Attention-Mechanismus des Modells ähnliche Ideen aus der Datenbankadaptieren kann, da die Datenbank durch hierarchische Speicherung und Retrieval-Mechanismen riesige Datenmengen effizient verwalten kann.

Deshalb haben wir angefangen, einen Teil der Caches vom GPU auf einen größeren Speicher, wie etwa den CPU-Speicher, zu verlagern und dann die benötigten Teile dynamisch über die Retrievalmethode aufzurufen. Diese "Hierarchie + Retrieval"-Struktur ermöglicht es dem System, auch bei langen Kontexten effizient zu arbeiten, ohne zusätzliche Trainingskosten zu verursachen.

Oasis: Haben Sie also die Logik von "Speicherung - Indexierung - Retrieval" aus dem Datenbanksystem auf die Attention-Berechnung des Modells übertragen?

Dr. Lu: Ja, der Ansatz ist sehr ähnlich. Das Ziel der Datenbank ist es, den Datenzugriff schneller zu machen, während wir uns darauf konzentrieren, dass das Modell den Prozess des "Findens von Informationen" beim Generieren schneller macht. Der Unterschied besteht darin, dass wir uns mit unstrukturierten und dynamisch veränderlichen Kontextinformationen befassen, sodass die Retrieval- und Cache-Strategien neu gestaltet werden müssen.

Man könnte sagen, dass wir die Erfahrungen aus der Leistungsoffizierung traditioneller Datenbanksysteme erneut auf den Attention-Mechanismus des Modells anwenden, um ihn der "Speicherungs- und Abrufprozess" in der realen Welt näher zu bringen.

Dies war auch der ursprüngliche Ausgangspunkt von Retrieval Attention.

Oasis: Können wir vielleicht mit den grundlegenden Teilen beginnen, um den Lesern die unbekannten Konzepte zu erläutern? Was ist die Kernlogik des "Datenverwaltungsmechanismus" in einem Datenbankmanagementsystem (Database Management System)? Kann es als der Kern eines "großen Datenbanksystems" betrachtet werden?

Dr. Lu: Ich kann Ihnen kurz eine Erklärung geben. Im Kontext des maschinellen Lernens oder der Kognitionswissenschaft ist der ursprüngliche Ausgangspunkt des "Attention-Mechanismus" eigentlich eine Nachahmung der menschlichen Aufmerksamkeit. Das heißt, wenn das Modell eine Reihe von Informationen verarbeitet, wie es entscheidet, welche Inhalte wichtiger sind und welche ignoriert werden können.

Nehmen wir beispielsweise an, wir stellen dem Modell eine Frage. Es wird nicht alle Eingabedaten gleichermaßen betrachten, sondern es wird automatisch "Attention-Gewichte" zuweisen und sich auf diejenigen Teile konzentrieren, die für diese Frage relevanter sind. Dies ist der Kerngedanke des Attention-Mechanismus: In einer Sequenz wird der Rechenaufwand auf diejenigen wenigen Elemente konzentriert, die wirklich wichtig sind, basierend auf dem Aufgabenziel.

Frühe Sequenzmodelle, wie etwa RNN oder Seq2Seq, hatten bereits diese Eigenschaft der "selektiven Fokussierung", aber damals war sie stärker von der sequentiellen Struktur abhängig.

Mit der Entstehung von Attention kann das Modell nun Beziehungen zwischen beliebigen Positionen herstellen: Es ist nicht mehr von der Zeit oder der Position beschränkt, sondern kann direkt globale Abhängigkeiten erfassen. Dies ist auch der Grund, warum Transformer zu einer bahnbrechenden Architektur geworden ist.

Von der Implementierung her kann der Attention-Mechanismus als ein "Abfrage - Übereinstimmung - Gewichtung"-Prozess verstanden werden. Das Modell teilt die Eingabe in drei Arten von Vektoren auf: Query, Key und Value.

Query repräsentiert die Informationen, die derzeit verarbeitet werden müssen, d. h. "wonach ich gerade suche".

Key repräsentiert alle möglichen Übereinstimmungskriterien.

Value repräsentiert den Inhalt, der diesen Kriterien entspricht.

Das Modell berechnet die Ähnlichkeit zwischen Query und jedem Key, um eine Reihe von Gewichten zu erhalten, und verwendet dann diese Gewichte, um alle Values zu gewichtet zu summieren. Das Ergebnis ist der Teil, den das Modell als "am wichtigsten" ansieht, d. h. die endgültige Attention-Ausgabe.

Wenn wir den Attention-Mechanismus in den Rahmen eines Datenbanksystems betrachten, kann er eigentlich als ein dynamisches Informationsabrufsystem verstanden werden. Jedes Mal, wenn das Modell ein neues Token generiert, muss es in dem vorhandenen semantischen Raum nach den relevantesten Informationen "suchen", was dem Prozess einer Datenbank bei der Ausführung einer Abfrageanfrage sehr ähnlich ist.

In dieser Analogie repräsentieren die Attention-Gewichte das "Übereinstimmungsmaß" einer Abfrage. Je höher das Gewicht ist, desto mehr sollte das aktuell generierte Token auf diesen Teil der Eingabe "achten",