💬 NLP

FlexDraft: Flexible Speculative Decoding via Attention Tuning and Bonus-Guided Calibration

FlexDraft ist ein verlustfreies Framework für spekulatives Decodieren, das die Inferenz von LLMs über unterschiedliche Batch-Größen hinweg beschleunigt, indem es Attention-Tuning für eine hochwertige Blockdiffusions-Drafting, bonusgesteuerte Kalibrierung zur Auflösung von Diskrepanzen bei der Draft-Verifizierung und einen dynamischen Umschaltmechanismus kombiniert, der zwischen parallelen und sequenziellen Ausführungsmodi optimiert.

Ursprüngliche Autoren: Yaojie Zhang, Jianuo Huang, Junlong Ke, Yuhang Han, Yongji Long, Tianchen Zhao, Biqing Qi, Linfeng Zhang

Veröffentlicht 2026-05-20

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Yaojie Zhang, Jianuo Huang, Junlong Ke, Yuhang Han, Yongji Long, Tianchen Zhao, Biqing Qi, Linfeng Zhang

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, eine lange Geschichte zu schreiben, aber Sie haben einen sehr strengen Redakteur (das Zielmodell), der unglaublich klug ist, aber sich sehr langsam bewegt, weil er jedes einzelne Wort überprüfen muss, bevor er das nächste schreibt. Dadurch dauert das Schreiben der Geschichte ewig.

Um die Dinge zu beschleunigen, stellen Sie einen schnellen, energiegeladenen Assistenten (den Entwerfer) ein, der für Sie die nächsten paar Wörter errät. Der Assistent schreibt einen ganzen Absatz im Voraus vor, und dann überprüft der langsame Redakteur schnell, ob diese Vermutungen richtig waren. Wenn sie richtig sind, können Sie die langsame Denkzeit des Redakteurs für diese Wörter überspringen. Dies wird als Spekulatives Dekodieren bezeichnet.

Die alte Methode, dies zu tun, hatte jedoch zwei große Probleme:

Das Warten: Der Assistent würde eine Vermutung schreiben und dann warten, bis der Redakteur fertig war, sie zu überprüfen, bevor er die nächste Vermutung schrieb. Sie machten abwechselnd weiter, was Zeit verschwendete.
Die „Was-wäre-wenn"-Verwirrung: Bei neueren, schnelleren Methoden versucht der Assistant, mehrere verschiedene zukünftige Absätze gleichzeitig zu schreiben, in der Hoffnung, alle Möglichkeiten abzudecken. Aber der Assistent weiß nicht, welchen Absatz der Redakteur tatsächlich genehmigen wird. Dies führt dazu, dass der Assistent Energie damit verschwendet, Absätze zu schreiben, die verworfen werden, und der Redakteur zu viele Optionen überprüfen muss, was alles wieder verlangsamt, wenn Sie viele Geschichten gleichzeitig schreiben müssen (große Batch-Größen).

Hier kommt FlexDraft ins Spiel: Der intelligente, flexible Assistent

Die Arbeit stellt FlexDraft vor, ein neues System, das diese Probleme behebt, damit Sie schneller schreiben können, ohne an Qualität zu verlieren. So funktioniert es, unter Verwendung einfacher Analogien:

1. Die „Spezialbrille" (Attention-Tuning)

Normalerweise müssen Sie, um den Assistenten besser im Erraten zu machen, sein gesamtes Gehirn neu trainieren, was teuer und riskant ist (es könnte dazu führen, dass er vergisst, wie man richtig spricht).

FlexDraft ist anders. Es setzt dem Assistenten nur für die sehr letzten Schritte des Denkens eine Paar „Spezialbrillen" auf die Augen.

Wie es funktioniert: Der Assistent behält sein gesamtes ursprüngliches Wissen (das eingefrorene Gehirn), lernt aber, diese neuen Brillen zu verwenden, um auf „leere Stellen" (Mask-Token) zu schauen und sie alle auf einmal auszufüllen.
Der Vorteil: Der Assistent wird zu einem schnellen, parallelen Rater, ohne dass eine vollständige Gehirntransplantation nötig ist. Er bleibt dem Stil des Redakteurs treu, sodass die Vermutungen von hoher Qualität und sicher sind.

2. Die „Magische Notiz" (Bonus-geführte Kalibrierung)

