Two-dimensional early exit optimisation of LLM inference

Probleemstelling

Grote Taalmodellen (LLM's) hebben de Natural Language Processing (NLP) gerevolutioneerd, maar hun inferentie vereist aanzienlijke rekenkracht, geheugen en energie. Voor relatief eenvoudige taken, zoals classificatie, is het gebruik van een volledig getraind model vaak overkill en inefficient. Bestaande oplossingen zoals early exit (vroege exit) strategieën proberen dit op te lossen door de berekening te stoppen in een tussenliggende laag zodra het model voldoende vertrouwen heeft in zijn voorspelling. Echter, deze methoden optimaliseren meestal slechts één dimensie: ofwel de diepte van het netwerk (laag-voor-laag), ofwel de lengte van de input (token-voor-token). Er ontbreekt een strategie die beide dimensies simultaan en synergetisch optimaliseert voor tekstuele input die natuurlijk in semantische eenheden (zinnen) kan worden opgesplitst.

Methodologie: 2D Vroege Exit Strategie

De auteurs introduceren een tweedimensionale (2D) vroege exit strategie die laag-voor-laag (layer-wise) en zin-voor-zin (sentence-wise) exit combineert.

1. Input Verwerking:

   • De inputtekst wordt opgesplitst in zinnen.
   • De inferentie verloopt iteratief: de eerste zin wordt verwerkt, daarna de tweede, enzovoort.
   • Dynamische Laagactivatie: Bij het verwerken van de $s$-de zin wordt het aantal actieve lagen in het model verhoogd. De diepte $L_s$ voor zin $s$ wordt berekend als:
     $$L_s = \min((s + 1) \times \Delta, L)$$
     waarbij $L$ het totaal aantal lagen is en $\Delta$ een stapgrootte is. Dit betekent dat elke nieuwe zin niet alleen de contextbreedte vergroot, maar ook de representatiediepte.

2. Architectuur en Training:

   • Er worden lichte classificatie-adapters (twee lagen volledig verbonden netwerken) toegevoegd aan elke laag van het LLM.
   • Training: Er zijn twee modi getest:
     • Adapter-only training: Alleen de adapters worden getraind, de LLM-weights blijven bevroren.
     • Fine-tuning: Het volledige model (inclusief adapters) wordt gefine-tuned met een gewogen verliesfunctie die zowel volledige als partiële input (prefixen van zinnen) belooft.
   • Embedding: De embedding van een zin wordt berekend als het gemiddelde van de token-embeddings binnen die zin. Voor de vroege exit wordt het gemiddelde van de embeddings van de eerste $j$ zinnen gebruikt.

3. Inferentie en Stopcondities:

   • De algoritme (Algorithm 1) verwerkt zinnen sequentieel.
   • Vertrouwen (Confidence): Voor elke laag en zin wordt het vertrouwen berekend als het verschil tussen de hoogste en de op één na hoogste softmax-kans.
   • Accumulatie: Als het vertrouwen een drempel ($\tau_{ignore}$) overschrijdt, wordt dit toegevoegd aan een cumulatieve score.
   • Vroege Stop: Zodra de cumulatieve score van een klasse een acceptatiedrempel ($\tau_{acc}$) bereikt, stopt de inferentie onmiddellijk en wordt de voorspelling teruggegeven.
   • Als geen enkele drempel wordt bereikt na het verwerken van alle zinnen en lagen, wordt het resultaat van de laatste laag van de laatste zin gebruikt.

Belangrijkste Bijdragen

1. Methodologische Innovatie: Integratie van laag-voor-laag stop en zin-voor-zin inputtrimming in één synergetisch proces. Dit resulteert in multiplicatieve (in plaats van additieve) rekenbesparingen.
2. Uitgebreide Benchmarking: Evaluatie van de 2D-strategie tegen optimale laag-voor-laag early exit en de LayerSkip-methode op vier state-of-the-art LLM's (Llama 3.1/3.2, Gemma, Qwen) met 3B tot 8B parameters.
3. Analyse van Fine-tuning: Een diepgaande analyse van hoe fine-tuning de voordelen van 2D early exit beïnvloedt, met de ontdekking dat fine-tuning de variatie in nauwkeurigheid over de lagen "plat" maakt, waardoor het vroege exit-potentieel afneemt.

Resultaten

De experimenten zijn uitgevoerd op drie sentimentanalyse datasets met verschillende moeilijkheidsgraden: Steam (2 klassen), MMS (3 klassen) en Amazon-5 (5 klassen).

• Snelheidswinst:
  • Voor eenvoudigere taken (Steam en MMS) met "vanilla" (niet gefine-tuned) modellen boekte de 2D-strategie een extra snelheidswinst van 1,4x tot 2,3x ten opzichte van de optimale laag-voor-laag early exit.
  • Voor complexere taken (Amazon-5) was de winst kleiner (tot 1,5x) of zelfs negatief voor bepaalde modellen (Llama), wat suggereert dat de strategie werkt wanneer semantische informatie voorspelbaar accumuleert over zinnen.
• Invloed van Fine-tuning:
  • Fine-tuning verbeterde de algehele nauwkeurigheid, maar verminderde de snelheidswinst van de 2D-strategie aanzienlijk (bijv. van 2,3x naar 1,9x op de MMS dataset).
  • De auteurs stellen dat fine-tuning de nauwkeurigheidsgradiënten over de lagen en deels van de input "vereffent", waardoor het model minder geneigd is om vroeg een betrouwbare beslissing te nemen.
• Vergelijking met LayerSkip: De 2D-strategie presteerde beter dan LayerSkip op de Steam en MMS datasets, maar LayerSkip deed het beter op de moeilijkste dataset (Amazon-5) bij een lage foutmarge (2%).

Betekenis en Conclusie

Dit paper toont aan dat er een significant potentieel ligt in het coördineren van exit-strategieën over meerdere dimensies (diepte en inputlengte) voor LLM-inferentie.

• Toepasbaarheid: De methode is model-agnostisch, vereist slechts lichte adapters en is orthogonaal (onafhankelijk) van andere efficiëntietechnieken zoals kwantisatie, pruning en Mixture-of-Experts.
• Beperkingen: De strategie werkt het beste bij taken waar de betekenis van de tekst lineair en voorspelbaar opbouwt (zoals sentimentanalyse). Bij complexe, meerdimensionale taken of bij modellen die zwaar gefine-tuned zijn, neemt het voordeel af.
• Toekomstperspectief: De auteurs pleiten voor het ontwikkelen van fine-tuning strategieën die de "gradiënten" voor vroege exit behouden, en voor het uitbreiden van de methode naar andere sequentieverwerkingstaken zoals intentie-detectie en vraagbeantwoording.

Kortom, de 2D vroege exit biedt een krachtige route om de inferentiekosten van LLM's te verlagen voor specifieke classificatietaken, mits de taakcomplexiteit en het trainingsprotocol hierop zijn afgestemd.