Continuous Depth Batching: Effizientes Batching für depth-adaptive Loop-Modelle

CDB: Erstes effizientes Batching für depth-adaptive Loops

Looped Language Models sind eine Architekturvariante, bei der ein Block geteilter Schichten innerhalb eines einzigen Forward-Passes mehrfach durchlaufen wird. Diese Wiederverwendung ermöglicht depth-adaptive Inference: Für „einfache“ Tokens kann der Loop nach wenigen Durchläufen verlassen werden, für „schwere“ Tokens wird er häufiger ausgeführt. Dadurch skaliert der Rechenaufwand mit der Komplexität der Eingabe, was potenziell die Inferenzkosten senkt. Die Autoren eines neuen Papers bezeichnen dies als ein Hauptversprechen von looped Language Models.

Allerdings besteht ein praktisches Hindernis: Standard-Batching-Systeme wie vLLM setzen voraus, dass alle Tokens in einem Batch denselben Forward-Pass durchlaufen, also dieselbe Anzahl an Schichten. Bei variabler Loop-Tiefe ist diese Annahme verletzt. Tokens mit unterschiedlicher Loop-Anzahl können keinen einheitlichen Forward-Pass teilen. Die Autoren stellen Continuous Depth Batching (CDB) als erste effiziente Batching-Methode für depth-adaptive looped Language Models vor. CDB bildet neue Batches zwischen den Loop-Schritten und adressiert damit das Problem unterschiedlicher Loop-Anzahlen. Die Methode plant geloopte und nicht-geloopte Teile der Architektur dynamisch, verwaltet looped KV-Caching und sagt voraus, welche Tokens den Loop vorzeitig verlassen, um Batches asynchron vorzubereiten.

Laut den Autoren hängt der praktische Wert depth-adaptiver Inference davon ab, ob Batching effizient gemacht werden kann. CDB wird als Methode beschrieben, die diese Lücke schließen soll, indem sie die Vorteile depth-adaptiver Modelle unter Batching-Bedingungen erhält.

Das Problem: Ungleiche Loop-Tiefen brechen Standard-Batching

Das Kernproblem liegt in der Funktionsweise von Standard-Batching. Beim Batching werden mehrere Sequenzen in einem einzigen Forward-Pass verarbeitet, um die Hardware-Auslastung zu maximieren. Dies funktioniert nur, wenn alle Tokens im Batch dieselbe Berechnung durchlaufen – also dieselbe Anzahl an Schichten oder, im Fall von Loop-Modellen, dieselbe Anzahl an Loop-Iterationen. Tokens mit unterschiedlicher Loop-Anzahl können keinen einheitlichen Forward-Pass teilen: Ein Token, das den Loop nach zwei Durchläufen verlässt, und ein anderes, das fünf Durchläufe benötigt, müssten zu unterschiedlichen Zeitpunkten unterschiedliche Schichten ausführen. In einem starren Batch-System wie vLLM, das auf eine feste Modelltiefe ausgelegt ist, lassen sich solche heterogenen Anforderungen nicht abbilden.

Die Konsequenz ist, dass entweder alle Tokens auf die maximale Loop-Tiefe gesetzt werden müssen, um einen einheitlichen Forward-Pass zu ermöglichen – was die Rechenersparnis der depth-adaptiven Inference zunichte macht – oder dass Tokens mit unterschiedlicher Tiefe in separaten Batches verarbeitet werden müssen. Da die Tokens innerhalb eines Batches dieselbe Loop-Anzahl aufweisen müssen, führt dies zu einer Fragmentierung der Batches: Je mehr unterschiedliche Loop-Tiefen im Eingabestrom vorkommen, desto kleiner werden die einzelnen Batches, da nur noch Tokens mit identischer Tiefe gruppiert werden können. Kleinere Batches reduzieren die Effizienz der Matrixmultiplikationen, die von großen Batches profitieren, und erhöhen den Overhead durch häufigere Kernel-Aufrufe. Die Autoren weisen darauf hin, dass der praktische Wert depth-adaptiver Inference davon abhängt, ob Batching effizient gemacht werden kann.

