Technische Samenvatting: MESANET: SEQUENTIËLE MODELLERING DOOR LOCAAL OPTIMAAL TEST-TIME TRAINEN

Probleemstelling

Sequentiële modellering wordt momenteel gedomineerd door causale transformer-architecturen die gebruikmaken van softmax self-attention. Hoewel effectief, lijden deze modellen aan een lineaire schaling in geheugen- en computationele vereisten ten opzichte van de sequentielengte tijdens inferentie. Recente inspanningen hebben geprobeerd dit aan te pakken door de softmax-operatie te lineariseren, wat resulteert in recurrente neurale netwerk (RNN) modellen met constante geheugen- en computekosten (bijv. DeltaNet, Mamba, xLSTM). Deze modellen worden vaak verenigd door het perspectief dat hun recurrente dynamiek voortkomt uit een in-context regressiedoelstelling, die benaderd wordt via online leerregels (bijv. gradiëntafdalingstappen). Echter, deze bestaande benaderingen vertrouwen doorgaans op incrementele updates (eerste-orde methoden) in plaats van het oplossen voor de optimale gewichten op elk tijdstip, wat potentieel de bekwaamheid kan beperken om snel aan te passen aan nieuwe bewijslast terwijl eerdere informatie behouden blijft. Bovendien werden eerdere iteraties van optimale test-time traininglagen (zoals de oorspronkelijke Mesa-laag) beperkt door numerieke instabiliteit bij het gebruik van forget gates en een gebrek aan parallelliseerbaarheid, wat schaalbare training verhinderde.

Methodologie

De auteurs introduceren MesaNet, een taalmodelarchitectuur gecentreerd rond een nieuwe, numeriek stabiele en chunkwise parallelliseerbare Mesa-laag. Deze laag is ontworpen om op elk tijdstip lokale optimale test-time training uit te voeren.

De Mesa-laag

