Technische Samenvatting: Depth-Recurrent Attention Mixtures (Dreamer)

Probleemstelling

Traditionele Chain-of-Thought (CoT) redenering in Large Language Models (LLMs) vertrouwt op het verbaliseren van gedachten in discrete natuurlijke taal. Deze aanpak beperkt de expressiviteit en brengt hoge computationele kosten met zich mee door lange sequentielengtes. Hoewel Depth Recurrence (DR)—het hergebruiken van parameters over de diepte om latent redeneren in een continue ruimte mogelijk te maken—een veelbelovend alternatief biedt, kampen bestaande implementaties met twee kritieke bottlenecks:

1. Layer-Size Bottleneck: Het schalen van dense DR-modellen is onhaalbaar omdat het volledige model bij elke dieptestap moet worden uitgevoerd.
2. Hidden-Size Bottleneck: Het verplaatsen van redenering van de sequentiedimensie naar de dieptedimensie verergert de beperking van constante hidden-state groottes. Terwijl aandachtmechanismen dit probleem hebben opgelost langs de sequentiedimensie (Sequence Attention), negeerde eerdere DR-werken deze bottleneck langs de dieptedimensie.

Bovendien ontbreekt het de voorgaande onderzoeken aan rigoureuze baselines die gematcht zijn op het gebied van Floating Point Operations (FLOPs), parameteraantallen en geheugengebruik, waarbij vaak wordt vertrouwd op gedeeltelijk vaste laagstapels of dense cores die een eerlijke vergelijking hinderen.

Methodologie: Dreamer Framework

De auteurs introduceren Dreamer, een modulair framework van Depth-Recurrent Attention Mixtures. Deze architectuur verenigt drie orthogonale aandachtdimensies in één enkel, herhaaldelijk toegepast laag:

1. Sequence Attention (SA): Standaard self-attention die opereert langs de token-sequentiedimensie.
2. Depth Attention (DA): Een nieuw mechanisme dat aandacht toepast over voorgaande dieptes van dezelfde token.
   • Implementatie: De sequentiedimensie wordt behanden als de batchdimensie, en de diepte als de sequentiedimensie.
   • Optimalisatie: In tegenstelling tot SA kan de Key-Value (KV) cache voor DA na elke token worden verwijderd of overschreven tijdens inferentie, waardoor het geheugengebruik constant blijft ten opzichte van de sequentielengte.
   • Encoding: Gebruikt RoPE (Rotary Positional Embedding) gecodeerd voor diepte, met een omgekeerde helft toegepast om voorwaartse informatie over de maximale diepte te coderen.
3. Expert Attention (EA): Een sparse Mixture of Experts (MoE) die opereert over MLP-experts.
   • Adaptatie voor DR: Om de diepterecurrente (depth recurrence) te verwerken, worden projectiegewichten niet direct gedeeld. In plaats daarvan worden projecties en aggregaties omgezet in lichtgewicht, sparse MoEs met gekoppelde aandachtsscores om training te stabiliseren.
   • Routing: Gebruikt een Top-K routingstrategie met een bias-vector voor balans (vergelijkbaar met DeepSeek-V3) en introduceert een gedeelde, altijd actieve expert geschaald door de maximale score om de gradiëntstroom te bevorderen zonder extra rekenkosten.

Architectuurontwerp:
Het model past deze aandachtmechanismen sequentieel toe (DA $\to$ SA $\to$ EA) of in parallelle varianten. De auteurs gebruiken RMSNorm op de residual stream om de residual scales over de dieptes te beheren. Om een eerlijke vergelijking te garanderen, gebruiken de auteurs een bivariate coördinaat-descent benadering om FLOPs, parameters en geheugen strikt te matchen tussen hun DR-modellen en klassieke gelaagde baselines.

Belangrijkste Bijdragen

