Semantic Token Clustering for Efficient Uncertainty Quantification in Large Language Models

Probleemstelling

Grote Taalmodellen (LLMs) presteren uitstekend op diverse taken, maar hun output is niet altijd feitelijk juist. Een kritieke beperking is hun neiging tot oververzekerdheid: ze produceren vaak plausibele maar onjuiste antwoorden met een hoge zekerheid. Om dit te adresseren is onzekerheidskwantificatie (uncertainty quantification) essentieel om onbetrouwbare output te identificeren.

Bestaande methoden voor onzekerheidskwantificatie hebben echter twee grote nadelen:

1. Hoge rekentijd: Veel methoden (zoals Semantic Entropy) vereisen herhaaldelijk genereren (sampling) van meerdere antwoorden om semantische diversiteit te meten.
2. Afhankelijkheid van externe modellen: Sommige methoden (zoals Claim Conditioned Probability) gebruiken externe Natural Language Inference (NLI) modellen, wat de complexiteit en kosten verhoogt.

Er is dus behoefte aan een methode die efficiënt is (geen herhaald genereren), zelfstandig werkt (geen externe modellen) en de semantische structuur van het LLM optimaal benut.

---

Methodologie: Semantic Token Clustering (STC)

De auteurs stellen Semantic Token Clustering (STC) voor, een methode die onzekerheid kwantificeert op basis van één enkele generatie, door gebruik te maken van de interne semantische representaties van het LLM.

De methode bestaat uit twee fasen:

1. Pre-computatie Fase (Offline)

• Token Embedding Clustering: Alle tokens in het vocabulaire van het model worden gegroepeerd in semantisch consistente clusters. Dit gebeurt met een onbewaakte clustering-algoritme (Agglomerative Clustering) op basis van de token-embeddings (combinatie van input- en output-embeddings).
• Doel: Tokens die semantisch vergelijkbaar zijn (bijv. "TV", "tv", "television") worden in dezelfde cluster geplaatst.
• Voordeel: Deze stap is eenmalig en offline, waardoor er tijdens de inferentie geen extra rekentijd nodig is voor het clusteren zelf.

2. Inferentie Fase (Online)

Tijdens het genereren van een antwoord wordt voor elk volgende token de onzekerheid als volgt berekend:

• Aggregatie van Kansmassa: In plaats van alleen naar de kans van het exact gegenereerde token te kijken, wordt de totale kansmassa opgeteld over alle tokens binnen de bijbehorende semantische cluster.
  • Reden: Als een model "TV" kiest, maar "television" had ook een hoge kans gehad, zou de kans op "TV" alleen de zekerheid van het model onderschatten. Door de kansen van alle semantisch equivalente tokens op te tellen, krijgt men een nauwkeurigere maatstaf voor de zekerheid.
• Prefix Matching: Om te zorgen dat de clusters ook rekening houden met de context en tokenisatie (bijv. als "television" wordt opgesplitst in "tele" en "vision"), wordt gecontroleerd of een kandidaat-token een prefix is van de gegenereerde tekst. Tokens die als prefix matchen, worden ook meegenomen in de aggregatie.
• Berekening Onzekerheid: De uiteindelijke onzekerheidsscore voor een volledige sequentie is $1 - \prod \hat{p}_c$, waarbij $\hat{p}_c$ de geaggregeerde kans is over de cluster voor elk decodestapje. Een lage geaggregeerde kans resulteert in een hoge onzekerheidsscore.

---

Belangrijkste Bijdragen

1. Efficiëntie: STC vereist slechts één generatie (single sample) en geen herhaaldelijk sampling, wat de inferentietijd drastisch verlaagt.
2. Zelfstandigheid: De methode maakt geen gebruik van externe modellen (zoals NLI-modellen) of extra trainingsdata. Het werkt puur op basis van de interne logits en embeddings van het bestaande LLM.
3. Semantische Awareness: In tegenstelling tot traditionele methoden die alleen kijken naar token-kansen (zoals Perplexity), integreert STC semantische informatie door tokens te groeperen op basis van betekenis.
4. Implementatiegemak: De methode vereist geen fine-tuning en kan direct worden toegepast op elke "white-box" LLM waar toegang is tot de token-embeddings.

---

Resultaten

De auteurs hebben STC getest op meerdere open-source modellen (Llama-2/3, Mistral, Qwen2.5) en datasets (TriviaQA, Natural Questions, WebQuestions).

• Prestatie (AUROC): STC presteert vergelijkbaar met of beter dan state-of-the-art baselines (zoals Semantic Entropy, EigenScore en CCP). Het behaalt consistent hoge AUROC-scores (bijv. ~0.90 op TQA voor Qwen2.5-7B), wat aangeeft dat het onzekerheidsscores zeer goed kan onderscheiden tussen juiste en onjuiste antwoorden.
• Efficiëntie:
  • In vergelijking met sampling-based methoden (zoals Semantic Entropy) die honderden procenten extra rekentijd vereisen, is STC extreem snel.
  • De relatieve uitvoeringstijd (overhead) is gemiddeld < 3% ten opzichte van basis-inferentie.
  • In vergelijking met CCP (een andere single-sample methode) reduceert STC de inferentie-tijd overhead met gemiddeld 98%, terwijl de prestaties vergelijkbaar blijven.
• Ablatie Studies: Het verwijderen van zowel embedding clustering als prefix matching leidt tot een aanzienlijke daling in prestatie, wat aantoont dat beide componenten cruciaal zijn voor het succes van de methode.

---

Betekenis en Toekomstperspectief

Deze paper introduceert een paradigmaverschuiving in onzekerheidskwantificatie voor LLMs. Door de interne semantische structuur van het model direct te benutten in plaats van externe heuristieken of dure sampling, biedt STC een praktische oplossing voor real-time toepassingen waar snelheid en betrouwbaarheid essentieel zijn (bijv. in gezondheidszorg, juridische domeinen of zoekmachines).

Beperkingen:

• De methode vereist toegang tot interne token-logits en embeddings, wat betekent dat deze niet direct toepasbaar is op gesloten API's (zoals GPT-4) zonder deze interne details.
• De methode gebruikt statische embeddings, wat kan leiden tot ruis bij polyseme woorden (woorden met meerdere betekenissen) die niet in context worden geanalyseerd. De auteurs suggereren dat toekomstig werk kan focussen op context-gevoelige embeddings om dit te verbeteren.

Concluderend biedt STC een krachtige balans tussen hoogwaardige prestaties en uiterst lage rekentijd, waardoor het een ideale kandidaat is voor het schalen van betrouwbare LLM-toepassingen.