In tegenstelling tot standaard lineaire RNN's die hun staat bijwerken via een enkele gradiëntafdalingstap (bijv. DeltaNet), definieert de Mesa-laag de lineaire afbeelding $\Phi_t$ als de oplossing voor een test-time optimalisatieprobleem. Specifiek minimaliseert het een cumulatieve, geregulariseerde kwadratische foutverlies over alle tot nu toe geziene data ($e_{1:t}$):
$$\hat{\Phi}^{mesa}_t = \arg \min_{\Phi} \left( \frac{1}{2} \sum_{t'=1}^t \zeta_{tt'} \|v_{t'} - \Phi k_{t'}\|^2 + \frac{1}{2} \text{Tr}(\Phi \Lambda \Phi^\top) \right)$$
waarbij $\zeta_{tt'}$ cumulatieve forget factors vertegenwoordigt, en $\Lambda$ een symmetrische positief definitieve regularisator is.

De output wordt berekend als $\Delta^{mesa}_t = \hat{\Phi}^{mesa}_t q_t$. Cruciaal is dat, in plaats van recursieve kleinste kwadraten (wat lijdt onder numerieke instabiliteit met forget gates en niet parallelliseerbaar is), de auteurs de optimale query-respons $q^*_t = (H_t + \Lambda)^{-1} q_t$ oplossen met behulp van de Conjugate Gradient (CG) methode.

• Parallellisering: De auteurs maken gebruik van de equivalentie tussen de CG-methode over meerdere tijdstappen en Gated Linear Attention (GLA). Door staatvariabelen $G_t$ en $H_t$ te onderhouden en het lineaire systeem parallel over chunks van de sequentie op te lossen, bereikt de laag een $O(T)$ complexiteit tijdens training en $O(1)$ tijdens inferentie.
• Dynamische Compute: Het aantal CG-iteraties ($k$) dat vereist is om een stopcriterium te bereiken, varieert per sequentie en token. Dit stelt het model in staat om dynamisch test-time compute toe te wijzen, waarbij de inferentiekosten worden afgewogen tegen de prestaties.

Architectuur (MesaNet)

MesaNet neemt een standaard decoder-only transformer-backbone aan (vergelijkbaar met Llama 2), maar vervangt de sequentie-mixende component (meestal Multi-Head Attention) door de Mesa-laag. De architectuur bevat:

• Residuele Blokken: Het stapelen van $N$ blokken die een sequentie-mixende laag (Mesa) en een kanaal-mixende laag (SwiGLU MLP) bevatten.
• Preprocessing: Keys en queries worden genormaliseerd en geconvolueerd met een kleine window voordat ze de Mesa-laag bereiken.
• Hybride Variant: Een "Hawk-Mesa"-model wordt ook verkend, waarbij wordt afgewisseld tussen LRU (Hawk) lagen en Mesa-lagen om de sterktes van korte-sequentie en lange-sequentie modellering te combineren.

Belangrijkste Bijdragen

1. Een Parallelliseerbare, Stabiele Mesa-laag: Het artikel presenteert een versie van de Mesa-laag die de sequentiële bottleneck en numerieke instabiliteit van de oorspronkelijke formulering overwint. Het maakt gebruik van een CG-solver met chunkwise parallellisering, wat efficiënte training op moderne hardware (GPU's/TPU's) mogelijk maakt, terwijl de theoretische voordelen van optimale test-time regressie behouden blijven.
2. Dynamische Toewijzing van Test-Time Compute: De laag roept expliciet een externe solver aan waarbij het aantal iteraties afhankelijk is van de moeilijkheid van het lineaire systeem voor een gegeven sequentie. Dit biedt een principieel mechanisme om de inferentie-compute te variëren op basis van de complexiteit van de input.
3. Uitgebreide Empirische Evaluatie: De auteurs trainen MesaNets op schalen van 140M, 440M en 1B parameters op de SlimPajama dataset. Ze bieden een rigoureuze, "apples-to-apples" vergelijking met Mamba2, xLSTM, DeltaNet, Gated DeltaNet en Transformers, waarbij architectuur, tokenizer en hyperparameters worden gecontroleerd.

Resultaten

• Taalmodellering Prestaties: MesaNets behalen een lagere validatie-perplexiteit (PPL) dan andere RNN-baselines (Mamba2, xLSTM, DeltaNet) over alle schalen. Op de 1B parameter schaal evenaart of overtreft MesaNet zelfs licht de prestaties van een Transformer-baseline met dezelfde architectuur.
• Token Positie Analyse: Een gedetailleerde analyse van de Negative Log-Likelihood (NLL) ten opzichte van de tokenpositie onthult een duidelijk patroon:
  • Vroege Tokens: De meeste RNN's, inclusief MesaNet, presteren beter dan Transformers op de vroege tokens in een sequentie.
  • Lange Contexten: Terwijl Transformers uitblinken bij latere tokens, breiden MesaNet en Hawk-Mesa hun voordeel aanzienlijk uit (voorbij 512 tokens) vergeleken met andere RNN's, hoewel ze uiteindelijk achter Transformers blijven bij zeer lange dieptes.
• Downstream Benchmarks:
  • Globale Redenering: MesaNet presteert beter dan alle andere RNN's op benchmarks die langetermijn-afhankelijkheden vereisen (bijv. Lambada, HellaSwag, RACE), hoewel het nog steeds achter de Transformer-baseline blijft.
  • In-Context Recall: MesaNet overtreft andere RNN's en zelfs een Sliding Window Attention (SWA) model met een venster van 1024 tokens, maar blijft achter de full-attention Transformer.
  • Few-Shot Learning: MesaNet vertoont sterke prestaties op taken zoals word scrambling en vertaling, waarbij het vaak Transformers overtreft op scrambling-taken.
• Efficiëntie: Ondanks dat het extra FLOP's vereist voor de CG-solver tijdens inferentie, behoudt MesaNet concurrerende token-throughput op H100 GPU's en TPUv5's, vooral wanneer een gematigd aantal CG-stappen (bijv. 10-30) wordt gebruikt.

Betekenis en Claims

Het artikel positioneert MesaNet als een belangrijke stap richting optimale test-time training binnen neurale netwerken. De auteurs beweren dat door sequentiële optimalisatieproblemen (specifiek het vinden van de beste lineaire associatieve geheugen-gewichten) tijdens de inferentie op te lossen, het model effectiever kan adapteren aan nieuwe bewijslast dan modellen die eenvoudige gradiëntupdates gebruiken.

Het werk benadrukt een afweging: de prestatiewinsten in perplexiteit en downstream redenering komen met de kosten van verhoogde inferentie-compute (FLOP's). De auteurs suggereren dat dit aansluit bij recente trends om de prestaties te verbeteren door de test-time compute te verhogen. Zij stellen dat de Mesa-laag een flexibel mechanisme biedt om deze compute dynamisch toe te wijzen, wat potentieel complexe sequentiële problemen binnen het netwerk zelf kan oplossen.

Het artikel concludeert dat hoewel MesaNets momenteel een compute-overhead ervaren, hun vermogen om andere RNN's te overtreffen op taken met lange-context begrip en globale redenering suggereert dat lokale optimale test-time training een veelbelovende richting is voor efficiënte sequentie-modellering, die potentieel de kloof kan overbruggen tussen de efficiëntie van RNN's en de prestaties van Transformers. De auteurs erkennen dat de huidige implementatie van de compute-kosten een nadeel is en suggereren dat toekomstig werk zich kan richten op warm-start optimalisaties of hybride inferentiestrategieën om dit te mitigeren.