Zusätzlich entstehen Probleme im Speicher und in der Zeit. Beim Speicher muss der KV-Cache für alle Tokens bis zur maximalen Loop-Tiefe vorgehalten werden, selbst wenn viele Tokens bereits nach wenigen Iterationen aussteigen. Das verschwendet GPU-Speicher, der für größere Batches oder längere Sequenzen fehlt. Zeitlich entsteht ein Synchronisationsproblem: Wenn ein Batch gemeinsam verarbeitet wird, muss die GPU auf das Token mit der höchsten Loop-Anzahl warten, während die bereits fertigen Tokens ungenutzt bleiben. Diese Ineffizienzen machen deutlich, warum eine spezialisierte Batching-Methode erforderlich ist.

Wie CDB Batches zwischen Loop-Schritten dynamisch neu bildet

Continuous Depth Batching setzt an dem beschriebenen Problem an. Statt alle Tokens eines Batches über die gesamte Loop-Tiefe zu zwingen, bildet CDB neue Batches zwischen den Loop-Schritten. Nach jedem Durchlauf durch den geloopten Block wird geprüft, welche Tokens den Loop verlassen und welche weiterlaufen. Die verbleibenden Tokens werden zu einem neuen Batch zusammengefasst, der den nächsten Loop-Schritt ausführt. Dieser Ansatz ermöglicht es, dass Tokens mit unterschiedlicher Loop-Tiefe nicht mehr in einem starren Batch gefangen sind, sondern dynamisch neu gruppiert werden.

Um diese dynamische Neuformierung effizient zu gestalten, plant CDB geloopte und nicht-geloopte Teile der Architektur getrennt. Nicht-geloopte Schichten wie Token-Embedding oder LM-Head werden nur einmal pro Token ausgeführt, während der geloopte Kern mehrfach durchlaufen wird. Diese Trennung verhindert, dass feste Kosten pro Token unnötig mit der Loop-Tiefe skaliert werden. Die Planung muss dabei berücksichtigen, welche Schichten Teil des rekurrenten Kerns sind und welche außerhalb liegen, um die Ausführungsreihenfolge zu optimieren.

Ein zentraler Baustein ist das Management des looped KV-Cachings. Die Key-Value-Zustände der rekurrenten Schichten müssen über die Loop-Iterationen hinweg korrekt gespeichert und wiederverwendet werden, ohne dass es zu Konflikten zwischen Tokens mit unterschiedlicher Tiefe kommt. Da Tokens, die früh aussteigen, weniger KV-Cache benötigen als solche, die lange im Loop bleiben, muss CDB den Speicher dynamisch zuweisen und freigeben. Dies ist notwendig, um den GPU-Speicher effizient zu nutzen und nicht für alle Tokens den maximalen Cache vorzuhalten.

Ein weiterer Baustein ist die Vorhersage von Token-Exits. CDB sagt asynchron voraus, welche Tokens den Loop vorzeitig verlassen werden, und bereitet die entsprechenden Batches bereits im Voraus vor. Die Vorhersage muss hinreichend genau sein, um die Batch-Planung nicht zu stören. Die Autoren berichten, dass diese Kombination aus dynamischer Planung, KV-Cache-Verwaltung und Exit-Prädiktion es ermöglicht, die Vorteile depth-adaptiver Inference unter Batching-Bedingungen zu erhalten.

Ergebnisse: 99 Prozent des maximalen Speedups auf Ouro und Huginn

Die Autoren evaluieren CDB auf zwei looped Language Models: Ouro 1.4B und Huginn 3.5B. In den Experimenten erreicht CDB bis zu 99 Prozent des geschätzten maximalen Speedups, der durch depth-adaptive Inference theoretisch möglich ist. Das bedeutet, dass die Methode nahezu die gesamte potenzielle Beschleunigung ausschöpft, die sich aus der variablen Loop-Tiefe ergibt. Die Autoren betonen, dass weitere Gewinne primär von der Modellarchitektur und dem Exit-Verhalten abhängen.

Die Ergebnisse zeigen jedoch auch, dass die Architektur des Modells einen erheblichen Einfluss hat. Vollständig geloopte Architekturen – also Modelle, bei denen nahezu alle Schichten im rekurrenten Kern geteilt werden – sind am besten für depth-adaptive Inference geeignet. Der Grund liegt darin, dass bei solchen Modellen der überwiegende Teil der Berechnung von der Loop-Tiefe abhängt und damit von CDB effizient skaliert werden kann. Wenn der geloopte Kern den größten Teil der Rechenlast ausmacht, kann die variable Tiefe ihre Wirkung voll entfalten.

