Technische Samenvatting: Geloopte State-Space Taalmodellen met Adaptieve Exit-State Selectie

Probleemstelling

Recente vooruitgang in Large Language Models (LLMs) is primair gedreven door het schalen van modelgrootte, trainingsdata en rekenintensieve inferentie. Echter, dit schalingsparadigma introduceert aanzienlijke implementatieuitdagingen, inclusief verhoogd geheugengebruik en snel groeiende inferentiekosten door het onderhoud van grote key-value caches in Transformer-gebaseerde self-attention mechanismen. Hoewel State Space Models (SSMs), in het bijzonder Mamba-2, een efficiënt alternatief bieden met gunstige scaling voor lange contexten, blijft het onduidelijk of het principe van "computationele diepte boven onafhankelijke parameters"—onlangs aangetoond in Looped Transformers—ook van toepassing is op SSMs. Bovendien heeft bestaand onderzoek naar geloopde architecturen zich bijna uitsluitend gericht op Transformer-backbones, wat een gat laat in het begrip van hoe recurrente parameterscharing functioneert binnen Mamba en hybride Mamba-Transformer architecturen.

Methodologie

De auteurs stellen Looped Mamba en Looped Hybrid Mamba–Transformer architecturen voor. Het kernmechanisme bestaat uit het delen van de parameters van een enkele Mamba (of hybride) block over lagen en het herhaaldelijk toepassen hiervan via expliciete eindige-diepte lussen.

Architectuur

• Notatie: Een model wordt aangeduid als $N \otimes R$, waarbij $N$ het aantal onafhankelijk geparametriseerde lagen is en $R$ het aantal recurrente stappen. De effectieve diepte is $NR$.
• Looped Mamba: Past herhaaldelijk een gedeelde $N$-laagse Mamba-backbone $M_\theta$ toe.
  • $h^{(r)} = M_\theta(h^{(r-1)})$ voor $r = 1, \dots, R$.
  • De uiteindelijke output wordt afgeleid van de hidden state op stap $R$.
• Looped Hybrid: Breidt de lus uit naar een backbone die zowel Mamba-2 lagen als een kleine subset van causale self-attention lagen bevat. De lus past toe op de gehele hybride block, waardoor iteratieve verfijning van verborgen representaties mogelijk is door zowel state-space updates als attention-gebaseerde interacties.
• Trainingsprotocol:
  1. Pre-training: Geoptimaliseerd met een final-loop objectief (loss berekend alleen op stap $R$).
  2. Exit-Gate Fase 1: Introduceert een entropie-geregulariseerd adaptief computationeel objectief (geadapteerd van Ouro) om een token-niveau distributie over exit-dieptes te leren met behulp van outputs van alle recurrente stappen.
  3. Exit-Gate Fase 2: Bevriest de backbone en head, en traint enkel de exit gate om stimulering te bevorderen wanneer de volgende stap de voorspellingsloss verbetert.

Adaptieve Exit-State Selectie

De paper adapteert de twee-fasen exit gate van Ouro naar Mamba. Tijdens inferentie controleert een drempelwaarde $q$ de selectie van de exit-stap $z_i$ voor elk token op basis van de cumulatieve exit-waarschijnlijkheid.

• Semantische Notitie: De auteurs verduidelijken dat in SSMs de state chain van alle $R$ stappen behouden moet blijven vanwege continuïteitsrestricties van de state. Daarom selecteert de "exit" welke recurrente stap de hidden state doorgeeft aan de taalmodellering-head; het vermindert inherent niet de wall-clock computation of het aantal uitgevoerde stappen zonder aanvullende state-handling mechanismen.

Belangrijkste Bijdragen

1. Gecontroleerde Empirische Studie: De auteurs presenteren een rigoureuze vergelijking van Looped Mamba en hybride modellen onder gematchte iso-parameter (dezelfde $N$, variërend $R$) en iso-FLOPs (dezelfde effectieve diepte $NR$) protocollen tegenover niet-geloopte Mamba en Transformer baselines.
2. Evaluatie van Redeneringstaken: De studie valideert de aanpak op synthetische redeneertaken (Mano en $p$-hop inductie) en taalmodel pre-training met downstream benchmarks.
3. Adaptief Exit-Mechanisme: De introductie van een drempelwaarde-gecontroleerd adaptief exit-state selectiemechanisme voor geloopde state-space modellen, wat aantoont dat het de downstream performance kan verbeteren bij intermediaire dieptes.

