On the Predictive Power of Representation Dispersion in Language Models

Probleemstelling

Grote Taalmodellen (LLMs) presteren opmerkelijk goed op diverse taken, maar hun inbeddingsruimte (embedding space) vertoont vaak tekenen van anisotropie of rank collapse. Dit betekent dat de verborgen staten (hidden states) van het model zich in een smalle kegel bevinden of een lage-dimensionale deelruimte bezetten, in plaats van de volledige ruimte te benutten. Hoewel bekend is dat deze geometrie de expressieve kracht kan beperken, was de precieze link tussen de breedte van deze inbeddingsruimte en de prestaties van het model bij autoregressieve tekstgeneratie (met name perplexiteit) niet volledig gekwantificeerd. De auteurs stellen de vraag: Hoe beïnvloedt de spreiding van representaties de voorspellende kracht van een taalmodel?

Methodologie

De kern van de studie is het kwantificeren van representatie-dispersie (representation dispersion).

1. Definitie van Dispersie:
   De dispersie wordt gedefinieerd als de gemiddelde paarsgewijze cosinus-afstand tussen de verborgen vectoren (hidden vectors) van een bepaald model (meestal de laatste laag) voor een reeks tekstfragmenten.

   • Formule: $D = \frac{1}{\binom{N}{2}} \sum_{1 \le i < j \le N} \left[ 1 - \frac{E_i \cdot E_j}{\|E_i\|\|E_j\|} \right]$
   • Een hogere waarde duidt op een bredere spreiding van representaties in de inbeddingsruimte.

2. Empirische Analyse:

   • De auteurs analyseren diverse model families (LLaMA, Qwen, Gemma, Mistral, Phi) en domeinen (Wikipedia, nieuws, wetenschappelijke abstracts).
   • Ze vergelijken de dispersie met de perplexiteit (PPL) van tekstsegmenten.
   • Ze onderzoeken de invloed van laagdiepte (layer-wise patterns) en fine-tuning (LoRA vs. full-parameter).
   • Ze analyseren de dispersie binnen semantische clusters (teksten met dezelfde 10-gram voortzetting) om te zien of sterke modellen ook semantisch vergelijkbare contexten beter onderscheiden.

3. Praktische Toepassingen (Zonder gelabelde data):

   • Hardheid Ranking: Gebruik van dispersie om voorbeelden te rangschikken op moeilijkheidsgraad.
   • Model Selectie: Een "dispersion gap" metric om modellen binnen een familie te screenen op hun aanpassing aan een specifiek domein (bijv. wiskunde of code).
   • Laag Selectie voor kNN-LM: Het identificeren van de beste sublaag (Attention vs. FFN) voor retrieval-augmented language modeling.
   • Trainingssignaal: Integratie van een "push-away" (afstotende) loss-functie in het trainingproces om de dispersie te vergroten.

Belangrijkste Resultaten

1. Sterke Negatieve Correlatie:
   Er is een sterke, negatieve correlatie tussen representatie-dispersie en perplexiteit. Modellen met een bredere spreiding van hun contextuele representaties bereiken een lagere perplexiteit (betere voorspelling). Dit geldt voor verschillende modelgroottes en domeinen.

   • Observatie: Tekstfragmenten die het model goed voorspelt (lage PPL) hebben bredere, meer gedispergeerde embeddings. Fragmenten die moeilijk te voorspellen zijn (hoge PPL) zijn meer geclusterd/comprimerd.

2. Laag- en Trainingseffecten:

   • De negatieve correlatie wordt sterker in de diepere lagen van het model.
   • Fine-tuning (zowel LoRA als full-parameter) verhoogt de gemiddelde scheiding van embeddings in het doeldomein, wat correleert met betere prestaties.
   • Zelfs binnen semantisch zeer vergelijkbare clusters (dezelfde 10-gram voortzetting) neemt de interne afstand toe tijdens het trainen, wat aangeeft dat sterke modellen subtiele verschillen kunnen onderscheiden.

3. Toepassingen:

   • Label-vrije Diagnostiek: Dispersie kan worden gebruikt om voorbeelden te rangschikken op moeilijkheid zonder ground-truth labels. Voorbeelden met lage dispersie zijn vaak "hard" voor het model.
   • Model Selectie: De "dispersion gap" (een maatstaf voor hoe goed domein-specifieke tokens van elkaar en van algemene tokens gescheiden zijn) correleert sterk met downstream-taakprestaties (bijv. MATH en HumanEval). Dit biedt een goedkope, label-vrije methode om de beste checkpoints te selecteren.
   • kNN-LM Optimalisatie: De output van de Attention-sublaag vertoont consistent hogere dispersie dan de Feed-Forward (FFN) laag en levert betere resultaten op voor kNN-LM retrieval.
   • Training Verbetering: Het toevoegen van een eenvoudige "push-away" loss (die de paarsgewijze afstand maximaliseert) tijdens het trainen leidt direct tot een vermindering van de perplexiteit, zowel in single-domain als cross-domain scenario's.

Bijdragen

1. Fundamenteel Inzicht: Het paper vestigt een causaal verband (onderbouwd door experimenten met "push-away" en "squeeze" loss) tussen de geometrische breedte van de inbeddingsruimte en de voorspellende kwaliteit van taalmodellen.
2. Praktische Tools: Het introduceert een reeks label-vrije methoden voor:
   • Het screenen van modelcheckpoints.
   • Het identificeren van moeilijke data-slices.
   • Het selecteren van de optimale laag voor retrieval-systemen.
3. Trainingssignaal: Het demonstreert dat het expliciet stimuleren van bredere representaties via een auxiliary loss de prestaties van het model kan verbeteren zonder extra gelabelde data.

Significantie

Deze studie biedt een nieuw perspectief op de interne werking van LLMs. In plaats van alleen te kijken naar de outputkwaliteit, toont het aan dat de geometrie van de verborgen staten een robuuste indicator is voor modelcapaciteit.

• Efficiëntie: Het biedt goedkope, berekeningsarme methoden om modellen te evalueren en te selecteren zonder zware inferentie of annotatie.
• Robuustheid: Het suggereert dat "rank collapse" (te sterke compressie) een beperkende factor is en dat het actief bevorderen van dispersie een veelbelovende richting is voor het trainen van robuustere en interpreteerbaardere modellen.
• Toepasbaarheid: De bevindingen zijn breed toepasbaar over verschillende modelarchitecturen en domeinen, wat de generaliseerbaarheid van de methode onderstreept.

Kortom, het paper bewijst dat een taalmodel dat zijn contexten "breder" in de vectorruimte spreidt, beter in staat is om tekst te voorspellen, en biedt concrete methoden om dit principe te benutten voor modelontwikkeling en evaluatie.