🤖 machine learning

Bastion: Budget-Aware Speculative Decoding with Tree-structured Block Diffusion Drafting

BASTION ist ein trainingsfreies, budgetbewusstes spekulatives Dekodierungsframework, das durch adaptive Best-First-Erweiterung abfrageabhängige Baumstrukturen dynamisch konstruiert, um die Entwurfsqualität mit Hardwarebeschränkungen in Einklang zu bringen, und erreicht eine bis zu 6,61-fache Beschleunigung gegenüber der standardmäßigen autoregressiven Dekodierung, während es bestehende Block-Diffusions-Baselines übertrifft.

Ursprüngliche Autoren: Soowon Oh, Nam Cao, Yujin Kim, Hojung Jung, Huzama Ahmad, Sangmin Bae, Se-Young Yun

Veröffentlicht 2026-05-29

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Soowon Oh, Nam Cao, Yujin Kim, Hojung Jung, Huzama Ahmad, Sangmin Bae, Se-Young Yun

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 Geschichte mit einem sehr weisen, aber unglaublich langsamen Lektor (dem Zielmodell) zu schreiben. Jedes Mal, wenn Sie ein einziges Wort schreiben, müssen Sie warten, bis der Lektor es liest, darüber nachdenkt und Ihnen das grüne Licht gibt, bevor Sie das nächste schreiben können. So funktionieren aktuelle KI-Chatbots: Sie schreiben ein Wort nach dem anderen und warten nach jedem einzelnen Schritt auf eine „Prüfung". Das ist zwar präzise, aber schmerzlich langsam.

Um dies zu beschleunigen, verwenden Forscher ein Entwurfsmodell (einen schnelleren, weniger weisen Assistenten), um die nächsten paar Wörter zu erraten, bevor der Lektor sie prüft. Wenn der Lektor dem Vorschlag zustimmt, können Sie mehrere Wörter auf einmal schreiben und die Wartezeit überspringen.

Allerdings gibt es einen Haken bei den neuesten, schnellsten Entwurfsmodellen (genannt Block Diffusion). Anstatt einen Satz Wort für Wort zu erraten, schreien sie einen ganzen Block von Wörtern auf einmal heraus. Das Problem ist, dass sie, weil sie alle gleichzeitig herausgeschrien werden, nicht immer sicher sind, wie die Wörter in einer logischen Reihenfolge zusammenpassen. Es ist wie ein Koch, der eine Menge Zutaten auf einmal auf die Theke wirft; einzeln sehen sie gut aus, aber wenn Sie einfach das oberste greifen, landen Sie möglicherweise bei einem seltsamen, sinnlosen Gericht.

BASTION ist ein neues System, das dieses Durcheinander beheben und den gesamten Prozess blitzschnell machen soll. So funktioniert es, unter Verwendung einfacher Analogien:

1. Der „Baum der Möglichkeiten" (anstatt eines einzelnen Pfades)

Ältere Methoden würden den Ausruf des Entwurfsmodells nehmen, das einzelne „beste" Wort auswählen, dann das nächste „beste" Wort und auf das Beste hoffen. Wenn sich dieser Pfad als falsch herausstellte, musste der Lektor das Ganze ablehnen, und Sie verloren Zeit.

BASTION ist anders. Stellen Sie sich vor, das Entwurfsmodell gibt Ihnen nicht nur einen Pfad; es gibt Ihnen einen Stammbaum der Möglichkeiten.

Im ersten Schritt sagt es: „Vielleicht ist das nächste Wort 'Katze' (80 % Wahrscheinlichkeit) oder 'Hund' (20 % Wahrscheinlichkeit)."
Anstatt einfach 'Katze' zu wählen, baut BASTION einen kleinen Baum: einen Ast für 'Katze', einen für 'Hund'.
Dann verzweigt es sich für das nächste Wort wieder von beiden, 'Katze' und 'Hund'.
Plötzlich haben Sie einen kleinen Wald potenzieller Sätze, der aus demselben Ausgangspunkt wächst.

2. Der „Kluge Gärtner" (Budgetbewusste Steuerung)

