🤖 AI

Looped State-Space Language Models with Adaptive Exit-State Selection

Diese Arbeit zeigt auf, dass die Anwendung von Loop-Architekturen auf Mamba-basierte Zustandsraummodelle die Fähigkeiten zum logischen Schlussfolgern und die Parametereffizienz durch eine erhöhte Rechentiefe verbessert, während sie eine adaptive Exit-State-Selektion zur Optimierung der Vorhersagetiefe einführt, wenngleich sie anmerkt, dass tatsächliche Einsparungen bei der Inferenz weitere Mechanismen zur Zustandsverwaltung erfordern.

Ursprüngliche Autoren: Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

Veröffentlicht 2026-07-14

📖 6 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

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 hätten einen brillanten, superintelligenten Roboterkoch namens Mamba. Dieser Koch ist berühmt dafür, unglaublich schnell und effizient im Umgang mit dem Speicher zu sein, fähig, ein ganzes Festmahl zu zaubern, ohne eine riesige Küche voller verschiedener Werkzeuge zu benötigen. Aber es gibt einen Haken: Wenn er mit einem wirklich kniffligen Rezept konfrontiert wird – wie einem komplexen mathematischen Rätsel oder einem „Finde die versteckte Zutat“-Spiel – bleibt Mamba manchmal stecken. Er muss tiefer nachdenken, aber er verfügt nicht über genügend „Schichten“ des Denkens, um dies zu tun, ohne eine ganze neue, riesige Küche zu bauen (was ein Vermögen kostet).

Die Forscher der Rikkyo University und der University of Tokyo stellten eine einfache, kühne Frage: Was wäre, wenn wir nicht eine größere Küche bauen würden, sondern Mamba dasselbe Gericht immer und immer wieder kochen ließen und den Geschmack mit jedem Durchgang verfeinerten?

Die „Loop“-Küche: Ein Koch, viele Durchgänge

In der Welt der KI sind die meisten Modelle wie ein Fließband in einer Fabrik. Ein Datenstück (wie ein Wort in einem Satz) bewegt sich die Linie hinunter, durchläuft 24 verschiedene Stationen (Schichten), und dann ist es fertig. Wenn man möchte, dass der Roboter intensiver nachdenkt, fügt man normalerweise einfach mehr Stationen hinzu.

Die Autoren versuchten etwas anderes. Sie bauten ein Looped Mamba. Stellen Sie sich vor, Mamba ist eine einzige, supertalentierte Station. Anstatt die Daten durch eine lange Linie von 24 verschiedenen Stationen zu schicken, schicken sie sie immer wieder zur selben Station zurück, 24 Mal hintereinander.

Die Magie: Der Roboter verwendet bei jedem einzelnen Durchgang dasselbe Gehirn (Parameter). Es ist wie ein Schüler, der denselben Absatz aus einem Lehrbuch viermal liest, um ihn zu verstehen, anstatt vier verschiedene, verwirrende Absätze zu lesen.
Das Ergebnis: Bei schwierigen Logikrätseln namens Mano (modulare Arithmetik) und p-hop induction (das Finden verborgener Muster) bewirkte diese „Wiederlese“-Strategie Wunder. Ein Modell, das nur 4 Schichten hatte, aber diese 6 Mal lochte (4 ⊗ 6), deklassierte ein Standardmodell mit 24 einzigartigen Schichten. Dies deutet darauf hin, dass für das logische Denken die Tiefe des Denkens wichtiger ist als die Anzahl der einzigartigen Werkzeuge.

Das „Exit Gate“: Die beste Version auswählen

Aber warten Sie, was, wenn der Roboter den Absatz ewig weiterliest? Das wäre eine Verschwendung von Zeit. Die Forscher fügten ein cleveres Feature hinzu, das Exit Gate (Ausgangstor) genannt wird.

Betrachten Sie dieses Tor als einen intelligenten Selektor. Während Mamba die Daten verarbeitet, fragt das Tor: „Wird die Antwort klarer?“

Wenn die Antwort nach 2 Durchgängen bereits klar ist, sagt das Tor: „Stopp! Wir verwenden das Ergebnis aus diesem Durchgang.“
Wenn das Problem schwierig ist, sagt das Tor: „Mach weiter!“ und wählt das Ergebnis aus dem 3. oder 4. Durchgang aus.

Dies wird Adaptive Exit-State Selection genannt. Die Studie fand heraus, dass dieses Tor für kleinere Modelle (um 140 Millionen Parameter) ein Game-Changer war. Es ermöglichte dem Modell, die „genau richtige“ Menge an Denkzeit für jedes spezifische Wort zu wählen, und performte oft besser, als den Roboter dazu zu zwingen, immer das Ergebnis der maximalen Anzahl an Durchgängen zu verwenden.