1. Verenigd Framework: De introductie van een modulaire architectuur die Sequence, Depth en Expert aandacht combineert, waardoor schalingdimensies (parameters, FLOPs, diepte, latent geheugen) worden ontkoppeld voor efficiënte schaling.
2. Mitigatie van Bottlenecks:
   • Sparse Expert Attention verlicht de layer-size bottleneck, waardoor DR efficiënt kan schalen.
   • Depth Attention verlicht de hidden-size bottleneck, waardoor DR effectief kan schalen voor complexe latente redenering.
3. Rigoureuze Evaluatie: Het artikel biedt nauw aansluitende vergelijkingen (FLOP, parameter en geheugen) tussen DR-varianten en SOTA MoEs, een gat dat in eerdere literatuur werd geïdentificeerd.
4. Analyse van Kennisgebruik: De auteurs analyseren de kennisallocatie over de dieptes, waarbij ze de diversiteit van expertselectie en hergebruikpatronen kwantificeren.

Experimentele Resultaten

De modellen werden geëvalueerd op natuurlijke taal redeneerbenchmarks (GSM8K, MATH, MMLU-Math, MathQA) en taalmodelleringstaken, getraind op ongeveer 100B tokens.

• Data-efficiëntie: Diepte-recurrente modellen (DR en DR+DA) hadden 2 tot 8$\times$ minder trainings-tokens nodig om dezelfde nauwkeurigheid te bereiken als FLOP-, parameter- en geheugengematchte SOTA-baselines.
• Prestaties vs. Schaal: Modellen met diepte 16 (ongeveer 1B parameters) gebruikend DR+DA presteerden beter dan klassiek gelaagde baselines met diepte 32 (ongeveer 2B parameters) over alle redeneerbenchmarks. Dit vertegenwoordigt een reductie van ongeveer 2$\times$ in parameters, FLOPs en geheugen voor superieure prestaties.
• Impact van Depth Attention: Hoewel pure DR (zonder DA) winst liet zien, leverde de combinatie met Depth Attention (DR+DA) de sterkste resultaten op, vooral in diepere modellen (diepte 32), waar het zowel de baseline als pure DR aanzienlijk versloeg.
• Expert Diversiteit: Analyse toonde aan dat DR-modellen een 2 tot 11$\times$ grotere expertselectie-diversiteit vertonen vergeleken met SOTA MoEs. Ongeveer 50% van de experts werd gebruikt over ten minste 7 verschillende dieptes, wat wijst op significante kennishergebruik en diepte-generalisatie.

Betekenis en Claims

Het artikel beweert dat Dreamer een volledig diepte-recurrente single-layer architectuur vertegenwoordigt die consequent nauw aansluitende SOTA Transformers in natuurlijke taal redenering verslaat.

• Efficiëntie: Het framework demonstreert dat latent redeneren via diepterecurrente, wanneer uitgerust met Depth Attention en sparse experts, superieure data- en parameterefficiëntie biedt.
• Expressiviteit: Depth Attention biedt een kwalitatieve uitbreiding van de modelexpressiviteit, waarbij het verder gaat dan uniforme skip-verbindingen om gerichte informatie-routing tussen dieptes mogelijk te maken.
• Modulariteit: Door toegang tot kennis (sequentie, diepte, experts) te formuleren als een verenigde aandachtsmix, stellen de auteurs een perspectief voor dat betere begrip, controle en modulaire uitbreidbaarheid van modelgedrag faciliteert.
• Schaalbaarheid: Het werk adresseert de kritieke behoefte om diepterecurrente te schalen voorbij de beperkingen van dense cores en constante hidden sizes, en suggereert dat diepte behandeld moet worden als een "first-class sequential dimension".

De auteurs nemen een bescheiden standpunt in met betrekking tot toekomstig werk, waarbij zij opmerken dat hoewel dynamische diepte en alternatieve positie-encodings veelbelovend zijn, betrouwbare generalisatie buiten de getrainde diepteregimes een nog onverkend uitdaging blijft. Ze erkennen ook dat latent redeneren de interpretatievraagstukken kan bemoeilijken met betrekking tot de specifieke paden die modellen nemen om tot conclusies te komen.