LoopMTP: A looped transformer guided by latent multi-token prediction

Technische Samenvatting: LoopMTP

Probleemstelling
Large Language Models (LLM's) hebben sterke redeneercapaciteiten getoond, die vaak gekoppeld zijn aan architectonische diepte. Echter, toenemende diepte vereist aanzienlijke aantallen parameters, wat onpraktisch is voor omgevingen met beperkte middelen (bijv. on-premise implementaties in de gezondheidszorg). Geloopte transformers bieden een parameterefficiënt alternatief door een enkele stapel lagen meerdere keren ($T$) te hergebruiken om effectieve diepte te simuleren. Ondanks hun potentieel lijden bestaande geloopte modellen onder twee kritieke beperkingen:

1. Latente Overthinking: Intermediaire representaties worden in elke iteratie overschreven. Als een voorspelling in een vroeg stadium correct is, kunnen daaropvolgende loops deze kunnen herzien tot een fout.
2. Onderscheideloze Computatie: Naarmate het aantal iteraties toeneemt, worden verborgen toestanden (hidden states) steeds gelijker, wat leidt tot redundante arbeid en verspilde rekenkracht.
   Bovendien ontbreekt het bestaande benaderingen aan sturing over de loops heen, en Multi-Token Prediction (MTP)—die vooruitziend redeneren stimuleert—heeft historisch gezien alleen voordelen getoond bij grote schaal of met dure volledige vocabulaire-projecties.

Methodologie: LoopMTP
De auteurs stellen LoopMTP voor, een geloopde transformer-architectuur die Multi-Token Prediction (MTP) integreert als een latent sturingssignaal om overthinking en onderscheideloze computatie aan te pakken. De kerncomponenten zijn:

• MTP-gestuurde Geloopte Blok: Het model loopt over de gehele transformer-stack $T$ keer. Om het overschrijven van nuttige informatie te voorkomen en de trajectorie te sturen, gebruikt de architectuur:

  • Iteratie-index Embedding: De genormaliseerde iteratie-index wordt aan token-vectoren toegevoegd om de huidige stap te signaleren.
  • Per-iteratie Normalisatie: LayerNorms worden specifiek toegepast op de toestand van elke iteratie om iteratie-specifieke statistieken af te handelen.
  • Input Fusie: De output van de vorige iteratie en de her-genormaliseerde token-embeddings worden geconcateneerd en geprojecteerd voordat ze de gedeelde backbone ingaan.
  • Loop-LNS (Looped Layer-Norm Scaling): Een aangepaste schaleringsfactor van $1/T$ wordt toegepast op de residu-stromen om stabiliteit te behouden over variabele unrollings, waardoor de neutralisatie van gradiënten wordt voorkomen die wordt waargenomen bij progressieve dieptescaling.

• Soft MTP Alignment (Latente Sturing): In plaats van volledige token-distributies te voorspellen (wat rekenintensief is), gebruikt LoopMTP een soft alignment objectief. Voor iteratie $t$ (waarbij $t \ge 2$) wordt de verborgen toestand uitgelijnd via cosinus-gelijkenis met de embedding van de token $t$ stappen vooruit ($u_{i+t}$). Dit fungeert als een representationele regularisator, die elke iteratie stuurt naar het anticiperen op specifieke toekomstige tokens zonder dat volledige vocabulaire-projecties nodig zijn. De eerste iteratie ($t=1$) blijft onbeperkt om het model toe te staan een rijke initiële representatie te leren.

• Geleerd Toestand-aggregatie: In plaats van intermediaire toestanden weg te gooien, aggregeert LoopMTP de outputs van alle $T$ iteraties met behulp van een token-gewogen som. Een lichtgewicht, gedeelde lineaire gate berekent per-iteratie gewichten, waardoor het model nuttige informatie uit vroege loops kan behouden terwijl het nieuwe berekeningen uit latere loops integreert. Een "ponder regularizer" (KL-divergentie) moedigt een uniforme distributie van aandacht over de iteraties aan om voortijdige convergentie te voorkomen.

• Trainingsdoelstelling: De totale loss combineert de standaard Next-Token Prediction (NTP) cross-entropy, de soft MTP alignment loss en de ponder regularizer.

Belangrijkste Bijdragen

1. Architectonische Innovatie: De introductie van LoopMTP, die structureel de link legt tussen het aantal loops en het anticiperen op toekomstige tokens, wat een natuurlijke remedie biedt voor latente overthinking en onderscheideloze computatie.
2. Efficiënte MTP Integratie: Het aantonen dat MTP effectief kan worden toegepast in parameterefficiënte geloopte modellen via een lichtgewicht cosinus-alignment signaal, waarbij de computationele overhead van standaard MTP wordt vermeden.
3. Stabiele Training: De voorgestelde architectuur maakt stabiele training mogelijk over een breed scala aan loop-aantallen (tot $T=15$) zonder de noodzaak van uitgebreide hyperparameter grid-searches per loop-aantal, een veelvoorkomend probleem bij eerdere geloopte modellen zoals LoopFormer.
4. Kleine Expert-modellen: Het aantonen dat LoopMTP kleine, domeinspecifieke expert-modellen (bijv. voor wiskunde) kan trainen die grotere niet-geloopte baselines overtreffen onder hetzelfde parameterbudget.

Resultaten
Geëvalueerd op een 260M parameter model getraind op 6.8B tokens:

• Prestatiewinsten: LoopMTP presteert tot 8,1% beter in relatieve gemiddelde nauwkeurigheid op algemene taken vergeleken met een parameter-gematchte niet-geloopte baseline. Het vermindert ook de Bits-Per-Byte (BPB) op QA, wiskunde en code taken met tot wel 7,5% vergeleken met de niet-geloopte baseline en presteert significant beter dan de state-of-the-art LoopFormer (winnend in 27 van de 28 gematchte vergelijkingen).
• Schaalbaarheid: In tegen tegenstelling tot LoopFormer, dat lijdt onder instabiele training en prestatievermindering bij hogere loop-aantallen, vertoont LoopMTP een monotone verbetering in perplexiteit en algemene taaknauwkeurigheid naarmate $T$ toeneemt.
• Mechanisme Analyse:
  • Representaties: Met MTP worden de verborgen toestanden over de iteraties heen meer onderscheidend, wat bevestigt dat het model unieke transformaties leert in plaats van redundante.
  • Ground-Truth Retrieval: De ground-truth token rang verbetert aanzienlijk (tot 35,6$\times$ beter) met MTP-sturing, wat aangeeft dat het alignment-signaal effectief wordt overgedragen naar de taalmodel-kop.
  • Gating: De geleerde gating-mechanisme past zich aan aan de diepte, waarbij latere iteraties hogere en variabele gate-waarden vertonen, wat wijst op een input-afhankelijke strategie voor informatiebehoud.
• Domeinexperts: In een gespecialiseerde wiskunde trainingssetting bereikte een LoopMTP-model met $T=15$ een nauwkeurigheid van 19,03% op GSM8K, een relatieve verbetering van ~170% ten opzichte van een niet-geloopte baseline (7,05%) bij hetzelfde aantal parameters.

Betekenis en Claims
Het artikel claimt dat LoopMTP erin slaagt de kloof te overbruggen tussen parameterefficiëntie en sterke redeneercapaciteiten. Door MTP te gebruiken als een latent sturend signaal, lost het de fundamentele instabiliteit en redundantie-problemen van eerdere geloopte transformers op. De auteurs benadrukken dat deze aanpak de inzet van sterke redeneermodellen op beperkte hardware (bijv. on-premise) mogelijk maakt, wat cruciaal is voor gevoelige domeinen zoals de gezondheidszorg waar data de locatie niet mag verlaten. Het werk suggereert dat "effectieve diepte" verkregen door looping, wanneer gestuurd door vooruitziende doelstellingen, de rauwe parameter-scaling kan overtreffen voor specifieke redeneertaken. De auteurs blijven bescheiden en merken op dat de voordelen in gespecialiseerde domeinen (buiten wiskunde) verdere verificatie vereisen en dat de bovengrenzen van nuttige loop-diepte een open vraag blijven.