MACRO: Markov Chain Routing of Transformer Layers

Technische Samenvatting: MACRO – Markov Chain Routing van Transformer-lagen

Probleemstelling
Standaard Large Language Models (LLMs) voeren transformer-blokken sequentieel uit, van de eerste naar de laatste laag, volgens een rigide computationeel schema. Deze statische aanpak negeert het potentieel voor dynamische computationele paden die de redeneernauwkeurigheid bij complexe taken zouden kunnen verbeteren zonder de onderliggende gewichten van het model te wijzigen. Eerder werk heeft adaptieve diepte verkend, primair voor inferentieversnelling (bijv. early exits, gestructureerd weg laten van lagen), maar het actief herstructureren van de executieroute van een bevroren model om prestaties te verbeteren, blijft een uitdagende grens. Bestaande methoden voor dynamische routing kampen met significante knelpunten: sommige vereisen dure offline Monte Carlo Tree Search (MCTS) voor datageneratie en training (bijv. Dr.LLM), andere vereisen ground-truth labels tijdens inferentie, en sommige vertrouwen op per-sample zoekloops die computationeel onhaalbaar zijn.

Methodologie
De auteurs stellen MACRO (Markov Chain Routing of Transformer Layers) voor, een framework dat taakspecifieke executieroutes leert over bevroren LLM-architecturen. MACRO modelleert laag-routing als een context-afhankelijk Markov-beleid, $\pi_\theta(a_t | s_t)$, dat de sequentie van laagtoepassingen, het overslaan, herhalen en toevoegen van verborgen toestanden (hidden states) dicteert.

• Toestands- en Actieruimte: De toestand $s_t$ is een compacte samenvatting bestaande uit de huidige laagindex ($\ell$), de resterende computationele budgetfase (vroeg/midden/laat), de binnenkomende laag-verplaatsing ($\delta$), en de operator-context (of de huidige verborgen toestand is geproduceerd door een gewone laag of een additie). De actieruimte omvat:

  1. Lokale Laagbeweging: Het toepassen van een nabijgelegen laag $\ell + \Delta$ (wat het overslaan, herhalen of achterwaarts bezoeken toestaat).
  2. Add-and-Apply: Het toevoegen van een eerdere verborgen toestand $h_{t-b}$ aan de huidige toestand $h_t$ (met een coëfficiënt $\gamma > 0$) vóór het toepassen van een lokale laag.
  3. Rejoin: Het beëindigen van de geleerde prefix en het uitvoeren van de resterende standaard suffix van de transformer.
     Haalbaarheidsmaskers zorgen ervoor dat alle gesamplede programma's binnen de laaggrenzen en budgetbeperkingen blijven.

• Leerprocedure: MACRO maakt gebruik van een estimation-of-distribution benadering. Vertrekkend vanuit een forward-biased prior (waarbij het standaard sequentiële pad het meest waarschijnlijk is), samplet het framework iteratief kandidaat-routingprogramma's, evalueert deze op een subset van de trainingsdata en werkt de Markov-transitiematrix bij met behulp van weighted maximum likelihood. Hoog-presterende routes worden exponentieel gewogen om het beleid te verfijnen. Een replay buffer zorgt voor stabiele convergentie door de dekking van voorbeelden uit vorige iteraties te behouden.

• Decoding en Selectie: In tegenstelling tot methoden die vertrouwen op stochastische sampling of per-test-time search, gebruikt MACRO een top-k Viterbi-algoritme om deterministisch de meest waarschijnlijke geldige routingprogramma's te decoderen uit het geleerde beleid. Structurele haalbaarheidsmaskers zijn geïntegreerd in de Viterbi-recursie. De uiteindelijke route wordt geselecteerd door de top-k kandidaten te evalueren op een held-out validatieset, waarbij de route met de hoogste nauwkeurigheid wordt gekozen.

Belangrijkste Bijdragen

1. Parameter-efficiënt Framework: MACRO introduceert een lichtgewicht, context-afhankelijk Markov-beleid voor adaptieve laag-routing in bevroren en gekwantiseerde LLM's, wat geen updates aan de gewichten van het model vereist.
2. Exact Decoding Schema: Het artikel presenteert een exact top-k Viterbi-decodeschema uitgerust met structurele haalbaarheidsmaskers. Dit elimineert de noodzaak voor per-sample test-time search, aanvullende neurale netwerktraining, of ground-truth inferentie-orakels.
3. Empirische Superieuriteit: Het framework vertoont consistente verbeteringen ten opzien van niet-gerouteerde baselines en bestaande dynamische routingmethoden over diverse redeneer- en kennisbenchmarks.

Resultaten
Geëvalueerd over meerdere open-weight LLM's (inclusief Qwen3-varianten, Mixtral en DeepSeek-R1-Distill) op benchmarks zoals GSM8K, MATH500, MedQA en MMLU-Pro:

• Nauwkeurigheidswinst: MACRO bereikt een gemiddelde verbetering in nauwkeurigheid van +5,0% ten opzien van ongerouteerde baselines. De winst is het meest uitgesproken bij kleinere modellen (bijv. +12,68% op Qwen3-1.7B).
• Vergelijking met SOTA: MACRO presteert +7,2% beter dan de leidende dynamische routingmethode, Dr.LLM, in gemiddelde nauwkeurigheid.
• Efficiëntie: Het route-zoekproces is aanzienlijk efficiënter, waarbij de zoektijd met 9,4× wordt verminderd (van 14,8 uur naar 1,6 uur per benchmark) vergeleken met Dr.LLM.
• Robuustheid: Statistische significantietesten bevestigen dat MACRO de prestaties verbetert over alle geteste modellen zonder enige benchmark te verslechteren.
• Mechanistisch Inzicht: Logit-lens analyse suggereert dat routing de prestaties niet verbetert door de interne redeneercapaciteiten van het model te veranderen, maar door antwoorden die al intern gerepresenteerd zijn beter bloot te leggen, waardoor ze niet worden onderdrukt nabij de output.

Betekenis en Claims
Het artikel claimt dat MACRO een belangrijke stap voorwaarts is in het praktisch maken van dynamische computation voor bevroren LLM's. Door routing te formaliseren als een Markov-proces, overwinnen de auteurs de computationele kosten en datavereisten van eerdere methoden (zoals MCTS-gestuurde training). Het werk betoogt dat routing dient als een mechanisme om reeds aanwezige antwoorden effectiever te "eliciteren" in plaats van de geleerde representaties fundamenteel te wijzigen. De auteurs stellen dat deze aanpak verder uitgebreid kan worden om input-afhankelijke dynamische routing mogelijk te maken, wat een nieuwe dimensie biedt voor flexibele test-time scaling voor LLM's. De methode is bijzonder opmerkelijk vanwege haar vermogen om de prestaties op kleine modellen te verbeteren en haar compatibiliteit met gekwantiseerde modellen, wat wijst op een brede toepasbaarheid in omgevingen met beperkte middelen.