Toward a First-Principles Update Geometry for the Language-Model Head

In het uitgestrekte landschap van de moderne kunstmatige intelligentie komen de meest zichtbare prestaties vaak voort uit systemen die mensachtige tekst genereren. Achter deze gesprekken schuilt een complex machine learning-model, een digitale hersenpan die getraind is op enorme hoeveelheden data. Aan het uiterste einde van de verwerkingsketen van deze machine bevindt zich een specifiek onderdeel dat bekend staat als de taalmodel-kop (language-model head). De taak ervan is misleidend eenvoudig: het neemt een verborgen interne staat — een wiskundige representatie van wat de computer denkt — en zet dit om in een lijst met waarschijnlijkheden voor elk woord in zijn vocabulaire. Als het model de zin "De lucht is" probeert af te maken, berekent dit onderdeel hoe waarschijnlijk het is dat het volgende woord "blauw", "bewolkt" of "vallend" zal zijn, en wijst het een score toe aan elke mogelijkheid. Jarenlang hebben ingenieurs de wiskunde die deze laatste conversie aanstuurt behandeld als een standaardprobleem, waarbij ze dezelfde algemene regels toepassen om de gewichten van het model aan te passen als bij de eerdere lagen van het netwerk.

Echter, een nieuw perspectief suggereußen dat deze laatste stap fundamenteel anders is dan de rest van de machine. Omdat de output een waarschijnlijkheidsverdeling is — een set percentages die samen honderd procent moeten vormen — zouden de regels voor het veranderen ervan de unieke geometrie van waarschijnlijkheden moeten respecteren, en niet alleen de regels voor het veranderen van ruwe getallen. Een recent artikel van Aditya Somasundaram van Columbia University onderzoekt dit idee door de laatste laag en de waarschijnlijkheidsconversie als één enkele, verenigde module te behandelen. De auteur betoogt dat om te begrijpen hoe we dit deel van het systeem kunnen verbeteren, we moeten kijken naar hoe de update de relatieve kansen tussen woorden verandert, in plaats van alleen naar hoeveel de getallen verschuiven. Door een specifieke wiskundige afstand te gebruiken die de verandering in deze relatieve kansen meet, leidt de studie een nieuwe manier af om de omvang van een update te meten. Het resultaat is een fris geometrisch beeld waarbij de "omvang" van een verandering wordt bepaald door de spreiding van de interne representaties van het model voor verschillende woorden, wat leidt tot een voorstel voor een nieuw soort optimizer die deze modellen efficiënter kan laten leren.

De kern van dit onderzoek begint met een vraag over hoe we verandering meten. Wanneer een computer zijn interne instellingen bijwerkt om te leren van een fout, maakt hij een minuscule aanpassing aan een enorme tabel met getallen. In standaard trainingsmethoden meten ingenieurs de omvang van deze aanpassing vaak door te kijken naar de grootste mogelijke verandering die het kan veroorzaken in de ruwe getallen voordat deze worden omgezet in waarschijnlijkheden. Maar de auteur wijst erop dat voor de laatste laag de ruwe getallen niet het eindproduct zijn; de waarschijnlijkheden wel. Het conversieproces, bekend als softmax, heeft een speciale eigenschap: als je hetzelfde bedrag aan elk getal in de lijst toevoegt, veranderen de resulterende waarschijnlijkheden helemaal niet. Dit betekent dat het meten van de ruwe omvang van de update misleidend is, omdat het veranderingen meet die geen effect hebben op de uiteindelijke output. Om dit op te lossen, wendt het artikel zich tot een concept genaamd de projectieve afstand van Hilbert. Dit is een manier om de afstand tussen twee sets waarschijnlijkheden te meten die zich volledig richt op hoe de ratio tussen twee woorden verandert. Het negeert de absolute grootte van de getallen en kijkt alleen naar de relatieve positie van het ene woord ten opzichte van het andere.