Große nicht-geloopte Schichten außerhalb des rekurrenten Kerns, etwa Token-Embedding, LM-Head oder ungeteilte Transformer-Blöcke, verlangsamen und verkomplizieren die Planung. Diese Schichten müssen für jedes Token unabhängig von der Loop-Tiefe ausgeführt werden und erzeugen feste Kosten pro Token, die nicht durch depth-adaptive Einsparungen reduziert werden können. Wenn solche Schichten einen signifikanten Anteil der Gesamtberechnung ausmachen, schmälert das den erreichbaren Speedup, weil selbst bei minimaler Loop-Tiefe noch ein hoher Grundaufwand bestehen bleibt. Zudem erschweren sie die dynamische Batch-Planung, da sie zwischen den Loop-Schritten nicht einfach übersprungen werden können. Die verbleibenden Lücken zum theoretischen Maximum hängen primär von der Modellarchitektur und dem Exit-Verhalten ab – also davon, wie gut das Modell vorhersagen kann, wann ein Token den Loop verlässt, und wie homogen die Loop-Tiefen im Batch verteilt sind. Damit liefert CDB nach Angaben der Autoren erstmals den Nachweis, dass depth-adaptive Inference nicht nur ein theoretisches Konzept ist, sondern unter realistischen Batching-Bedingungen nahezu verlustfrei umgesetzt werden kann.

Loop-Modelle im Kontext: Von Mixture-of-Recursions bis LoopWM

Die Arbeit ordnet sich in ein wachsendes Forschungsfeld zu Loop-Modellen und adaptiver Berechnung ein. Die kuratierte Liste „Awesome Loop Models“ definiert Loop-Modelle streng: Innerhalb eines einzelnen Forward-Passes wird eine geteilte gelernte Schicht, ein Block, ein Modul oder ein Operator wiederverwendet. Diese Definition grenzt Loop-Modelle von bloßen Iterationen auf höherer Ebene wie Agent-Loops oder wiederholten Full-Model-Aufrufen ab. Die Liste dient als Referenzkatalog für Arbeiten, die dieser Definition entsprechen.

Adam Fischs „Mixture-of-recursions“ implementiert dynamische rekursive Tiefen für adaptive Token-Level-Berechnung in LLMs. Laut einer Profilseite zu Fisch ist er Senior Research Scientist bei Google DeepMind und hat an Methoden wie Conformal risk control gearbeitet. Die Arbeit zeigt, wie rekursive Tiefen genutzt werden können, um den Rechenaufwand pro Token anzupassen.

Looped World Models (LoopWM) werden als erste geloopte Architekturen für Weltmodellierung beschrieben. Sie verfeinern latente Umgebungszustände iterativ durch einen parametergeteilten Transformer-Block und erreichen nach Angaben der Autoren bis zu 100-fache Parametereffizienz gegenüber konventionellen Ansätzen. Die adaptive Berechnung skaliert die Tiefe automatisch entsprechend der Komplexität jedes Vorhersageschritts.

Eine Survey von Soumajyoti Sarkar unterscheidet adaptive Compute-Techniken von Mixture-of-Experts: Bei MoE bleibt der Rechenaufwand pro Token fix, während echte adaptive Compute-Verfahren die Berechnung je nach Eingabekomplexität variieren. Diese Unterscheidung ist zentral für das Verständnis von Loop-Modellen als adaptive Compute-Architekturen.

LoopMDM zeigt, dass selektives Loopen früher-mittlerer Transformer-Schichten in maskierten Diffusionsmodellen Trainingseffizienz und Modellleistung verbessert. Die Autoren berichten von bis zu 3,3-fach weniger Trainings-FLOPs bei gleicher Leistung und einer Verbesserung von bis zu 8,5 Punkten auf GSM8K. Inference-Compute kann durch variable Loop-Anzahl skaliert werden.

Schließlich findet eine Studie zu looped Language Models für agentische Tool-Nutzung, dass adaptive Inference einen günstigeren Compute-Performance-Trade-off bietet, indem zusätzliche Berechnung nur bei Bedarf zugewiesen wird. In kontrollierten Experimenten auf API-Bank, BFCL und NESTful profitierte die kompositionelle Tool-Nutzung von rekurrenter Berechnung, während adaptive Inference den Trade-off verbesserte. CDB fügt dieser Landschaft die fehlende Infrastrukturkomponente hinzu: ein Batching-Verfahren, das die variablen Tiefen dieser Modelle in der Praxis handhabbar macht.