Retrofitting Linear Attention into Diffusion Language Models

Stellen Sie sich die Welt der künstlichen Intelligenz als eine riesige, belebte Bibliothek vor, in der ein superintelligenter Roboter-Bibliothekar versucht, eine Geschichte zu schreiben. Jahrelang war die Standardmethode, mit der dieser Bibliothekar arbeitete, „autoregressiv“: Er schrieb ein Wort, hielt inne, dachte über alles nach, was er gerade geschrieben hatte, schrieb das nächste Wort, hielt inne und dachte wieder nach. Es war, als würde man einen Turm Stein für Stein bauen und warten, bis der Kleber getrocknet ist, bevor man den nächsten Stein hinzufügt. Während diese Methode zuverlässig war, war sie besonders bei langen Geschichten langsam.

Vor Kurzem kam ein neuer Stil namens „Diffusion“ auf. Anstatt Stein für Stein zu bauen, beginnt der Diffusion-Bibliothekar mit einer leeren Seite voller Fragezeichen und versucht, die ganze Geschichte auf einmal zu erraten, indem er die Vermutungen immer wieder verfeinert, bis die Wörter Sinn ergeben. Dies ist, als hätte man ein Team von Künstlern, die gleichzeitig ein Wandgemälde malen; sie können viel schneller arbeiten, weil sie nicht darauf warten müssen, dass der vorherige Stein trocknet. Doch selbst diese schnelle Methode hat einen Haken. Wenn das Wandgemälde größer wird, muss das Team ständig auf jedes einzelne Wort zurückblicken, das sie bereits festgelegt haben, um sicherzustellen, dass die neuen Wörter passen. Es ist, als müsste man versuchen, ein 10.000-seitiges Buch im Kopf zu behalten, während man Seite 101 schreibt, was schließlich alle ausbremst und die Kapazität der Bibliothek füllt.

Dies ist das Rätsel, das ein Forscherteam gelöst hat. Sie fragten: „Können wir diese schnelle Diffusion-Methode noch schneller machen, indem wir ändern, wie der Bibliothekar sich an die Vergangenheit erinnert?“ Ihre Antwort war ein kluger Trick namens „Block-Hybrid-Attention“. Sie erkannten, dass der Bibliothekar zwar den aktuellen Absatz perfekt im Blick behalten muss (um den Fluss der Sätze zu gewährleisten), er aber nicht jedes einzelne Wort der vorherigen Kapitel in Hochauflösung neu lesen muss. Stattdessen kann er die alten Kapitel in einem winzigen, fest definierten „Zusammenfassungsblatt“ zusammenfassen. Dies ermöglicht es dem Bibliothekar, die neuen Teile der Geschichte mit Lichtgeschwindigkeit zu schreiben, ohne durch das Gewicht der gesamten Geschichte ausgebremst zu werden.

Die große Idee: Ein Zwei-Geschwindigkeiten-Speichersystem

Das Paper stellt eine Methode vor, um ein bestehendes, leistungsstarkes Diffusion-Sprachmodell (speziell ein 16-Milliarden-Parameter-Modell namens LLaDA 2.1) aufzurüsten (oder zu verbessern), damit es diese „Zusammenfassungsblatt“-Strategie nutzen kann, ohne von Grund auf neu trainiert werden zu müssen.

Stellen Sie sich das Gehirn des Modells wie ein System aus 20 verschiedenen Denkschichten vor. Im Originalmodell agiert jede einzelne dieser 20 Schichten wie ein superaufmerksamer Bibliothekar, der jedes zuvor geschriebene Wort liest, um das aktuelle zu verstehen. Das ist präzise, aber schwerfällig. Die Innovation der Forscher, die Block-Hybrid-Attention, teilt die Arbeit in zwei Modi auf:

1. Der „Jetzt“-Modus (Exakte Attention): Wenn das Modell an dem aktuellen Block von Wörtern arbeitet, den es gerade aktiv generiert, behält es seinen „superaufmerksamen“ Modus bei. Es nutzt das Standard-Speicherverfahren mit hoher Kapazität, um jedes Wort im aktuellen Absatz zu betrachten, damit die Sätze einen perfekten Sinn ergeben.
2. Der „Vergangenheit“-Modus (Lineare Attention): Wenn das Modell sich an die vorherigen Blöcke erinnern muss (die fertigen Kapitel), wechselt es zur „Linearen Attention“. Anstatt das ganze Buch neu zu lesen, konsultiert es einen fest definierten Zusammenfassungszustand. Das ist vergleichbar mit dem Blick in ein Register oder eine einseitige Zusammenfassung eines Buches, anstatt den gesamten Text erneut zu lesen. Diese Zusammenfassung behält ihre Größe bei, egal wie lang das Buch wird.

Wie sie es gemacht haben: Das Zwei-Stufen-Upgrade