Hier kommt der schwierige Teil: Sie können keinen unendlichen Wald wachsen lassen. Der Lektor (Zielmodell) hat ein Limit, wie viele Äste er gleichzeitig prüfen kann. Wenn Sie einen Baum wachsen lassen, der zu breit oder zu tief ist, dauert die Prüfung des Baumes länger, als die Wörter einfach nacheinander zu schreiben.

Hier kommt BASTIONs „Kluger Gärtner" ins Spiel.

Das Budget: Der Gärtner weiß genau, wie viel Zeit der Lektor hat (das „Budget").
Die Strategie: Anstatt den Baum auf eine feste Größe wachsen zu lassen (wie „immer 10 Äste"), betrachtet der Gärtner die Zuverlässigkeit jedes Astes.
- Wenn ein Ast vielversprechend aussieht (hohe Zuverlässigkeit), lässt der Gärtner ihn tiefer wachsen.
- Wenn ein Ast schwach aussieht, hört er dort auf zu wachsen.
Das Stoppschild: Der Gärtner fragt ständig: „Wird das Hinzufügen eines weiteren Astes uns mehr Geschwindigkeit bringen, oder wird es nur Zeit verschwenden, Sackgassen zu prüfen?" Sobald die Kosten für die Prüfung eines neuen Astes den Nutzen überwiegen, hört der Gärtner auf zu wachsen und schickt den Baum zum Lektor.

3. Der „Tachometer" (Hardwarebewusstsein)

Verschiedene Computer (GPUs) sind wie verschiedene Autos. Ein Sportwagen (eine leistungsstarke GPU) kann einen riesigen Baum sehr schnell prüfen. Ein Kompaktwagen (eine schwächere GPU) könnte mit demselben Baum zu kämpfen haben.

BASTION verfügt über einen eingebauten Tachometer, der genau weiß, wie schnell Ihr spezifischer Computer ist. Es errät nicht einfach; es misst, wie lange es dauert, einen Baum bestimmter Größe auf Ihrem Gerät zu verifizieren. Es verwendet diese Echtzeitdaten, um genau zu entscheiden, wie groß der Baum für Ihr spezifisches Setup sein sollte, und stellt sicher, dass Sie die maximale Geschwindigkeit erreichen, ohne Ihren Computer zu überlasten.

Das Ergebnis

Durch die Kombination dieser Ideen erreicht BASTION das, was die Studie als 6,61-fache Beschleunigung bezeichnet.

Standard-KI: Schreibt 1 Wort, wartet, schreibt 1 Wort, wartet. (Geschwindigkeit: 1x)
Alte schnelle Methoden: Errät ein paar Wörter, bleibt aber oft auf dem falschen Pfad stecken. (Geschwindigkeit: ~2-3x)
BASTION: Zieht einen intelligenten, maßgeschneiderten Baum von Vermutungen, prüft die vielversprechendsten Pfade und stoppt genau dann, wenn es am effizientesten ist. (Geschwindigkeit: ~6,6x)

Kurz gesagt, ist BASTION wie ein kluger Projektmanager für das KI-Schreiben. Anstatt blind zu raten oder eine starre Struktur aufzubauen, baut es dynamisch einen flexiblen „Baum der Optionen", der perfekt an die Geschwindigkeit des Computers angepasst ist und sicherstellt, dass die KI so schnell wie möglich schreibt, ohne Fehler zu machen.

Technische Zusammenfassung: BASTION

Problemstellung

Large Language Models (LLMs) leiden aufgrund der sequenziellen Natur des autoregressiven (AR) Decodierens unter hohen Serving-Kosten, da für jedes generierte Token ein separater Vorwärtsdurchlauf des Zielmodells erforderlich ist. Spekulative Decodierung adressiert dies, indem ein leichtgewichtiger Drafter mehrere Kandidat-Token vorschlägt, die das Zielmodell parallel verifiziert.

Während jüngste Fortschritte bei block-diffusionsbasierten Draftern (z. B. DFlash, TiDAR) die parallele Vorhersage mehrerer zukünftiger Token in einem einzigen Schritt ermöglichen, führen sie zu einer kritischen Einschränkung: Diese Modelle generieren positionsspezifische Randverteilungen anstelle von vollständig bedingten Sequenzen. Folglich führt das Festlegen auf einen einzigen gierigen (Top-1)-Pfad oft zu inkohärenten Sequenzen unter der Verteilung des Zielmodells. Bestehende baumbasierte Methoden zur spekulativen Decodierung (z. B. EAGLE-3) verlassen sich typischerweise auf statische, heuristisch vordefinierte Baumtopologien, die nicht auf die spezifische Abfrage oder Hardwarebeschränkungen angepasst sind, was zu suboptimalen Kompromissen zwischen Draft-Qualität und Verifizierungsaufwand führt.

Methodik: BASTION

BASTION (Budget-Aware Speculative Decoding with Tree-structured Block Diffusion Drafting) ist ein trainingsfreies Framework, das abfrageabhängige Baumtopologien dynamisch konstruiert, um die Decodierbeschleunigung zu maximieren. Es integriert drei synergistische Komponenten:

1. Baumgestützte block-diffusionsbasierte Drafting

Im Gegensatz zu autoregressiven Draftern, die Token sequenziell generieren, nutzt BASTION einen block-diffusionsbasierten Drafter, um in einem einzigen Vorwärtsdurchlauf positionsspezifische Randverteilungen $\{q_k(\cdot)\}$ zu erzeugen.

Topologie-Entkopplung: Die Drafting-Kosten ( $T_{draft}$ ) sind unabhängig von der Baumstruktur (Ketten, breite Bäume oder tiefe Bäume), da der Block parallel verarbeitet wird.
Randverteilungs-Unabhängigkeit: Kandidaten an verschiedenen Positionen sind unter der validierten Kontextbedingung voneinander unabhängig, was ihre freie Kombination zu einem Präfixbaum $\mathcal{T}$ ermöglicht.
Pfad-Bewertung: Jeder Knoten $i$ im Baum in der Tiefe $d(i)$ erhält einen Pfad-Score $\rho(i) = \prod_{k=1}^{d(i)} q_k(x_{ik})$ , der die gemeinsame Wahrscheinlichkeit des Pfades unter der Verteilung des Drafters darstellt.

2. Akzeptanz-Surrogat und Best-First-Erweiterung

Um festzustellen, welche Knoten in den Verifizierungsbaum aufgenommen werden sollen, ohne die Kosten einer Zielmodell-Verifizierung zu verursachen, setzt BASTION ein Pfad-Vertrauens-Surrogat ein.

Surrogat-Schätzung: Die erwartete Akzeptanzlänge $\mathbb{E}[A(\mathcal{T})]$ wird durch die Summe der Pfad-Scores approximiert: $\hat{A}(\mathcal{T}) = \sum_{i \in \mathcal{T}} \rho(i)$ . Dies basiert auf der Prämisse, dass Pfad-Scores mit hoher Wahrscheinlichkeit des Drafters mit der bevorzugten Trajektorie des Zielmodells übereinstimmen.
Optimale Konstruktion: Die Arbeit beweist, dass eine Best-First-Erweiterungsstrategie unter einem festen Budget nachweislich optimal ist, um $\hat{A}(\mathcal{T})$ zu maximieren. Der Algorithmus fügt iterativ den gültigen Knoten mit dem höchsten $\rho(i)$ hinzu und stellt sicher, dass der Baum zunächst entlang der vielversprechendsten Trajektorien wächst. Dies ergibt eine konkave Beziehung zwischen Baumgröße und erwarteter Akzeptanz mit abnehmenden Grenzerträgen.

3. Hardwarebewusster Online-Budget-Controller

BASTION bestimmt dynamisch die optimale Baumgröße $N$ in jedem Decodierschritt, indem es den marginalen Gewinn an Akzeptanzlänge gegen die marginale Zunahme der Verifizierungs-Latenz abwägt.

Latenz-Modellierung: Das Framework verwendet ein hardwarebewusstes Roofline-Modell, um die Verifizierungs-Latenz $T_{verify}(N)$ zu schätzen. Dieses Modell berücksichtigt Speicherbandbreite und Rechenbeschränkungen und wird durch beobachtete Laufzeiten kalibriert, um Kernel-Fusion und Planungs-Overhead zu korrigieren.
Adaptives Stoppen: Der Controller erweitert den Baum, bis die geschätzte Beschleunigung $S(N) = \frac{\hat{A}(\mathcal{T}_N) \cdot L_{AR}}{C(N)}$ zu sinken beginnt. Dies stellt sicher, dass das System die Erweiterung genau dann stoppt, wenn die Kosten für die Verifizierung zusätzlicher Knoten den Nutzen einer längeren akzeptierten Sequenz überwiegen.

Hauptbeiträge

Baumgestützte block-diffusionsbasierte Drafting: Einführung eines Best-First-Präfixbaum-Konstruktionsalgorithmus, der Pfad-Scores aus Drafter-Logits nutzt, um hochwertige, präfixgeschlossene Kandidatenbäume aus block-diffusionsbasierten Logits zu erstellen und damit über die gierige Decodierung einzelner Pfade hinauszugehen.
Hardwarebewusster Online-Budget-Controller: Ein neuartiger Controller, der die optimale Baumgröße in jedem Schritt dynamisch bestimmt, indem er erwartete Akzeptanzgewinne gegen hardwarekalibrierte Verifizierungs-Latenz abwägt, wodurch statische, festgelegte Baumgrößen überflüssig werden.
Umfassende empirische Validierung: Ausführliche Evaluierung über diverse Zielmodelle (Qwen3, Llama-3.1), Benchmarks (Mathematik, Code, Chat) und GPU-Architekturen (A100, H100, A6000, RTX PRO 6000), die signifikante Beschleunigungen ohne Anpassung pro Einstellung nachweisen.

Experimentelle Ergebnisse

Evaluiert auf einer NVIDIA A100 GPU mit gieriger Decodierung (Temperatur 0):

Beschleunigung: BASTION erzielt im Durchschnitt eine 6,61-fache Beschleunigung gegenüber der Standard-autoregressiven Decodierung bei Qwen3-8B.
Vergleich mit Baselines:
- Übertrifft EAGLE-3 (state-of-the-art AR-basierte Baum-Decodierung) um den Faktor 2,45 (6,61-fach vs. 2,70-fach im Durchschnitt).
- Übertrifft DFlash (state-of-the-art block-diffusionsbasierte Einzelfadendecodierung) um den Faktor 1,39 (6,61-fach vs. 4,76-fach im Durchschnitt).
Akzeptanzlänge: BASTION erreicht deutlich größere durchschnittliche akzeptierte Längen (AAL) im Vergleich zu Baselines mit einzelnen Pfaden (z. B. 8,48 vs. 6,24 bei Qwen3-8B), was bestätigt, dass die baumgestützte Verifizierung erfolgreich längere gültige Fortsetzungen erfasst.
Generalisierung: Die Leistungsgewinne sind über verschiedene GPU-Generationen (einschließlich H100 und Blackwell-Architekturen) und Modellfamilien hinweg konsistent, was darauf hindeutet, dass die Methode nicht an spezifische Hardware- oder Modellkonfigurationen gebunden ist.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass BASTION einen bedeutenden Fortschritt in der spekulativen Decodierung darstellt, indem es die Spannung zwischen der parallelen Effizienz block-diffusionsbasierter Drafters und der sequenziellen Kohärenz auflöst, die von Zielmodellen erforderlich ist. Durch den Ersatz statischer, heuristischer Baumformen durch eine dynamische, budgetbewusste Konstruktion nutzt BASTION verfügbare parallele Rechenressourcen vollständig aus, vermeidet gleichzeitig die Latenzstrafen einer übermäßigen Verifizierung unwahrscheinlicher Pfade.

Entscheidend ist, dass die Methode trainingsfrei ist und keine Anpassung pro Einstellung erfordert, was eine sofortige Bereitstellung über diverse Hardware- und Modellkonfigurationen hinweg ermöglicht. Die Autoren behaupten, dass ihr Ansatz die Lücke zwischen dem theoretischen Potenzial block-diffusionsbasierter Drafting und praktischer, hochdurchsatzfähiger Inferenz effektiv schließt und nahezu optimale Beschleunigungen erreicht, indem es die Draft-Struktur adaptiv an Echtzeit-Systembeschränkungen anpasst.