PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding
Das Papier stellt PRESTO vor, ein fundiertes Framework, das das diffusionsbasierte spekulative Dekodieren durch die Implementierung von präfix-ausgerichteter Bewertung und prioritätsbasierter Baumsuche verbessert, um die Diskrepanz zwischen Diffusions-Marginalen und autoregressiver Verifizierung zu lösen und dadurch den End-to-End-Durchsatz signifikant zu steigern.
Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen
Stellen Sie sich vor, Sie versuchen, das nächste Wort in einer Geschichte vorherzusagen. Lange Zeit haben die klügsten Computer (genannt Large Language Models) dies Wort für Wort getan, wie ein Mensch, der ein Buch laut vorliest, und nach jedem einzelnen Wort innehält, um darüber nachzudenken, was als Nächstes kommt. Das ist genau, aber langsam. Kürzlich haben Wissenschaftler einen neuen Weg entdeckt, diese Geschichten mithilfe von „Diffusionsmodellen“ zu schreiben. Denken Sie dabei an einen Bildhauer, der mit einem Marmorblock beginnt und die ganze Statue auf einmal ausarbeitet, anstatt immer nur ein winziges Stückchen nach dem anderen abzutragen. Dies ermöglicht es dem Computer, viele Wörter gleichzeitig zu erraten, was unglaublich schnell ist.
Es gibt jedoch einen Haken. Wenn Sie viele Wörter gleichzeitig erraten, können Sie ein paar falsch machen. Um dies zu beheben, gibt es einen cleveren Trick namens „Speculative Decoding“. Es ist wie ein schneller, juniorartiger Assistent, der die nächsten paar Wörter errät, und dann ein langsamer, super-intelligenter Chef, der überprüft, ob diese Vermutungen richtig sind. Wenn der Chef den Vorschlägen zustimmt, akzeptiert er den ganzen Stapel an Wörtern sofort, was eine Menge Zeit spart. Das Problem ist, dass der juniorartige Assistent (das Diffusionsmodell) zwar gut darin ist, einzelne Wörter zu erraten, aber nicht immer weiß, wie diese Wörter in einer bestimmten Reihenfolge zusammenpassen. Es ist, als würde der Assistent hervorragend darin sein, einzelne Zutaten für einen Kuchen auszuwählen, aber er wüsste nicht immer, welche Kombination von Zutaten tatsächlich gut schmecken würde.
Hier kommt ein neues Paper ins Spiel. Die Forscher, unter der Leitung von Zheng Wang und Kollegen, erkannten, dass die aktuelle Art und Weise, diese schnellen Assistenten zu verwenden, viel Geschwindigkeit ungenutzt lässt. Sie fanden heraus, dass der Assistent zwar eine riesige Vielfalt an möglichen Wortkombinationen generieren kann, die aktuelle Methode jedoch nur einen einzigen Pfad überprüft, als würde man einen einzigen Flur entlanggehen und hoffen, dass die Tür am Ende offen ist. Die Autoren schlagen ein neues System namens PRESTO (Prefix-Aligned Tree Drafting) vor. Anstatt nur einen Flur entlangzugehen, baut PRESTO einen Baum von Möglichkeiten auf und erkundet viele verschiedene Pfade gleichzeitig. Aber hier liegt die Magie: Es behebt einen grundlegenden Fehler in der Art und Weise, wie das Vertrauen des Assistenten gemessen wird. Das ursprüngliche Vertrauen des Assistenten ist „präfix-blind“, was bedeutet, dass es ihm egal ist, welche Wörter zuvor kamen. PRESTO fügt einen „präfix-ausgerichteten“ Score hinzu, der wie ein Kompass wirkt und sicherstellt, dass die für die Erkundung gewählten Pfade diejenigen sind, die am wahrscheinlichsten vom Chef akzeptiert werden.
Das Ergebnis ist ein signifikant schnelleres System. In ihren Tests half PRESTO dem Computer, mehr Wörter in jeder Runde des Ratens zu akzeptieren. Bei einigen der besten bestehenden Setups machte es den gesamten Prozess 1,5-mal schneller. Bei anderen Setups ermöglichte es eine 1,12-fache Beschleunigung. Das Paper legt nahe, dass, indem wir den Rateprozess wie ein baumartiges Abenteuer statt wie eine gerade Linie behandeln und sicherstellen, dass die Äste danach ausgewählt werden, wie gut sie zur bisherigen Geschichte passen, wir das Beste aus beiden Welten bekommen: die Geschwindigkeit der Diffusion und die Genauigkeit sorgfältiger Überprüfung.
Das Problem: Die „Ein-Pfad“-Falle
Um zu verstehen, warum PRESTO benötigt wird, stellen Sie sich vor, Sie spielen ein Spiel wie „Mad Libs“ mit einem Freund, der versucht, die fehlenden Wörter zu erraten. Ihr Freund ist ein Diffusionsmodell. Er ist großartig darin, in eine Lücke zu schauen und zu sagen: „Ich wette, das Wort hier ist ‚Katze‘!“ oder „Vielleicht ist es ‚Hund‘?“ oder „Oder vielleicht ‚Rakete‘?“. Er kann all diese Optionen zur exakt gleichen Zeit ausrufen.
Die aktuelle Methode, diesen Freund zu verwenden, ist jedoch sehr starr. Sie nimmt die Top-Vermutung des Freundes, schreibt sie auf und fragt dann den „Chef“ (das Zielmodell), ob diese Vermutung richtig ist. Wenn der Chef sagt „Nein“, wird das Ganze verworfen und man muss von vorne beginnen. Wenn der Chef „Ja“ sagt, geht man zum nächsten Wort über und wiederholt den Vorgang. Dies wird lineares Drafting genannt. Es ist, als würde man durch einen Wald wandern und nur den Pfad betrachten, der direkt vor einem liegt.
Die Autoren beobachteten, dass dieser Ansatz ineffizient ist. Da das Diffusionsmodell viele Optionen gleichzeitig generiert, gibt es einen massiven „kombinatorischen Raum“ an Möglichkeiten. Es ist, als hätte man eine Karte mit tausend verschiedenen Pfaden, aber man darf nur auf einem einzigen wandern. Das Paper zeigt, dass man durch das Festhalten an nur einem Pfad viele gültige Routen verpasst. Tatsächlich akzeptierte die aktuelle Methode bei mathematischen Problemen wie GSM8K durchschnittlich etwa 6,5 Wörter, aber die Forscher berechneten, dass sie fast 10 Wörter akzeptiert hätten, wenn sie alle besten Pfade prüfen könnten. Das ist eine riesige Lücke!
Der Mismatch: Der „blinde“ Kompass
Die Forscher gruben tiefer und fanden einen spezifischen Grund, warum das einfache Prüfen von mehr Pfaden (das Bauen eines Baums) mit den alten Methoden nicht perfekt funktionierte. Sie identifizierten einen „fundamentalen Mismatch“.
In der Welt der Standard-KI (autoregressive Modelle) hängt der Konfidenz-Score für ein Wort stark von den Wörtern ab, die vor ihm kamen. Wenn der Satz lautet: „Die Katze saß auf der...“, weiß das Modell, dass „Matte“ ein sehr wahrscheinliches nächstes Wort ist, aber „Pizza“ nicht. Dies ist präfix-ausgerichtet.
Aber Diffusionsmodelle funktionieren anders. Sie generieren eine „marginale“ Wahrscheinlichkeit für jede Position unabhängig vone-inander. Es ist, als würde das Modell sagen: „An Position 5 ist ‚Katze‘ zu 80 % wahrscheinlich“, ohne zu berücksichtigen, ob an Position 4 „Die“ oder „Die schnelle braune“ stand. Dies ist präfix-blind.
Wenn man versucht, einen Baum aus Vermutungen mit diesen blinden Scores aufzubauen, entsteht ein Ranking-Problem. Man wählt vielleicht einen Pfad, der für das erste Wort großartig aussieht, aber für das zweite Wort schrecklich ist, weil das Modell nicht realisiert hat, dass das erste Wort den Kontext verändert hat. Es ist wie ein GPS, das Ihnen Wegbeschreibungen basierend nur auf der aktuellen Straße gibt, ohne zu berücksichtigen, dass Sie gerade links abgebogen sind und sich nun in einer Einbahnstraße befinden. Das Paper argumentiert, dass die Verwendung dieser blinden Scores zum Aufbau eines Baums zu einem „unzuverlässigen Pfad-Ranking“ führt, was bedeutet, dass das System die falschen Zweige erkundet und Zeit verschwendet.
Die Lösung: PRESTO
PRESTO (Prefix-Aligned Scoring and priority-based Tree search for diffusion Speculative decOding) löst dies, indem es einen „Korrekturfaktor“ zum Score des Diffusionsmodells hinzufügt.
- Präfix-ausgerichtetes Scoring: Die Autoren erkannten, dass sie das starke „marginale“ Signal des Diffusionsmodells (wie wahrscheinlich ein Wort für sich allein ist) mit einem „präfix-konditionierten“ Signal (wie wahrscheinlich es angesichts der vorherigen Wörter ist) kombinieren müssen. Sie entwickelten eine neue Scoring-Formel, die die Diffusionswahrscheinlichkeit mit einem Korrekturfaktor multipliziert, der aus einem einfachen n-Gramm-Modell (einem leichtgewichtigen Werkzeug, das Wortkombinationen betrachtet) abgeleitet wird. Dies erzeugt einen Score, der dem Fluss der Geschichte respektiert.
- Prioritätsbasierte Baumsuche: Anstatt nur den Top-Pfad zu wählen, baut PRESTO einen Baum auf. Es verwendet die neuen, korrigierten Scores, um zu entscheiden, welche Zweige wachsen. Es priorisiert Pfade, die das höchste Potenzial haben, vom Chef akzeptiert zu werden. Es ist wie ein Wanderer, der, anstatt einfach geradeaus zu gehen, eine Karte betrachtet und den Pfad wählt, der am wahrscheinlichsten zum Gipfel führt, selbst wenn dieser Pfad am Anfang nicht der offensichtlichste ist.
Das Paper testete zwei Arten, diesen Baum wachsen zu lassen: Beam Search (Beibehalten einer festen Anzahl von Top-Pfaden bei jedem Schritt) und Best-First Search (immer den bisher besten gefundenen Pfad erweitern). Sie fanden heraus, dass für ihren spezifischen Aufbau Beam Search genauso gut wie die komplexere Best-First Search funktionierte, also hielten sie sich für die einfachere, effizientere Option.
Die Ergebnisse: Schneller und Schlauer
Die Autoren setzten PRESTO bei verschiedenen Aufgaben ein, darunter mathematische Probleme (GSM8K, Math500), Programmierherausforderungen (HumanEval, LiveCodeBench) und Chat-Gespräche. Sie verwendeten zwei verschiedene Arten von Systemen:
- Dedizierte Diffusions-Drafts: Ein kleines, schnelles Diffusionsmodell, das für ein größeres autoregressives Modell rät (wie dFlash).
- Self-Speculative Diffusion LLMs: Ein einzelnes Diffusionsmodell, das selbst rät und selbst prüft (wie Nemotron-Labs-Diffusion).
Die Ergebnisse waren konsistent über alle Bereiche hinweg. PRESTO erhöhte durchweg die durchschnittliche Akzeptanzlänge, also die Anzahl der Wörter, die der Chef in einem Durchgang akzeptiert.
- Beim dFlash-System (unter Verwendung von Qwen3-8B) sprang die durchschnittliche Akzeptanzlänge von etwa 6,6 Wörtern auf 9,6 Wörter. Dies übersetzte sich in eine End-to-End-Beschleunigung um das 1,5-fache.
- Beim Nemotron-Labs-Diffusion-System erhöhte sich die Akzeptanzlänge von 8,8 auf 9,9 Wörter, was eine 1,12-fache Beschleunigung zur Folge hatte.
Besonders beeindruckend war, dass das Paper zeigte, dass PRESTO auch funktioniert, wenn das System „stochastisch“ (randomisiert) arbeitet, was normalerweise schwieriger vorherzusagen ist. In diesen Fällen war die Beschleunigung sogar noch ausgeprägter, wobei einige Benchmarks fast die doppelte Durchsatzrate erreichten.
Die Autoren überprüften auch die „Kosten“ dieser neuen Methode. Sie fanden heraus, dass die zusätzliche Arbeit, die zum Aufbau des Baums und zur Berechnung der neuen Scores erforderlich war, minimal war – weniger als 4 % der Gesamtzeit. Der weitaus größte Teil der Zeit (über 90 %) wurde immer noch für die eigentliche Verifizierung durch das Chefmodell aufgewendet. Das bedeutet, dass PRESTO ein hocheffizientes Upgrade ist, das das System nicht durch Overhead verlangsamt.
Was PRESTO NICHT ist
Es ist wichtig zu beachten, was das Paper nicht behauptet. Die Autoren stellen explizit klar, dass das bloße Anwendung einer naiven Baumstruktur (oh ohne ihr präfix-ausgerichtetes Scoring) suboptimal ist. Wenn man einfach die Roh-Scores des Diffusionsmodells nimmt und einen Baum aufbaut, erhält man nicht den vollen Nutzen, da die Scores „blind“ sind. PRESTO bezieht sich spezifisch darauf, diesen Scoring-Mismatch zu beheben.
Darüber hinaus behauptet das Paper nicht, das Problem der Diffusionsmodelle vollständig gelöst zu haben. Sie räumen ein, dass ihre Methode auf ein „berechenbares präfix-ausgerichtetes Signal“ (wie das n-Gramm-Modell, das sie verwendeten) angewiesen ist, um die Hauptarbeit zu leisten. Sie schlagen vor, dass zukünftige Arbeiten noch reichhaltigere Signale untersuchen könnten, aber für den Moment reicht ihre einfache Korrektur aus, um massive Gewinne zu erzielen.
Warum das wichtig ist
Im Wettlauf darum, KI schneller und effizienter zu machen, zählt jedes bisschen Geschwindigkeit. Speculative Decoding ist ein heiß diskutiertes Thema, weil es uns erlaubt, die besten, genauesten Modelle zu nutzen, ohne den vollen Zeitaufwand zu bezahlen. Die aktuellen Methoden waren jedoch dadurch begrenzt, dass sie Diffusionsmodelle wie lineare Maschinen behandelten und dabei ihre einzigartige Fähigkeit ignorierten, viele Optionen gleichzeitig zu generieren.
PRESTO ändert das Spiel, indem es Diffusionsmodelle als die vielpfadigen Entdecker behandelt, die sie sind. Indem es das Scoring an die Art und Weise ausrichtet, wie der Chef die Arbeit prüft, schaltet es das volle Potenzial der parallelen Generierung des Diffusionsmodells frei. Das Ergebnis ist ein System, das nicht nur etwas schneller ist, sondern signifikant effizienter, was es ermöglicht, Text, mathematische Probleme und Code mit Geschwindigkeiten zu generieren, die für diese Arten von Modellen zuvor als unmöglich galten. Wie die Autoren sagen, haben sie einen „Einzelpfad-Spaziergang“ in eine „baumartige Expedition“ verwandelt und damit sichergestellt, dass jeder Schritt ein Schritt in Richtung der richtigen Antwort ist.
Ertrinken Sie in Arbeiten in Ihrem Fachgebiet?
Erhalten Sie tägliche Digests der neuesten Arbeiten passend zu Ihren Forschungsbegriffen — mit technischen Zusammenfassungen, in Ihrer Sprache.