Es gibt jedoch eine Wendung. Die Forscher waren sehr sorgfältig darauf bedacht, darauf hinzuweisen, dass dies bisher noch keinen Strom oder Zeit auf dem Computer spart. Obwohl das Tor ein früheres Ergebnis auswählt, läuft der Roboter im Hintergrund physisch alle Durchgänge ab. Dies liegt daran, dass Mambas Speicherzustand kontinuierlich ist: Das Ergebnis von Durchgang 3 hängt von dem Zustand ab, den Durchgang 2 hinterlassen hat, welcher wiederum von Durchgang 1 abhängt. Man kann den Roboter nicht einfach aus den späteren Durchgängen „herausfallen lassen“, ohne die Kette für alle nachfolgenden Wörter zu unterbrechen. Daher wählt das Tor die Vorhersagetiefe (welche Version der Antwort zu verwenden ist), aber die Rechenzeit (Wall-Clock Computation) bleibt gleich. Um tatsächlich Zeit zu sparen, bräuchte man eine neue Art, den Speicherzustand des Roboters zu handhaben, was die Autoren als Aufgabe für die zukünftige Arbeit beschreiben.

Der Realitätscheck: Was nicht funktionierte (und was ein Vielleicht bleibt)

Das Paper ist erfrischend ehrlich darüber, was es nicht gelöst hat.

Größer ist nicht immer besser (für diesen Trick): Als sie ihr Looped Mamba mit einem Standard-, nicht-geloopten Modell verglichen, das genauso „teuer“ in der Ausführung war (gleiche Anzahl an Berechnungen oder „iso-FLOPs“), gewann das Standardmodell beim Basistest „Wie verwirrt ist das Modell?“ (Perplexity). Das geloopte Modell war großartig bei Aufgaben des logischen Schlussfolgerns, aber das tiefere Standardmodell war immer noch etwas besser darin, einfach das nächste Wort in einem Satz vorherzusagen.
Die Größe spielt eine Rolle: Das „Exit Gate“ funktionierte wunderbar für die 140-Millionen-Parameter-Modelle. Aber als sie es bei den größeren 370-Millionen-Parameter-Modellen versuchten, verbesserte das Tor die Dinge nicht wesentlich. Die größeren Modelle schienen es vorzuziehen, die Durchgänge einfach alle durchzulaufen. Die Autoren legen nahe, dass größere Modelle sich eher auf den finalen, tiefen Denkdurchgang verlassen.
Es ist kein Allheilmittel für alles: Die Forscher testeten dies an synthetischen Rätseln (wie mathematischen Bäumen und Musterfindung). Während diese beweisen, dass das Konzept funktioniert, geben sie zu, dass wir noch nicht wissen, ob dieser „Looping“-Trick auch bei offenem, realweltlichem Denken wie dem Schreiben eines Romans oder dem Lösen eines komplexen Rechtsfalls so gut funktionieren wird.

Das Fazentelemnt

Das Paper legt nahe, dass das wiederholte Gebrauch desselben intelligenten Gehirns ein leistungsstarker Weg ist, um KI beim logischen Denken intelligenter zu machen, ohne ein größeres, teureres Gehirn zu bauen. Es ist eine „parametereffiziente“ Skalierungsachse.

Bewiesen: Bei spezifischen Logikrätseln schlägt das Loopen eines kleinen Mamba-Modells ein Standardmodell gleicher Größe und erreicht sogar die Leistung eines viel größeren Modells.
Vorgeschlagen: Dieser Ansatz könnte ein Schlüssel dazu sein, zukünftige KI effizienter zu machen, insbesondere für Aufgaben, die tiefes, schrittweises Denken erfordern.
Noch nicht gelöst: Wir müssen noch herausfinden, wie man das „Exit Gate“ dazu bringt, tatsächlich Rechenzeit zu sparen (anstatt nur die Ausgabe zu selektieren) und ob dies auch bei den massiven Milliarden-Parameter-Modellen funktioniert, die die heutige größte KI antreiben.

Kurz gesagt: Die Autoren haben nicht einfach nur einen größeren Roboter gebaut; sie haben einem kleineren Roboter beigebracht, tiefer nachzudenken, indem sie ihm erlaubten, sich das Problem ein zweites (oder drittes, oder viertes) Mal anzusehen. Und vorerst sieht dieser zweite Blick sehr vielversprechend aus.

Technisches Resümee: Looped State-Space Language Models mit adaptiver Exit-State-Selektion

Problemstellung