Die Forscher haben das Modell nicht einfach nur mit neuen Teilen bestückt und gehofft, dass es funktioniert; sie nutzten einen sorgfältigen Zwei-Stufen-Prozess, um sicherzustellen, dass das Modell während des Upgrades nicht seine Intelligenz verliert.

• Stufe 1: Das Schatten-Training. Sie nahmen das ursprüngliche, intelligente Modell (den „Lehrer“) und frierten es ein. Dann ersetzten sie 6 der 20 Attention-Schichten durch ihre neue „Block-Hybrid“-Version (den „Schüler“). Der Schüler wurde darauf trainiert, die Ausgabe des Lehrers exakt zu imitieren, wobei eine „korrumpierte“ Version des Textes als Input diente. Es war, als würde ein Schüler einem Meisterkoch Schatten stellen, versucht dessen exakte Bewegungen zu kopieren und dieselben Aromen zu schmecken, aber nur für die spezifischen Gerichte, für die der Schüler zuständig ist. Diese Phase dauerte etwa 24 Stunden.
• Stufe 2: Das Fine-Tuning. Nachdem der Schüler gelernt hatte, den Lehrer zu imitieren, ließen die Forscher das gesamte Modell wieder zusammenarbeiten. Sie verwendeten eine Technik namens LoRA (Low-Rank Adaptation), um winzige, präzise Anpassungen an den Verbindungen des Modells vorzunehmen. Dies war wie ein schnelles „Polieren“ des aufgerüsteten Modells, um kleine Unstimmigkeiten zu beheben, die entstanden, weil die neuen Teile nun mit den alten Teilen kommunizierten. Diese Phase dauerte etwa 6 Stunden.

Die Ergebnisse: Schneller, leichter und dennoch klug

Das Team testete das aufgerüstete Modell, das sie LLaDA-HYBRID nannten, um zu sehen, ob es immer noch gut im Schreiben und im Lösen von Problemen ist und ob es tatsächlich schneller ist.

Kann es immer noch denken?
Überraschenderweise ja. Obwohl sie 6 von 20 Schichten mit dieser „Zusammenfassungsblatt“-Methode ersetzt haben, blieb die Leistung des Modells der des Originals sehr nahe.

• Bei einem Coding-Test namens HumanEval erreichte das Original 75,6 %, während das Hybrid-Modell 72,0 % erreichte.
• Bei einem Mathe-Test namens CMATH erreichte das Original 88,3 %, das Hybrid-Modell hingegen 86,7 %.
• Interessanterweise verbesserte sich das Hybrid-Modell bei einem anderen Coding-Test, MBPP+, von 57,7 % auf 63,0 %.
  Das Paper legt nahe, dass das Modell bei den schwierigsten Denkaufgaben (wie einem schwierigen Wissenschaftsquiz namens GPQA-Diamond, bei dem es von 38,9 % auf 30,8 % sank) zwar etwas weniger perfekt ist, aber seine Fähigkeit zum Programmieren und zum Lösen von mathematischen Problemen weitgehend bewahrt hat.

Ist es tatsächlich schneller?
Hier geschieht die Magie. Da das Modell nicht mehr für jedes neue Wort die gesamte Historie des Gesprächs neu lesen muss, kann es mehr Anfragen gleichzeitig bearbeiten.

• Bei der Testung, wie viele Wörter das Modell pro Sekunde generieren kann, war das Hybrid-Modell bei seiner Spitzenleistung (beim Umgang mit 128 gleichzeitigen Anfragen) 1,7-mal schneller als das Original.
• Der Speicherverbrauch des Originalmodells wuchs mit der Länge der Geschichte an, bis es an eine Grenze stieß, an der es nicht mehr mehr Nutzer bewältigen konnte. Das Hybrid-Modell konnte dank seines fest definierten „Zusammenfassungsblatts“ mehr gleichzeitige Anfragen bearbeiten, bevor der Speicher voll war.

Das Fazit

Dieses Paper legt nahe, dass wir nicht immer eine neue, schnellere KI von Grund auf neu bauen müssen. Stattdessen können wir eine leistungsstarke, bestehende KI nehmen und ihr ein „hybrides Gedächtnis-System“ geben. Indem wir den scharfen Fokus auf den gegenwärtigen Moment behalten und gleichzeitig die Vergangenheit in eine kompakte Form zusammenfassen, können wir diese Modelle signifikant schneller und effizienter machen. Die Forscher fanden heraus, dass dieses Upgrade mit etwa 60 Stunden Training auf öffentlichen Daten durchgeführt werden konnte, was einen vielversprechenden Weg aufzeigt, KI-Dienste schneller und kostengünstiger zu machen, ohne zu viel von ihrer Intelligenz einzubüßen. Die Autoren merken jedoch an, dass dies erst der Anfang ist; sie haben nur 6 Schichten aufgerüstet, und zukünftige Arbeiten könnten Wege finden, noch mehr Schichten aufzurüsten oder noch längere Geschichten zu verarbeiten.