Hier kommt der knifflige Teil: Wenn der Redakteur die Vermutungen des Assistenten überprüft, sagt der Redakteur manchmal: „Nein, dieses Wort ist falsch, aber hier ist das korrekte Wort, mit dem stattdessen begonnen werden soll." Dieses korrekte Wort wird als Bonus-Token bezeichnet.

Bei den alten „schnellen" Methoden musste der Assistent den nächsten Absatz bevor er diese Magische Notiz sah, erraten. Also könnte der Assistent eine Geschichte über einen Piraten erraten, während der Redakteur eine Geschichte über einen Koch wollte. Die Vermutung des Assistenten war falsch, weil er die Notiz nicht hatte.

FlexDraft fügt ein Magische-Notiz-System hinzu:

Wie es funktioniert: Sobald der Redakteur das Bonus-Token enthüllt, verwendet FlexDraft einen winzigen, schnellen Rechner, um die vorherigen Vermutungen des Assistenten sofort an diese neue Richtung anzupassen.
Der Vorteil: Selbst wenn der Assistent zunächst blind geraten hat, „steuert" das System die Vermutung schnell so aus, dass sie mit der Korrektur des Redakteurs übereinstimmt. Das bedeutet, dass weniger Vermutungen abgelehnt werden.

3. Der „Verkehrspolizist" (Flex-Decoding)

Das größte Problem bei den alten „schnellen" Methoden war, dass sie versuchten, jede mögliche zukünftige Geschichte gleichzeitig zu schreiben. Wenn Sie nur eine Geschichte zu schreiben haben (kleiner Batch), ist das in Ordnung. Aber wenn Sie 16 Geschichten gleichzeitig schreiben müssen (großer Batch), wird der Assistent überfordert, wenn er versucht, 16 verschiedene Zukünfte für jede einzelne Geschichte zu schreiben, und der Redakteur gerät ins Stocken, wenn er sie alle überprüft.

FlexDraft fungiert wie ein intelligenter Verkehrspolizist:

Leichter Verkehr (kleiner Batch): Wenn es nur wenige Geschichten gibt, sagt der Polizist: „Los! Schreiben Sie alle möglichen Zukünfte gleichzeitig!" Dies überlappt das Schreiben und Überprüfen, um Zeit zu sparen.
Schwerer Verkehr (großer Batch): Wenn es viele Geschichten gibt, sagt der Polizist: „Stopp! Schreiben Sie nicht alle Möglichkeiten. Schreiben Sie nur die eine, die am wahrscheinlichsten genehmigt wird."
Der Vorteil: Es schaltet automatisch die Strategien um. Es vermeidet die „verschwendete Energie", zu viele Optionen zu schreiben, wenn das System beschäftigt ist, und verhindert, dass die Geschwindigkeit einbricht.

Das Ergebnis

Durch die Kombination dieser drei Tricks ermöglicht FlexDraft dem langsamen Redakteur, Text viel schneller zu verarbeiten.

Kein Qualitätsverlust: Die endgültige Geschichte ist genau dieselbe, als hätte der langsame Redakteur sie Wort für Wort geschrieben.
Riesige Beschleunigung: Bei Tests mit einem beliebten KI-Modell (Qwen3-8B) machte FlexDraft das System 4,59-mal schneller als die Standard-langsame Methode.
Skalierbarkeit: Es funktioniert gut, egal ob Sie eine Geschichte schreiben oder eine große Menge an Geschichten verwalten, im Gegensatz zu früheren Methoden, die langsam und unhandlich wurden, wenn die Menge groß wurde.

Kurz gesagt ist FlexDraft eine flexible, verlustfreie Möglichkeit, einem schnellen Assistenten zu helfen, einem langsamen Redakteur zu assistieren, und stellt sicher, dass sie reibungslos zusammenarbeiten, ohne Zeit zu verschwenden oder Fehler zu machen.

Technische Zusammenfassung: FlexDraft

Problemstellung

Spekulative Dekodierung beschleunigt die Inferenz von Large Language Models (LLM), indem ein leichtgewichtiger Drafter mehrere Kandidatentokens vorschlägt, die parallel von einem Zielmodell verifiziert werden. Obwohl effektiv, leiden herkömmliche sequenzielle spekulative Dekodierungsverfahren unter gegenseitigen Wartezeiten zwischen den Phasen des Draftings und der Verifizierung sowie unter einem hohen Overhead beim Speicherzugriff aufgrund wiederholter Zustandsaustausche.