Jüngste Fortschritte bei Large Language Models (LLMs) wurden primär durch die Skalierung der Modellgröße, der Trainingsdaten und des rechenintensiven Inferenz-Verfahrens vorangetrieben. Dieses Skalierungsparadigma führt jedoch zu erheblichen Herausforderungen beim Deployment, einschließlich des steigenden Speicherverbrauchs und der rasant wachsenden Inferenzkosten aufgrund der Aufrechterhaltung großer Key-Value-Caches in Transformer-basierten Self-Attention-Mechanismen. Während State Space Models (SSMs), insbesondere Mamba-2, eine effiziente Alternative mit günstiger Skalierung für lange Kontexte bieten, bleibt unklar, ob das Prinzip von „rechnerischer Tiefe statt unabhängiger Parameter“ – wie es kürzlich bei Looped Transformers demonstriert wurde – auch für SSMs gilt. Zudem hat sich die bisherige Forschung zu Looped-Architekturen fast ausschließlich auf Transformer-Backbones konzentriert, was eine Wissenslücke hinsichtlich der Funktionsweise von rekurrentem Parameter-Sharing innerhalb von Mamba- und hybriden Mamba-Transformer-Architekturen hinterlässt.

Methodik

Die Autoren schlagen Looped Mamba und Looped Hybrid Mamba–Transformer Architekturen vor. Der Kernmechanismus besteht darin, die Parameter eines einzelnen Mamba- (oder Hybrid-) Blocks über Schichten hinweg zu teilen und diesen durch explizite endliche Tiefen-Schleifen (Loops) wiederholt anzuwenden.

Architektur

Notation: Ein Modell wird als $N \otimes R$ bezeichnet, wobei $N$ die Anzahl der unabhängig parametrisierten Schichten und $R$ die Anzahl der rekurrenten Schritte ist. Die effektive Tiefe beträgt $NR$.
Looped Mamba: Wendet wiederholt ein geteiltes $N$ $N$ -schichtiges Mamba-Backbone $M_\theta$ $M_{θ}$ an.
- $h^{(r)} = M_\theta(h^{(r-1)})$ für $r = 1, \dots, R$ .
- Der finale Output wird aus dem Hidden State des Schritts $R$ abgeleitet.
Looped Hybrid: Erweitert den Loop auf ein Backbone, das sowohl Mamba-2-Schichten als auch eine kleine Teilmenge von Causal Self-Attention-Schichten enthält. Der Loop wendet sich auf den gesamten hybriden Block an, was eine iterative Verfeinerung der Hidden Representations durch sowohl State-Space-Updates als auch Attention-basierte Interaktionen ermöglicht.
Trainingsprotokoll:
1. Pre-training: Optimiert mittels eines Final-Loop-Objective (Loss wird nur bei Schritt $R$ berechnet).
2. Exit-Gate Stufe 1: Einführung eines entropie-regulierten adaptiven Computation-Objective (adaptiert von Ouro), um eine Token-ebene Verteilung über die Exit-Tiefen unter Verwendung der Outputs aller rekurrenten Schritte zu lernen.
3. Exit-Gate Stufe 2: Das Backbone und der Head werden eingefroren, und es wird nur das Exit-Gate trainiert, um eine Fortsetzung nur dann zu fördern, wenn der nächste Schritt den Vorhersage-Loss verbessert.

Adaptive Exit-State-Selektion

Das Paper adaptiert das zweistufige Exit-Gate von Ouro auf Mamba. Während der Inferenz steuert ein Schwellenwert $q$ die Auswahl des Exit-Schritts $z_i$ für jeden Token basierend auf der kumulativen Exit-Wahrscheinlichkeit.

Semantischer Hinweis: Die Autoren stellen klar, dass in SSMs die State-Chain aller $R$ Schritte aufgrund von Kontinuitätsbeschränkungen des Zustands erhalten bleiben muss. Daher wählt der „Exit“ aus, welcher rekurrenten Schritts Hidden State an den Language Modeling Head übergeben wird; dies reduziert nicht inhärent die Rechenzeit (Wall-Clock Computation) oder die Anzahl der ausgeführten Schritte ohne zusätzliche State-Handling-Mechanismen.

Zentrale Beiträge

Kontrollierte empirische Studie: Die Autoren präsentieren einen rigorosen Vergleich von Looped Mamba und Hybrid-Modellen unter Verwendung von Iso-Parameter (gleiches $N$ , variierendes $R$ ) und Iso-FLOPs (gleiche effektive Tiefe $NR$) Protokollen gegenüber Nicht-Looped Mamba- und Transformer-Baselines.
Evaluierung von Reasoning-Aufgaben: Die Studie validiert den Ansatz auf synthetischen Reasoning-Aufgaben (Mano und $p$ -hop Induction) sowie auf Language Model Pre-training mit Downstream-Benchmarks.
Adaptiver Exit-Mechanismus: Die Einführung eines Schwellenwert-gesteuerten adaptiven Exit-State-Selektionsmechanismus für geloopte State-Space-Modelle, der die Fähigkeit zeigt, die Performance bei intermediären Tiefen zu verbessern.