Door deze specifieke afstandsmaat toe te passen, ontdekte de onderzoeker een verrassende en heldere relatie tussen de geometrie van de update en het gedrag van het model. De studie laat zien dat de maximale verandering die een update kan veroorzaken in de waarschijnlijkheidsverdeling direct gekoppeld is aan de fysieke spreiding van de rijen in de update-matrix. Stel je de update voor als een collectie vectoren, waarbij elke vector overeenkomt met een specifiek woord in het vocabulaire. De "omvang" van de update, in termen van hoeveel het de waarschijnlijkheden kan doen schudden, wordt bepaald door de afstand tussen de twee vectoren die het verst uit elkaar liggen. Als de vectoren voor verschillende woorden dicht bij elkaar geclusterd zijn, is de update klein en veilig. Als ze ver uit elkaar liggen, is de update groot en kan deze wilde schommelingen in de voorspellingen van het model veroorzaken. Deze bevinding herkadert het probleem van het bijwerken van de taalmodel-kop: in plaats van je zorgen te maken over de algehele magnitude van de getallen, wordt het doel het beheren van de diameter van de wolk van punten die de woorden vertegenwoordigen.

Dit geometrische inzicht leidt tot een nieuw voorstel voor hoe deze updates te construeren, met inspiratie uit een recente methode genaamd Muon die succesvol is geweest in andere delen van neurale netwerken. Muon werkt door de verschillende richtingen van een update in balans te brengen zodat ze allemaal even sterk zijn, wat voorkomt dat het model vast komt te zitten in smalle dalen van het leerlandschap. De auteur suggereert dat een vergelijkbaar principe zou moeten gelden voor de taalmodel-kop, maar dan met een twist. In plaats van de kracht van enkelvoudige richtingen in balans te brengen, is het doel om de afstanden tussen alle paren woord-vectoren zo gelijk mogelijk te maken. De ideale update zou de vectoren zo verspreiden dat elk woord ongeveer dezelfde afstand tot elk ander woord heeft, waardoor een perfect gebalanceerde wolk ontstaat. Dit zou ervoor zorgen dat het model het onderscheid tussen elk paar woorden met gelijke gevoeligheid behandelt.

Het artikel identificeert echter ook een harde fysieke limiet voor dit ideaal. In de taalmodel-kop is het aantal woorden in het vocabulaire vele malen groter dan het aantal dimensies dat beschikbaar is om hen te representeren. Het is wiskundig onmogelijk om een enorm aantal punten in een kleine ruimte te plaatsen zodat ze allemaal exact dezelfde afstand van elkaar hebben. Net zoals je niet honderd punten op een plat vel papier kunt plaatsen zodat ze allemaal even ver uit elkaar liggen, kun je de woord-vectoren niet perfect equidistant maken wanneer het vocabulaire enorm is en de verborgen ruimte klein is. De studie erkent deze beperking en suggereert dat het doel moet zijn om de best mogelijke benadering te vinden. De voorgestelde optimizer zou proberen de kleinste afstand tussen twee woorden te maximaliseren terwijl de grootste afstand binnen een veilige limiet wordt gehouden. Deze aanpak beoogt een update te creëren die zo gelijkmatig verdeeld is als de geometrie toelaat, om er zeker van te zijn dat geen enkel paar woorden wordt genegeerd of als ononderscheidbaar wordt behandeld, terwijl anderen te ver van elkaar worden weggeduwd.

De implicaties van dit werk zijn primair theoretisch en architecturaal, en bieden een nieuwe lens om naar de laatste stap van taalgeneratie te kijken. De auteur beweert niet dat hij een volledig getraind model heeft gebouwd dat bewijst dat deze methode in de praktijk beter werkt, maar biedt eerder een rigoureuze afleiding van hoe de update-geometrie eruit zou moeten zien als men de eerste principes van de waarschijnlijkheid volgt. Het artikel betoogt dat huidige methoden, die deze laatste laag vaak aan standaard optimalisatietechnieken overlaten, een kans missen om de specifieke functie van de module te respecteren. Door de conversie naar waarschijnlijkheden als een integraal onderdeel van het updateproces te behandelen, en door een afstandsmaat te gebruiken die de invariantie van de softmax-functie respecteert, biedt de voorgestelde geometrie een natuurlijkere manier om door het leerlandschap te navigeren. De studie concludeert dat hoewel de perfecte equidistante ordening onmogelijk is voor grote vocabulaire, het streven naar een benaderend equidistante configuratie een helder, principieel pad biedt voor het ontwerpen van toekomstige optimizers die specifiek zijn afgestemd op de unieke eisen van de taalmodel-kop.