TIDE: Token-Informed Depth Execution for Per-Token Early Exit in LLM Inference

Probleemstelling

Grote Taalmodellen (LLMs) behandelen momenteel elk token uniform door alle lagen van het netwerk te doorlopen, ongeacht de complexiteit van het token. Een functioneel woord zoals "de" krijgt dezelfde 32-lagen verwerking als een complexe redeneerstap in een wiskundige afleiding. Deze uniforme computetoewijzing is inefficiënt en kostbaar:

• Verspilling: Onderzoek toont aan dat voor een groot deel van de tokens de verborgen staten (hidden states) al vroeg in het proces convergeren naar de uiteindelijke uitkomst. De resterende lagen voegen weinig tot niets toe aan de output-distributie.
• Kosten: Het serveren van grote modellen (bijv. 70B parameters) is duur en vertraagt de inferentie. Prefill-latentie en autoregressieve decoding domineren de totale doorlooptijd.
• Beperkingen van bestaande methoden: Bestaande "early exit"-technieken zijn vaak beperkt tot encoder-modellen (zoals BERT), vereisen dure hertraining van het model, of vertrouwen op onbetrouwbare heuristieken zoals softmax-entropie, die bij generatie vaak hoog is en dus geen goed signaal geeft.

Methodologie: TIDE

TIDE (Token-Informed Depth Execution) is een post-training systeem dat per-token vroege exit mogelijk maakt zonder het onderliggende model te hoeven hertrainen of modificeren.

1. Werkingsprincipe:
Het systeem plaatst kleine, geleerde routers (classificatoren) op periodieke checkpoint-lagen. Tijdens de inferentie wordt voor elk token bepaald op welke laag de verborgen staat "geconvergeerd" is. Het token stopt dan met verwerken en gebruikt de output van die laag voor de volgende stap.

2. Calibratie (Offline):

• Data: Het systeem gebruikt een klein datasetje (2.000 WikiText-teksten) om staten te verzamelen van een bevroren model.
• Signaal: In plaats van entropie, gebruikt TIDE de cosine-相似iteit tussen de verborgen staat op een checkpoint-laag ($h_k$) en de finale laag ($h_L$). Als deze相似iteit een drempel ($\tau$, standaard 0.98) overschrijdt, wordt het token als "geconvergeerd" beschouwd.
• Training: Er worden lichte binary classifiers (routers) getraind, bestaande uit een twee-laags MLP met een bottleneck van 128 dimensies. Dit proces duurt minder dan 3 minuten op één GPU en resulteert in een checkpoint van slechts ~4 MB.

3. Inferentie (Online):

• Post-Hoc Evaluatie: Tijdens de generatie draait het model de volledige forward pass (alle lagen) om de KV-cache (Key-Value cache) intact te houden en corruptie te voorkomen.
• Routering: Na de forward pass evalueren de routers de verborgen staten van de checkpoint-lagen. De router selecteert de eerste laag waarvoor de score boven de drempel $\theta$ ligt.
• Output: De logits worden berekend op basis van de verborgen staat van die vroegere laag, in plaats van de laatste laag.

4. Technische Implementatie:

• Universele Adapter: TIDE bevat een UniversalAdapter die automatisch de structuur van 17 verschillende HuggingFace-architecturen (LLaMA, GPT-2, Phi, Falcon, etc.) detecteert zonder specifieke code per model.
• CUDA Kernels: Voor prestaties zijn er gefuseerde CUDA-kernels ontwikkeld die RMSNorm en router-evaluatie in één launch uitvoeren. Dit ondersteunt float32, float16 en bfloat16.
• Compatibiliteit: Het werkt met elke versie van de transformers-library en vereist geen aanpassingen aan de modelcode.

Belangrijkste Bijdragen

1. Post-training Systeem: Een early-exit systeem dat geen modelhertraining vereist en werkt met bestaande checkpoints.
2. Per-Token Adaptiviteit: Het systeem herkent dat convergentie een eigenschap is van het token en niet van het hele model, wat leidt tot fijnere optimalisatie dan globale heuristieken.
3. Universele Compatibiliteit: Ondersteuning voor een breed scala aan architecturen via automatische detectie.
4. Efficiënte Implementatie: Gefuseerde CUDA-kernels en een open-source release (3.097 regels code, 74 tests) die direct via PyPI beschikbaar is.
5. KV-Cache Integriteit: In tegenstelling tot andere methoden die de cache kunnen beschadigen, behoudt TIDE de cache-integriteit door alle lagen te draaien maar alleen de output te selecteren.

Resultaten

De evaluatie vond plaats op een NVIDIA A100 met modellen zoals DeepSeek R1 Distill 8B en Qwen3 8B.

• Exit Rates:
  • Prefill: 100% van de tokens vindt een exit-punt. Bij DeepSeek R1 verlaten 5% van de tokens het model al bij laag 11, de rest bij laag 31.
  • Decoding: 98–99% van de tokens tijdens autoregressieve generatie stopt vroeg (meestal bij laag 31 van 32).
• Performance:
  • Latentie: Reductie van prefill-latentie met 5,5% tot 7,2%.
  • Throughput: Een toename van 6,6% bij batch size 1 en 8,1% bij batch size 8 (voor Qwen3).
  • Opmerking: Bij DeepSeek R1 met batch size 8 was er een negatieve schaling (-16,3%), waarschijnlijk door de overhead van het verzamelen van alle hidden states bij grote batches.
• Kwaliteit: De kwaliteit van de output bleef intact. Bij een meervoudige wiskundige opgave (256 tokens) behield het model de juiste redenering en unieke tokens, zelfs met een exit-ratio van 99,6%.

Betekenis en Toekomstperspectief

TIDE bewijst dat per-token vroege exit haalbaar is voor autoregressieve LLM-inferentie zonder de complexiteit van modelhertraining of auxiliaire modellen (zoals bij speculative decoding).

• Efficiëntie: Het biedt directe besparingen in latentie en energie, wat cruciaal is voor het schalen van LLM-diensten.
• Toekomstige verbeteringen: De huidige "post-hoc" aanpak draait nog alle lagen (geen fysieke overslaan van lagen in de tijd), wat de maximale snelheid beperkt. Toekomstig werk richt zich op "true skip mode" om lagen fysiek over te slaan en de cache-discontinuïteit te managen. Ook wordt onderzocht of dynamische drempels ($\tau$) per laag kunnen leiden tot nog vroegere exits.

Kortom, TIDE is een praktische, plug-and-play oplossing voor het versnellen van LLM-inferentie die breed toepasbaar is op bestaande modellen.