Experimentelle Ergebnisse

Synthetische Reasoning-Aufgaben

Mano Task (Modulare Arithmetik): Looped Mamba übertrifft konsistent die parameter-gepaarten Nicht-Looped Baselines. In der schwierigsten Einstellung ( $L=24$ ) erreichte ein $4 \otimes 6$ Modell eine Genauigkeit von 99,35 % im Vergleich zu 7,50 % der $4 \otimes 1$ Baseline. Unter Iso-FLOPs-Bedingungen (effektive Tiefe 24) übertraf das $4 \otimes 6$ Modell die $24 \otimes 1$ Nicht-Looped Baseline (99,35 % vs. 37,45 %) signifikant, was darauf hindeutet, dass rekurrentes Parameter-Sharing einen stärkeren induktiven Bias für kompositorische Berechnungen liefert als die Erhöhung ungeteilter Schichten.
$p$ -hop Induction (Rekursive Retrieval): Flache Nicht-Looped Mamba-Modelle hatten Schwierigkeiten (z. B. $1 \otimes 1$ bei ~~38-41 %), während das Hinzufügen von Loops die Performance dramatisch verbesserte. Das $1 \otimes 12$ Modell erreichte eine nahezu perfekte Genauigkeit (~~99,89 %), was zeigt, dass rekursive Tiefe entscheidender ist als die Anzahl der distinkten Parameter für rekursives Retrieval.

Sprachmodellierung und Downstream-Benchmarks

Iso-Parameter: Rekurrente Tiefe verbesserte die Validierungs-Perplexität (PPL) für sowohl Mamba-2 als auch LLaMA-2 Backbones unter festem Parameterbudget.
Iso-FLOPs: Tiefere Nicht-Looped Modelle ( $24 \otimes 1$ ) behielten einen Vorteil in der Validierungs-Perplexität gegenüber geloopten Modellen mit gleicher effektiver Tiefe. Dennoch blieben geloopte Modelle in Downstream-Benchmarks wettbewerbsfähig und übertrafen teilweise tiefere Nicht-Looped Baselines.
Adaptive Exit:
- 140M Skala: Adaptives Exit-Gating verbesserte die Downstream-Performance. Für das 140M Mamba-2 Modell erreichte das Stage 2 Exit-Gate einen Durchschnittswert von 38,87 mit einem durchschnittlich gewählten Exit-Schritt von 3,00 und übertraf damit die Fixed-Loop4-Baseline (37,79).
- 370M Skala: Die beste Downstream-Performance blieb konzentriert nahe der Fixed-Loop4-Einstellung. Die Exit-Gated-Varianten übertrafen die Peak-Scores der Fixed-Loop-Modelle nicht, was darauf hindeutet, dass größere Modelle stärker auf den finalen rekurrenten Schritt angewiesen sind oder eine bessere Gate-Kalibrierung benötigen.
- Trade-off: Exit-Gating bietet eine steuerbare Grenze zwischen gewählter Tiefe und Performance, wodurch das Modell in der Lage ist, intermediäre rekurrenten States effektiv zu nutzen, anstatt eine feste maximale Tiefe erzwingen zu müssen.

Bedeutung und Behauptungen

Das Paper behauptet, dass rekurrente Tiefe eine vielversprechende, parametereffiziente Skalierungsachse für Mamba und hybride Sprachmodelle darstellt. Die Kernergebnisse sind:

Tiefe vs. Parameter: Viele Reasoning-Aufgaben profitieren eher von größerer Rechentiefe als von zusätzlichen unabhängigen Parametern. Die rekursive Wiederverwendung von Mamba-Layern verbessert die kompositorischen und rekursiven Berechnungsfähigkeiten erheblich, ohne die Parameteranzahl zu erhöhen.
Hybride Viabilität: Das Prinzip der geloopten Tiefe lässt sich effektiv auf hybride Mamba-Transformer-Architekturen anwenden, bei denen Attention-Layer gemeinsam mit Mamba-Layern geloopt werden.
Adaptive Selektion: Während aktuelle SSM-Implementierungen die Ausführung des vollen States erfordern, ermöglicht die adaptive Exit-State-Selektion die Wahl optimaler Vorhersagetiefen, was die Downstream-Performance (insbesondere auf der 140M Skala) verbessert.

Die Autoren formulieren ihre Ansprüche bescheiden und weisen darauf hin, dass die Experimente auf Modelle bis zu 370M Parametern beschränkt sind und dass tatsächliche Einsparungen bei der Inferenz-Rechenzeit zukünftige Arbeit an State-Handling-Mechanismen erfordern, um rekurrenten Schritte komplett zu überspringen. Die Ergebnisse werden als kontrollierte empirische Evidenz und nicht als vollständige Charakterisierung der großskaligen Performance präsentiert.