Technische Samenvatting: Laplace-hoofden verbeteren Transformers door tokenrepresentaties te gladstrijken

Probleemstelling

De standaard Transformer-architectuur werkt tokenrepresentaties bij via multi-head attention en residuverbindingen. In deze formulering berekent elk attention-head $i$ een softmax-attentionmatrix $P^{(i)}$ en werkt de input $X$ bij via een gewogen som van waarde-projecties: $X \leftarrow X + \sum_i P^{(i)}XW_{Vi}W_{oi}$.

Hoewel dit mechanisme effectief is, ontbreekt er expliciete controle over de variantie binnen de sequentie van tokenrepresentaties. Het heersende standpunt in de literatuur, met name met betrekking tot Graph Neural Networks (GNN's) en diepe Transformers, behandelt "oversmoothing" – waarbij tokenrepresentaties te veel op elkaar gaan lijken en hun onderscheidend vermogen verliezen – als een schadelijke faalmodus. Bijgevolg richt het merendeel van het eerdere werk zich op het voorkomen van gladstrijking door normalisatie, het weglaten van randen of diversiteitsverliezen. Dit artikel daagt de aanname uit dat token-oversmoothing inherent schadelijk is, en stelt in plaats daarvan dat gecontroleerde gladstrijking gunstig kan zijn voor representatieleren.

Methodologie

De auteurs stellen een parameterloze architecturale modificatie voor: het vervangen van een subset van standaard attention-heads door Laplace-hoofden.

De Laplace-operator

In de standaard update berekent een hoofd $P^{(i)}X$, wat een door attention gewogen gemiddelde van tokens is. Het voorgestelde Laplace-hoofd vervangt de attentionmatrix $P^{(i)}$ door de graf-Laplace-operator $I - P^{(i)}$, waarbij $I$ de eenheidsmatrix is. De updatevergelijking wordt:
$$X \leftarrow X + \sum_{i \in A} P^{(i)}XW_{Vi}W_{oi} + \sum_{i \in L} (I - P^{(i)})XW_{Vi}W_{oi}$$
waarbij $A$ de verzameling standaard attention-heads is en $L$ de verzameling Laplace-heads.

Motiveringen

Het artikel biedt twee theoretische motiveringen voor deze modificatie:

1. Directe controle van gemiddelde en variantie: Standaard attention-heads werken het gemiddelde van tokenrepresentaties bij. Daarentegen berekent de term $(I - P^{(i)})X$ het verschil tussen elk token en zijn door attention gewogen gemiddelde. Het toevoegen van dit verschil stelt het model in staat om direct de variantie binnen de sequentie van tokenrepresentaties te controleren, een vermogen dat ontbreekt in standaard Transformers.
2. Graf-diffusie: Als tokens worden beschouwd als knopen in een graf met randgewichten gedefinieerd door $P^{(i)}$, dan is $I - P^{(i)}$ de random-walk genormaliseerde graf-Laplace. Het toepassen van deze operator komt overeen met een enkele stap van warmtediffusie op de token-graf, waarbij het signaal (tokenrepresentaties) progressief wordt gladgestreken over de sequentie.

Belangrijkste bijdragen

1. Architecturaal voorstel: Een eenvoudige, parameterloze modificatie van de Transformer-architectuur die Laplace-heads integreert naast standaard attention-heads.
2. Empirische validatie: Aantonen dat deze modificatie de prestaties verbetert over drie verschillende leerparadigma's: supervisie classificatie, taalmodelleren en zelftoezichthoudend leren.
3. Geometrische karakterisering: Een gedetailleerde analyse van hoe Laplace-heads de geometrie van geleerde tokenrepresentaties veranderen, waarbij wordt aangetoond dat ze gunstige vormen van "instorting" en gladstrijking induceren in plaats van schadelijke oversmoothing.

Experimentele resultaten

1. Beeldclassificatie (supervisie leren)

Met gebruik van de DeiT-3-familie (ViT-B) op CIFAR-10, CIFAR-100 en ImageNet-1k:

• Prestaties: Het toevoegen van Laplace-heads verbeterde consistent de Top-1 nauwkeurigheid. De verbetering was bijna monotoon met het aantal Laplace-heads.
  • Op ImageNet behaalde het model met 12 Laplace-heads (ViT-B-12L) een nauwkeurigheid van 82,32% vergeleken met de baseline van 81,45%.
  • Op CIFAR-100 was de winst bijna 2% (77,04% versus 75,12%).
• Geometrie:
  • ANOVA: Laplace-heads verlaagden de variantie binnen de sequentie significant, terwijl ze de variantie tussen klassen verhoogden. Dit geeft aan dat tokens binnen dezelfde sequentie instorten naar een gemeenschappelijk gemiddelde, terwijl de gemiddelden van verschillende sequenties effectiever van elkaar scheiden.
  • Neural Collapse: De modellen vertoonden een sterkere convergentie naar de Neural Collapse-geometrie. Tokenrepresentaties vormden goed gescheiden, simplex-achtige clusters (Simplex ETF) voor elke klasse, waarbij sequentiegemiddelden uitkwamen op de globale klassengeometrie.

2. Taalmodelleren (autoregressief)

Met gebruik van GPT-2-achtige modellen (561M parameters) die vooraf zijn getraind op FineWeb-Edu en vervolgens zijn fine-gefine-tuned op redenerings/code-benchmarks:

• Prestaties: Modellen met Laplace-heads presteerden gemiddeld beter dan de baseline op ARC-Easy, ARC-Challenge, MMLU, GSM8K en HumanEval.
  • Het model met 5 Laplace-heads behaalde de hoogste gemiddelde prestatie (30,06%), met significante winsten op ARC-Easy (+16%) en ARC-Challenge (+7%).
• Geometrie:
  • Spectrale afname: Laplace-heads produceerden tokenrepresentaties met sneller afnemende spectra. Het aantal hoofdcomponenten dat nodig was om 99% van de variantie te verklaren, daalde met 21% (van 547,7 naar 430,4) toen alle heads Laplace-heads waren.
  • Linguïstische instorting: De modellen vertoonden een toegenomen scheidbaarheid tussen tokenrepresentaties die dezelfde volgende-tokenvoorspelling delen, wat wijst op een sterkere vorm van "linguïstische instorting" waarbij tokens met vergelijkbare semantische rollen samenkomen.

3. Zelftoezichthoudend leren (DINO)

Met gebruik van ViT-S-modellen op ImageNet-1k voor DINO vooraftraining:

• Prestaties: Laplace-heads verbeterden de nauwkeurigheid van lineaire probing op ImageNet en de segmentatieprestaties op ADE20K.
  • ViT-S-5L behaalde de beste ADE20K mIoU (32,27%) vergeleken met de baseline (29,81%).
• Geometrie:
  • Segmentatie: PCA-visualisaties toonden aan dat Laplace-heads scherpere objectgrenzen produceerden en fijnere details behielden (bijv. helmvormen, shirtnummers) in vergelijking met de baseline.
  • Spectrale afname: Net als bij taalmodelleren vertoonden deze modellen een snellere spectrale afname, waarbij tokenrepresentaties binnen een beeld werden geconcentreerd in een lager dimensionale deelruimte.

Betekenis en claims

Het artikel betoogt dat het heersende standpunt – dat token-oversmoothing universeel schadelijk is – onvolledig is. In plaats daarvan claimen de auteurs dat gladstrijking een mechanisme is voor representatieleren dat gunstig kan zijn wanneer het gecontroleerd wordt.

• Opdrachtspecifieke voordelen: De "gladstrijking" die door Laplace-heads wordt veroorzaakt, is geen faalmodus maar een functie die aansluit bij opdrachtdoelstellingen:
  • Bij supervisie classificatie faciliteert het Neural Collapse door tokens binnen sequenties te laten instorten en sequentiegemiddelden uit te lijnen met de klassengeometrie.
  • Bij taalmodelleren verbetert het Linguïstische Instorting, waardoor de scheidbaarheid van tokens met identieke volgende-tokendoelen toeneemt.
  • Bij zelftoezichthoudend leren concentreert het representaties in lager dimensionale deelruimten, wat downstream-taken zoals segmentatie ondersteunt.
• Mechanisme: Door expliciet graf-diffusie (via de Laplace) op te nemen in de architectuur, krijgt het model de mogelijkheid om direct variantie binnen de sequentie te controleren, wat leidt tot representaties die robuuster en geometrisch gestructureerd zijn.

De auteurs concluderen dat hoewel hun implementatie met een vaste strategie waarschijnlijk suboptimaal is, de resultaten suggereren dat gladstrijking verder moet worden onderzocht als een fundamenteel onderdeel van representatieleren, in plaats van uitsluitend als een probleem dat moet worden verzacht.