Task-Centric Acceleration of Small-Language Models

Probleemstelling

Kleine Taalmodellen (SLMs) zijn efficiënte alternatieven voor grote Taalmodellen (LLMs), vooral voor taakspecifieke toepassingen zoals tekstclassificatie, samenvattingen en vraag-antwoordtaken. Deze modellen worden vaak ingezet in omgevingen met hoge doorvoer en lage latentie-eisen (bijvoorbeeld op apparaten of in multi-agent systemen).

De huidige optimalisatiemethoden voor LLMs richten zich voornamelijk op het verminderen van geheugeneisen (quantisatie) of het versnellen van de inputverwerking (KV-caching, attention-optimalisatie) om langere contexten te hanteren. Echter, voor SLMs die zijn gefinetuned op specifieke taken, is de output-variabiliteit vaak aanzienlijk lager dan bij open-ended generatie. Dit betekent dat de output-takken van deze modellen sterk geconcentreerd zijn rond een klein aantal veelvoorkomende n-grams (woordenreeksen).

Het paper identificeert dat de huidige inefficiëntie in autoregressieve generatie voortkomt uit het aantal decoding-stappen dat nodig is om een antwoord te genereren. Bestaande methoden zoals speculative decoding vereisen vaak een "drafter"-model dat dezelfde vocabulaire deelt als het doelmodel, wat de toepasbaarheid beperkt. Het paper stelt de vraag: hoe kunnen we de inferentie-efficiëntie van SLMs maximaliseren door gebruik te maken van de lage output-variabiliteit en de concentratie van specifieke n-grams in de taakoutput?

Methodologie: TASC Framework

De auteurs introduceren TASC (Task-Adaptive Sequence Compression), een raamwerk dat twee verschillende methoden biedt afhankelijk van de gebruiksscenario's:

1. TASC-ft (Tokenizer Augmentation tijdens Finetuning)

Deze methode wordt toegepast wanneer een SLM wordt gefinetuned op een specifieke taak.

• Concept: Het uitbreiden van de tokenizer-vocabulaire met taakspecifieke n-grams.
• Proces:
  • Er wordt een "merge reward" berekend voor veelvoorkomende n-grams in het trainingscorpus van de verwachte outputs. De reward is gebaseerd op de frequentie en de potentiële reductie in token-aantal (bijvoorbeeld: het samenvoegen van een bigram bespaart 1 token).
  • Om te voorkomen dat het samenvoegen van tokens de kans op alternatieve continuaties onterecht onderdrukt (prefix-collisie), wordt een Prefix Collision Score (PCS) gebruikt. Alleen n-grams met een lage PCS-score worden toegevoegd.
  • De tokenizer wordt iteratief verrijkt met deze nieuwe tokens tot een bepaald budget is bereikt.
  • De embeddings voor deze nieuwe tokens worden initialiseerd als het gemiddelde van de embeddings van de oorspronkelijke tokens en vervolgens geoptimaliseerd tijdens het finetuning-proces (vaak met QLoRA).
• Doel: Het verminderen van het aantal decoding-stappen door veelvoorkomende zinsdelen als één token te behandelen.

2. TASC-spec (Inference-time Acceleratie zonder Training)

Deze methode is bedoeld voor scenario's waar geen finetuning mogelijk is of wanneer een kant-en-klaar gefinetuned model wordt gebruikt.

• Concept: Een training-vrije, lichte speculative decoding methode die gebruikmaakt van n-gram modellen.
• Proces:
  • Er wordt een "drafter" model gebouwd dat puur bestaat uit n-gram statistieken afgeleid van het corpus van de taakoutput ($E_\tau$).
  • Dit drafter model combineert globale informatie (uit het volledige corpus) met lokale contextinformatie (uit de huidige prompt). Dit gebeurt via een gemengde verdeling: $p_{mixed} = \lambda p_{corpus} + (1-\lambda) p_{prompt}$.
  • Het drafter genereert een reeks tokens (draft) met verwaarloosbare kosten.
  • Het doelmodel (target) verifieert deze reeks in één enkele forward pass.
• Voordelen: Geen extra training nodig, geen vereiste dat het drafter en doelmodel dezelfde vocabulaire delen (tokenizer-agnostisch), en zeer licht van gewicht.

Kernbijdragen

1. Karakterisering van Output-variabiliteit: De auteurs formaliseren het concept dat taakspecifieke SLMs een lage output-variabiliteit hebben, wat leidt tot een kleine "typische set" van n-grams. Dit wordt kwantitatief gemeten via entropie en de grootte van de typische set.
2. TASC-ft: Een algoritme dat de tokenizer van een SLM verrijkt met taakspecifieke n-grams, naadloos geïntegreerd in parameter-efficiënt finetuning. Dit resulteert in aanzienlijke latency-reducties zonder kwaliteitsverlies.
3. Predictieve Metriek: De auteurs tonen aan dat de 2-Rényi entropie van de token-distributie een betrouwbare voorspeller is voor het potentieel van versnelling via TASC-ft. Een lagere entropie correleert sterk met een grotere mogelijke versnelling.
4. TASC-spec: Een training-vrije speculative decoding methode die gebruikmaakt van n-gram drafters. Deze methode overwint de beperkingen van bestaande methoden (zoals vocabulaire-afstemming) en levert significante snelheidswinsten op.

Resultaten

De methoden zijn getest op diverse datasets, waaronder EUR-LEX (juridische classificatie), Asclepius (medische QA), en MASSIVE (intentie-classificatie).

• TASC-ft Resultaten:

  • Bereikte een reductie in runtime van maximaal 2.1x (bijv. op de EUR-LEX dataset).
  • De taakprestaties (F1-score, Bert-score, etc.) bleven behouden of verbeterden lichtjes.
  • De toevoeging van ongeveer 500 nieuwe tokens bleek vaak het optimale punt te zijn voor de afweging tussen versnelling en trainingskosten.
  • De 2-Rényi entropie bleek een zeer nauwkeurige voorspeller (92% succesrate) voor de richting van de runtime-reductie.

• TASC-spec Resultaten:

  • Bereikte snelheidswinsten tot 3.15x op de Asclepius dataset en 2.17x op EUR-LEX.
  • Dit is significant beter dan bestaande n-gram methoden (zoals n-grammys), die vaak slechts een kleine winst opleveren omdat ze niet taak-geoptimaliseerd zijn.
  • De acceptatiegraad van de gedraaide tokens was hoog in de beginposities van de generatie, wat de efficiëntie bevestigt.

Betekenis en Conclusie

Dit paper biedt een fundamenteel nieuw perspectief op het versnellen van Small Language Models. In plaats van te focussen op hardware-optimalisatie of complexe architecturale wijzigingen, benut TASC de inherente statistische eigenschappen van taakspecifieke outputs (lage variabiliteit).

De significance ligt in de dualiteit van de aanpak:

1. Voor ontwikkelaars die modellen kunnen finetunen, biedt TASC-ft een manier om de modelarchitectuur zelf te optimaliseren voor specifieke domeinen.
2. Voor gebruikers die kant-en-klare modellen willen gebruiken, biedt TASC-spec een plug-and-play oplossing voor versnelling zonder extra training of zware drafter-modellen.

De methode is bijzonder relevant voor edge-deployments en real-time systemen waar lage latentie cruciaal is, en demonstreert dat eenvoudige statistische benaderingen (n-grams) nog steeds krachtig kunnen zijn in het tijdperk van transformer-modellen, mits slim toegepast op de juiste use-cases.