Improving LLM Predictions via Inter-Layer Structural Encoders

Titel: Verbetering van LLM-predicties via Inter-Layer Structurele Encoders (ILSE)

1. Het Probleem

De standaardpraktijk bij Large Language Models (LLMs) is om voorspellingen te baseren op de token-representaties van de laatste laag (final-layer). Dit gaat uit van de aanname dat diepere lagen de meest taakrelevante informatie bevatten. Echter, recent onderzoek toont aan dat:

• Intermediaire lagen aanzienlijke informatie bevatten die vaak relevanter is dan die van de laatste laag.
• De optimale laag voor een specifieke taak varieert en niet universeel is.
• Bestaande methoden om lagen te combineren beperkingen hebben:
  • ELMo (lineaire weging): Kan geen complexe, niet-lineaire interacties tussen lagen vastleggen.
  • Depth-Wise Attention (DWAtt): Richt zich te veel op de laatste laag als "query" en behandelt de andere lagen niet gelijkwaardig.
  • Architecturale wijzigingen: Methoden die de transformer-architectuur zelf aanpassen, vereisen vaak training van nul af, wat onhaalbaar is voor grote, voorgeöefende modellen.

Er is behoefte aan een efficiënte methode om representaties van alle lagen van een bevroren (frozen) LLM te combineren zonder de onderliggende architectuur te wijzigen.

2. Methodologie: Inter-Layer Structural Encoders (ILSE)

De auteurs introduceren ILSE, een wiskundig onderbouwde aanpak die een effectieve representatie leert door de interne lagen van een LLM te combineren.

De Werkstroom:

1. Lagere Representaties: Voor elke laag $\ell$ van de LLM wordt een enkele representatie $z_\ell$ gegenereerd door middel-pooling over alle token-representaties in die laag.
2. Structuurdefinitie: ILSE introduceert een grafstructuur over deze laag-representaties en gebruikt DeepSets of Graph Neural Networks (GNNs) om te leren hoe deze te combineren.
3. De Encoders: Er worden drie varianten geïntroduceerd:
   • Set-Encoder: Behandelt lagen als een verzameling (set) zonder randen (lege graf) en gebruikt DeepSets voor permutatie-invariante pooling.
   • Fully-Connected (FC) Encoder: Construeert een volledig verbonden graf waar elke laag met elke andere laag verbonden is, gevolgd door een GNN.
   • Cayley-Encoder (De kerninnovatie): Dit is een wiskundig onderbouwde geometrische encoder.
     • Het mapt lagen af op knopen in een Cayley-graf gebaseerd op de speciale lineaire groep $SL(2, \mathbb{Z}_n)$.
     • Cayley-graffen zijn reguliere grafen (alle knopen hebben hetzelfde aantal verbindingen) met een logaritmische diameter.
     • Dit zorgt voor een bottleneck-vrije informatieflow: informatie kan efficiënt tussen alle knopen (lagen) worden uitgewisseld zonder dat de expressiviteit van de GNN afneemt (het "over-squashing"-probleem wordt vermeden).
     • De auteurs gebruiken de methode van Wilson et al. (2025) waarbij virtuele knopen worden toegevoegd om de volledige Cayley-structuur te behouden zonder truncatie.

Training:
De basis-LLM blijft volledig bevroren. Alleen de ILSE-encoder en een eenvoudige classifier (voor classificatie) of regressie-head (voor semantische gelijkenis) worden getraind.

3. Belangrijkste Bijdragen

1. ILSE Framework: Een rigoureuze methode om inter-lagenrepresentaties te combineren voor elke LLM, zonder de architectuur aan te passen.
2. Uitgebreide Evaluatie: Testen op 13 taken (classificatie en semantische tekstgelijkenis/STS) met 9 verschillende LLMs (van 14 miljoen tot 8 miljard parameters).
3. Superieure Prestaties: ILSE overtreft consistent bestaande baselines (zoals Last-Layer, Best-Layer, DWAtt) in bijna alle scenario's.
4. Data-efficiëntie: ILSE presteert uitstekend in few-shot scenario's (weinig trainingsdata).
5. Grootte-onafhankelijkheid: ILSE maakt kleine modellen (bijv. 14M parameters) concurrerend met veel grotere modellen (bijv. 2.8B parameters).

4. Resultaten

De experimenten tonen overtuigende verbeteringen:

• Algemene Prestaties:
  • ILSE behaalde de beste scores in alle 15 classificatie-configuraties.
  • Bij semantische gelijkenis (STS) behaalde ILSE de beste scores in 17 van de 21 zero-shot en supervised combinaties.
  • Verbetering: Tot 44% verbetering in nauwkeurigheid ten opzichte van de "Last-Layer" baseline en 25% in gelijkenismetrics.
• Vergelijking met Baselines:
  • ILSE overtreft DWAtt op alle 13 taken, terwijl het 3-5 keer minder parameters toevoegt.
  • De Cayley-Encoder (met GIN of GCN) presteerde over het algemeen het beste, vooral bij STS-taken, wat suggereert dat de bottleneck-vrije, reguliere structuur ideaal is voor semantische aggregatie.
• Few-Shot Learning:
  • Met slechts 32 voorbeelden per label presteerde ILSE beter dan baselines die getraind waren op volledige datasets (ca. 10.000 voorbeelden).
• Schaalbaarheid:
  • Een Pythia-14M model met ILSE bereikte vergelijkbare prestaties als een Pythia-2.8B model zonder ILSE (bijv. 95-96% vs 98-99% nauwkeurigheid op MTOP Domain).

5. Betekenis en Conclusie

Dit paper demonstreert dat de interne diepte van LLMs waardevolle, complementaire signalen bevat die vaak worden genegeerd door alleen de laatste laag te gebruiken.

• Wiskundige Onderbouwing: De toepassing van Cayley-graffen voor inter-lagencommunicatie biedt een wiskundig solide oplossing voor het "over-squashing"-probleem in GNNs, wat zorgt voor efficiënte informatieflow.
• Efficiëntie: De methode voegt slechts een zeer klein aantal trainbare parameters toe (minder dan 0,1% van de basisgrootte), wat het zeer schaalbaar maakt.
• Praktische Toepassing: ILSE is een krachtige tool voor het verbeteren van bestaande, bevroren LLMs, vooral in scenario's met beperkte data (few-shot) of wanneer het gebruik van zeer grote modellen te duur is (door kleine modellen te "upgraden").

De conclusie is dat gestructureerde aggregatie van alle lagen, met name via geometrisch onderbouwde grafen zoals Cayley-graffen, een effectieve inductieve bias biedt om de prestaties van Large Language Models aanzienlijk te verbeteren zonder de onderliggende modelarchitectuur te wijzigen.