Neuere parallele spekulative Dekodierungsverfahren versuchen, das Drafting und die Verifizierung innerhalb eines einzigen Forward-Passes des Zielmodells zu überlappen, um gegenseitige Wartezeiten zu minimieren. Dieses Paradigma führt jedoch zu zwei grundlegenden Einschränkungen, die die Leistung beeinträchtigen, insbesondere bei großen Batch-Größen:

Unsicherheit bezüglich des Bonus-Tokens: Beim parallelen Dekodieren werden Draft-Tokens generiert, bevor das Ergebnis der Verifizierung bekannt ist (insbesondere das „Bonus-Token", das am ersten Ablehnungspunkt gesampelt wird). Dies führt zu einer Diskrepanz zwischen der Draft-Verteilung und den Verifizierungsbedingungen, was die Akzeptanzraten verringert.
Unsicherheit bezüglich der Akzeptanzlänge: Da die Anzahl der akzeptierten Tokens während des Draftings unbekannt ist, müssen parallele Methoden Kandidatenverzweigungen für alle möglichen akzeptierten Präfixe vorbereiten. Dies zwingt das Zielmodell, eine schnell wachsende Anzahl von „Mask-Tokens" (zukünftige Drafts) zusammen mit den Verifizierungstokens zu verarbeiten, wodurch sich der Rechenaufwand quadratisch ( $O(N^2)$ ) mit der Länge des Draft-Blocks skaliert. Diese Redundanz führt dazu, dass die Durchsatzgewinne bei größeren Batch-Größen zusammenbrechen.

Methodik: FlexDraft

FlexDraft ist ein verlustfreies Rahmenwerk für spekulative Dekodierung, das entwickelt wurde, um sich flexibel an variierende Batch-Größen anzupassen und gleichzeitig die Verteilung des Zielmodells zu bewahren. Es adressiert die oben genannten Einschränkungen durch drei Kernkomponenten:

1. Attention-Tuning (Attn Tuning)

Um ein Block-Diffusions-Drafting (Generierung mehrerer Tokens parallel) zu ermöglichen, ohne das gesamte Modell neu zu trainieren oder die Qualität zu beeinträchtigen:

Architektur: Das Rahmenwerk hält den autoregressiven Pfad des Zielmodells (FFN-Schichten und ursprüngliche Attention-Projektoren) eingefroren.
Trainierbare Parameter: Nur die Attention-Projektoren ( $Q, K, V, O$ ) für die letzten paar Schichten werden speziell für Mask-Tokens angepasst.
Mechanismus: Mask-Tokens werden an das saubere Präfix angehängt. In tiefen Schichten achten diese Mask-Tokens bidirektional innerhalb ihres Blocks und auf das saubere Präfix, bleiben jedoch von anderen Blöcken isoliert. Dies ermöglicht es dem Drafter, das vortrainierte Wissen des Zielmodells über die eingefrorenen FFNs zu erben, während er lernt, mehrere zukünftige Tokens parallel mit einer minimalen Anzahl zusätzlicher Parameter (ca. 6 % der Größe des Zielmodells) vorherzusagen.

2. Bonus-gesteuerte Kalibrierung

Um die durch die Unsicherheit des Bonus-Tokens verursachte Diskrepanz zu beheben:

Mechanismus: Ein leichtgewichtiges Multi-Layer Perceptron (MLP) wird nach dem Language-Modeling-(LM)-Kopf angewendet.
Eingabe: Das MLP nimmt die Konkatenation der aufgelösten Embedding des Bonus-Tokens und des versteckten Zustands der aktuellen Mask-Position entgegen.
Ausgabe: Es erzeugt eine additive Kalibrierungs-Bias für die Draft-Logits. Dies lenkt die Draft-Verteilung so, dass sie mit dem vom Verifizierer bestimmten Bonus-Token konsistent ist, und richtet den Draft effektiv an den Verifizierungsbedingungen aus, obwohl der Draft ohne Kenntnis des Bonus-Tokens generiert wurde.

3. Flex-Decoding

Um den Kompromiss zwischen gegenseitigem Warten (kleine Batch-Größen) und redundanter Berechnung (große Batch-Größen) zu handhaben:

Dynamisches Umschalten: Das System schaltet dynamisch zwischen Parallelem Draft & Verifizieren (für kleine Batch-Größen) und Sequentiellem Draft gefolgt von Verifizieren (für große Batch-Größen) basierend auf einem Schwellenwert für die Batch-Größe um.
Selektive Verifizierung (kleine Batch-Größen): Im parallelen Modus schätzt das System die Wahrscheinlichkeit akzeptierter Längen unter Verwendung der kumulativen Draft-Confidence ab. Es beschneidet Kandidatenverzweigungen mit geringer Confidence und verifiziert nur eine Teilmenge hoch-Confidenter Präfixe, um den Overhead für den Forward-Pass des Zielmodells zu reduzieren, ohne die Länge des Draft-Blocks zu ändern.
Entkoppelte Ausführung (große Batch-Größen): Im sequentiellen Modus generiert der Drafter, sobald die akzeptierte Länge aufgelöst ist, nur noch den einzelnen Zweig, der auf dem verifizierten Präfix basiert. Dies wiederverwendet den KV-Cache aus dem Verifizierungsschritt des Ziels und eliminiert die redundanten Verzweigungen, die den $O(N^2)$ -Overhead verursachen.

Hauptbeiträge

Verlustfreies Block-Diffusions-Drafting: Eine Methode zur Ermöglichung der parallelen Token-Generierung durch Anpassung nur der Attention-Projektoren der letzten Schicht, wodurch die Zielverteilung bewahrt und kostspieliges kontinuierliches Vortraining vermieden wird.
Bonus-gesteuerte Kalibrierung: Ein Mechanismus zur Minderung der Diskrepanz zwischen Draft und Verifizierung durch Kalibrierung der Draft-Logits basierend auf dem aufgelösten Bonus-Token, was die Ausrichtung erheblich verbessert.
Batch-adaptive Ausführung: Ein einheitliches Rahmenwerk, das zwischen parallelen und sequentiellen Strategien wechselt und eine beschneidungsbasierte Pruning-Methode einsetzt, um einen Zusammenbruch der Beschleunigung bei großen Batch-Größen zu verhindern.

Experimentelle Ergebnisse

Experimente wurden an Qwen3-Modellen (1,7B, 4B und 8B) über Benchmarks hinweg durchgeführt, darunter GSM8K, MATH, HumanEval, MBPP und MT-Bench.

Leistung: FlexDraft erzielt eine durchschnittliche 4,59-fache Beschleunigung gegenüber der autoregressiven Dekodierung auf Qwen3-8B ohne Qualitätsverlust.
Vergleich: Es übertrifft starke Baselines, darunter EAGLE-3, DFlash, BiTA und Apple MTP. Insbesondere erzielt FlexDraft auf Qwen3-8B eine 5,88-fache Beschleunigung auf GSM8K und 5,79-fach auf MATH und übertrifft damit DFlash (jeweils 4,55-fach und 5,50-fach).
Robustheit gegenüber Batch-Größen: Im Gegensatz zu bestehenden parallelen Methoden, die aufgrund redundanter Berechnungen bei großen Batch-Größen unter einem Zusammenbruch der Beschleunigung leiden, behält FlexDraft über Batch-Größen (1 bis 16) hinweg eine überlegene Beschleunigung bei, indem es bei Bedarf auf sequentielle Ausführung umschaltet und Verzweigungen mit geringer Confidence beschneidet.
Ablationsstudien:
- Attn Tuning allein liefert signifikante Gewinne durch die Wiederverwendung von Zielwissen.
- Bonus-gesteuerte Kalibrierung verbessert die Beschleunigung weiter durch Ausrichtung der Drafts an den Verifizierungsbedingungen.
- Selektive Verifizierung reduziert redundante Forward-Berechnungen des Ziels und trägt zur finalen Beschleunigung bei.
- Die Verwendung von 10 Schichten für das Draft-Modell erwies sich als bester Kompromiss zwischen Akzeptanzlänge und Dekodierungs-Overhead.

Bedeutung

Die Arbeit behauptet, dass FlexDraft einen praktischen Weg zu einer effizienten LLM-Inferenz unter diversen Serving-Bedingungen bietet. Durch die Kombination einer leichten Parameteranpassung mit dynamischen Ausführungsstrategien adressiert es erfolgreich die inhärenten Kompromisse der parallelen spekulative Dekodierung. Das Rahmenwerk erzielt eine hohe Beschleunigung, ohne die Verteilung des Zielmodells zu beeinträchtigen (verlustfrei), und vermeidet die Qualitätsverschlechterung, die oft mit Methoden verbunden ist, die umfangreiches kontinuierliches Vortraining erfordern. Die Fähigkeit, über variierende Batch-Größen hinweg einen hohen Durchsatz aufrechtzuerhalten, macht es besonders für Produktionsumgebungen geeignet, in denen die Arbeitslastmerkmale schwanken.