Experimentele Resultaten

Synthetische Redeneringstaken

• Mano Taak (Modulaire Rekenkunde): Looped Mamba presteert consistent beter dan parameter-gematchte niet-geloopte baselines. In de moeilijkste setting ($L=24$) bereikte een $4 \otimes 6$ model een nauwkeurigheid van 99,35% vergeleken met 7,50% voor de $4 \otimes 1$ baseline. Onder iso-FLOPs condities (effectieve diepte 24), presteerde het $4 \otimes 6$ model significant beter dan de $24 \otimes 1$ niet-geloopte baseline (99,35% vs. 37,45%), wat suggereert dat recurrente parameterscharing een sterkere inductieve bias biedt voor compositionele computation dan het vergroten van ongedeelde lagen.
• $p$-hop Inductie (Recursieve Retrieval): Ondiepe niet-geloopte Mamba modellen hadden moeite (bijv. $1 \otimes 1$ bij 38-41%), terwijl het toevoegen van lussen de prestaties drastisch verbeterde. Het $1 \otimes 12$ model bereikte een bijna perfecte nauwkeurigheid (99,89%), wat aangeeft dat recursieve diepte kritischer is dan een verschillend aantal parameters voor recursieve retrieval.

Taalmodellering en Downstream Benchmarks

• Iso-Parameter: Recurrente diepte verbeterde de validatie perplexiteit (PPL) voor zowel Mamba-2 als LLaMA-2 backbones onder een vast parameterbudget.
• Iso-FLOPs: Diepere niet-geloopte modellen ($24 \otimes 1$) behielden een voordeel in validatie perplexiteit ten opzichte van geloopte modellen met een gematchte effectieve diepte. Echter, geloopte modellen bleven competitief op downstream benchmarks en presteerden soms beter dan diepere niet-geloopte baselines.
• Adaptieve Exit:
  • 140M Schaal: Adaptieve exit gating verbeterde de downstream performance. Voor het 140M Mamba-2 model bereikte de Stage 2 exit gate een gemiddelde score van 38,87 met een gemiddeld geselecteerde exit-stap van 3,00, wat beter is dan de fixed-loop4 baseline (37,79).
  • 370M Schaal: De beste downstream performance bleef geconcentreerd nabij de fixed loop4 setting. Exit-gated varianten overtroffen de piek fixed-loop scores niet, wat suggereert dat grotere modellen meer vertrouwen op de laatste recurrente stap of een betere gate-kalibratie vereisen.
  • Trade-off: Exit gating biedt een controleerbare grens tussen geselecteerde diepte en performance, waardoor het model effectief intermediaire recurrente states kan gebruiken in plaats van een vaste maximale diepte af te dwingen.

Betekenis en Claims

Het artikel claimt dat recurrente diepte een veelbelovende, parameterefficiënte schalingsas biedt voor Mamba en hybride taalmodellen. Belangrijkste bevindingen zijn:

• Diepte vs. Parameters: Veel redeneringstaken profiteren van een grotere computationele diepte in plaats van extra onafhankelijke parameters. Het recursief hergebruiken van Mamba-lagen verbetert de computationele capaciteiten voor compositionele en recursieve berekeningen aanzienlijk zonder het aantal parameters te verhogen.
• Hybride Leefbaarheid: Het principe van geloopde diepte is ook effectief toepasbaar op hybride Mamba-Transformer architecturen, waarbij attention-lagen gezamenlijk met Mamba-lagen worden geloopt.
• Adaptieve Selectie: Hoewel huidige SSM-implementaties een volledige state executie vereisen, maakt adaptieve exit-state selectie de selectie van optimale voorspellingsdieptes mogelijk, wat de downstream performance verbetert (met name op de 140M schaal).

De auteurs blijven bescheiden over hun claims en merken op dat de experimenten beperkt zijn tot modellen tot 370M parameters en dat daadwerkelijke besparingen in de inferentie-tijd een nieuw mechanisme vereisen om de recurrent stappen volledig over te slaan. De resultaten worden gepresenteerd als gecontroleerd empirisch bewijs in plaats van een volledige karakterisering van prestaties op grote schaal.