Technische Samenvatting: LEDE – Ervaringsgestuurde Dynamische Exits voor LLM's met Reinforcement Learning

1. Probleemstelling

Large Language Models (LLM's) lijden onder aanzienlijke latentie tijdens autoregressieve inferentie omdat elke gegenereerde token de volledige stack van transformer-lagen moet doorlopen. Hoewel Self-Speculative Decoding (SSD) is voorgesteld om dit proces te versnellen door de eigen vroege lagen van een model te gebruiken als een interne "drafter" om kandidaat-tokens parallel te genereren, zijn huidige SSD-methoden fundamenteel beperkt door statische configuraties.

Bestaande benaderingen vertrouwen op vaste draft-dieptes en vaste speculatie-lengtes die uniform over het gehele generatieproces worden toegepast. Deze rigiditeit houdt geen rekening met contextuele sparsiteit, waarbij de moeilijkheid van het voorspellen van tokens varieert afhankelijk van de lokale context. Bijgevolg kunnen statische methoden de trade-off tussen draft-kwaliteit (acceptatiepercentage) en speculatiesnelheid niet optimaal balanceren, wat leidt tot suboptimale inferentieprestaties.

2. Methodologie: Het LEDE-framework

De auteurs stellen LEDE (Learning-based Dynamic Exit) voor, een framework dat de dynamische controle van SSD herformuleert als een Markov-beslissingsproces (MDP). In plaats van rigide, vooraf gedefinieerde regels te gebruiken, zet LEDE offline reinforcement learning (RL) in om een beleid (policy) te trainen dat zowel de optimale exit-laag als de speculatie-lengte dynamisch selecteert op basis van de lokale context van de gegenereerde sequentie.

2.1. Formulering van het Markov-beslissingsproces

Het framework modelleert het besluitvormingsproces tijdens een enkele generatiestap als een MDP met de volgende componenten:

• Toestandsruimte ($S$): De agent observeert een toestandsvector $s_l$ bij elke kandidaat-laag $l$. Deze vector wordt afgeleid van de interne token-distributie van het model en bevat:
  • Token Betrouwbaarheid ($C_k^l$): De gemiddelde negatieve log-waarschijnlijkheid van de top-$k$ tokens.
  • Token Entropie ($H^l$): De onzekerheid van de volgende-token voorspellingsdistributie.
  • Gemiddelde Trace van Token Betrouwbaarheid ($C_{trace}$): De gemiddelde negatieve log-waarschijnlijkheid van de top-1 token over een glijdend venster van voorgaande stappen.
• Actieruimte ($A$): Bij elke kandidaat-laag beslist de policy (een Q-netwerk) tussen twee acties:
  • Doorgaan ($a=0$): Doorgaan naar de volgende kandidaat-laag.
  • Exiten ($a=1$): De huidige laag aanwijzen als de optimale exit-laag ($l^*$) voor de draft-fase.
• Beloningsfunctie ($R$): Ontworpen om de inferentie-efficiëntie te optimaliseren.
  • Positieve Beloning: Wordt gegeven als de draft-token wordt geaccepteerd ($x_t^{l^*} = x_t^L$), geschaald met de bespaarde diepte ($1 + \frac{L-l^*}{L}$).
  • Negatieve Beloning: Wordt gegeven als de draft-token wordt afgewezen ($-2 + \frac{L-l^*}{L}$).
  • Straf: Een kleine straf ($-0,01$) wordt toegepast voor het doorgaan, wat vroegtijdig exiten stimuleert wanneer passend.

2.2. Policy Optimalisatie via Offline RL

LEDE maakt gebruik van een Deep Q-Network (DQN) die offline is getraind met behulp van een replay buffer van eerdere inferentie-ervaringen:

1. Dataverzameling: Een initiële exploratiefase verzamelt diverse ervaringstuppels $(s_t, a_t, r_t, s_{t+1})$ waarbij de agent met de omgeving interacteert.
2. Offline Training: Het Q-netwerk wordt getraind door de n-staps Temporal Difference (TD) error te minimaliseren. De doelwaarde wordt berekend met een stabiel target-netwerk, waarbij een balans wordt gezocht tussen de directe beloning van een vroege exit en de lange-termijnwaarde van de toestand.
3. Dual-Control Mechanisme: Tijdens online inferentie voert de getrainde agent twee adaptieve controles uit:
   • Dynamische Selectie van de Exit-laag: Het bepaalt de optimale diepte $l^*$ om de drafting te starten.
   • Adaptieve Drafting Lengte: Na het selecteren van $l^*$ evalueert de agent de toestand opnieuw na elke draft-token. Een "exit" actie zet het draften voort (indicatie van hoge betrouwbaarheid), terwijl een "continue" actie (indicatie van lage betrouwbaarheid) het draften onmiddellijk beëindigt voor verificatie.

3. Belangrijkste Bijdragen

Het paper beschrijft drie primaire bijdragen:

1. MDP Formulering: Het is het eerste werk dat de dynamische controle van draft-diepte en speculatie-lengte in SSD formuleert als een MDP-probleem, waarbij reinforcement learning wordt toegepast om een controlebeleid af te leiden.
2. LEDE Framework: Het ontwerp en de implementatie van een framework dat gebruikmaakt van offline RL en experience replay om een context-bewust beleid te leren, waarmee de beperkingen van zowel statische configuraties als eenvoudige heuristiek-gebaseerde methoden worden overwonnen.
3. Empirische Validatie: Uitgebreide experimenten die significante versnellingen laten zien ten opzichte van autoregressieve generatie en substantiële efficiëntiewinsten ten opzichte van statische baselines (zoals LayerSkip).

4. Experimentele Resultaten

Het framework werd geëvalueerd op Llama-2 (7B, 13B) en Llama-3 (1B) modellen, evenals CodeLLaMA (7B, 34B), over diverse taken inclusief instructie-opvolging, samenvatting, taalmodellering en codegeneratie.

• Versnellingsprestaties:
  • LEDE bereikt een gemiddelde versnelling van 2,32× ten opzichte van standaard autoregressieve decodering.
  • Pieksnelheden bereiken 2,72× (bijv. op Llama-2-7B voor taalmodellering).
  • LEDE biedt een extra 17% versnelling ten opzichte van de beste statische speculatieve baseline (LayerSkip).
• Vergelijking met Baselines:
  • LayerSkip (Statisch): Prestaties zijn zeer gevoelig voor vaste configuraties en generaliseren slecht over verschillende taken.
  • Regelgebaseerde Methoden (LITE, Draft & Verify): Beperkt door vooraf gedefinieerde drempelwaarden, wat leidt tot suboptimale trade-offs tussen diepte en acceptatiepercentages.
  • LEDE: Door een fijnmazig controlebeleid te leren, presteert het consistent beter dan alle baselines qua versnelling, terwijl het een vergelijkbare generatiekwaliteit behoudt (gemeten via Rouge-L).
• Ablatie Studie: Het verwijderen van ofwel de dynamische selectie van de exit-laag of de adaptieve drafting-lengte verslechtert de prestaties aanzienlijk (waardoor de versnelling daalt naar ~2,0×), wat bevestigt dat beide mechanismen synergetisch en essentieel zijn.

5. Betekenis en Claims

Het paper claimt dat LEDE een conceptuele verschuiving vertegenwoordigt van rigide, door heuristiek gestuurde inferentie naar adaptieve, door beleid gestuurde inferentie. Door te leren de computationele diepte en speculatie-lengte te co-optimaliseren op basis van de realtime context, maakt het framework een robuustere en efficiëntere decodering mogelijk over diverse modellen en taken.

De auteurs positioneren dit werk als een stap naar het bouwen van intelligentere en computationeel efficiëntere LLM-systemen. Zij merken op dat toekomstig werk zich zal richten op het opschalen van deze voordelen naar grotere modellen (70B schaal en hoger), maar de huidige bijdrage vestigt de levensvatbaarheid van het gebruik van offline reinforcement learning om de inferentie-pipeline van grote taalmodellen dynamisch te beheren.