Beyond Perplexity: A Geometric and Spectral Study of Low-Rank Pre-Training

Technische Samenvatting: Voorbij Perplexiteit: Een Geometrische en Spectrale Studie van Low-Rank Pre-training

Probleemstelling
Het pre-trainen van Large Language Models (LLM's) wordt beperkt door de memorykosten die gepaard gaan met het opslaan van full-rank weights, gradients en optimizer states. Low-rank pre-training is naar voren gekomen als oplossing om deze kosten te verlagen, waarbij methoden zijn geëvolueerd tot drie categorieën: optimizer-gebaseerd (bijv. GaLore, Fira), architectuur-gebaseerd (bijv. CoLA, SLTrain) en adapter-gebaseerd (bijv. ReLoRA). De evaluatie van deze methoden blijft echter beperkt. Huidige praktijken vertrouwen bijna uitsluitend op validatie-perplexiteit uit runs met één seed. De auteurs betogen dat perplexiteit een slechte proxy is voor de kwaliteit van de oplossing; modellen met identieke perplexiteit kunnen convergeren naar fundamenteel verschillende regio's van het verlieslandschap met onderscheiden interne representaties. Bovendien falen bestaande vergelijkingen vaak om onderscheid te maken tussen hoe goed een methode een pre-getraind model fine-tunt versus hoe goed een model dat met die methode is gepre-train downstream generaliseert.

Methodologie
De auteurs introduceren een diagnostisch kader om de oplossingen die worden gevonden door vijf low-rank pre-training methoden (GaLore, Fira, CoLA, SLTrain, ReLoRA) te karakteriseren in vergelijking met full-rank training over drie modelgroottes (60M, 130M, 350M). De studie evalueert deze methoden met behulp van 16 metrieken over vier dimensies:

1. 1-D Verlieslandschap (Willekeurige Richtingen): De auteurs perturbëren modelparameters langs 100 willekeurige richtingen om het gemiddelde en de variantie van de validatie-loss te meten. Zij leiden scalair metrieken af voor Sharpness (de snelheid van loss-groei weg van het checkpoint) en Direction Variance (anisotropie van het loss-oppervlak).
2. 1-D Verlieslandschap (Top-K PCA Richtingen): Perturbaties worden toegepast langs de top singuliere richtingen van de weight matrices om de geometrie langs de meest significante assen te analyseren.
3. 1-D Interpolatie: De auteurs meten de loss langs lineaire paden tussen checkpoints. Zij introduceren twee specifieke metrieken:
   • Consecutive Checkpoints Barrier Height (CCBH): Meet de barrière tussen sequentiële checkpoints van dezelfde methode.
   • Inter-Method Barrier Height (IMBH): Een nieuwe metriek die de geometrische scheiding kwantificeert tussen verschillende trainingsmethoden op hetzelfde trainingsstadium. Dit onthult of verschillende methoden convergeren naar dezelfde basin of geometrisch onderscheiden oplossingen.
4. Spectrale en Rank Metrieken: De studie volgt de spectrale structuur van weights en updates, inclusief Effective Rank, Stable Rank, Spectral Gap en Threshold Rank.
5. Activatie Similariteit: De auteurs vergelijken hidden states van low-rank modellen met full-rank baselines met behulp van Activation L2 distance, Cosine Similarity en Linear CKA (Centered Kernel Alignment).

Tot slot testen de auteurs of deze geometrische en spectrale kenmerken downstream zero-shot prestaties beter kunnen voorspellen dan perplexiteit alleen, met behulp van lineaire regressie met leave-one-size-out (LOSO) en leave-one-method-out (LOMO) cross-validatie.

Belangrijkste Bijdragen

1. Diagnostisch Kader: Het paper stelt een uitgebreide evaluatiesuite voor, inclusief de nieuwe IMBH-metriek, die de geometrische scheiding tussen trainingsalgoritmen kwantificeert in plaats van alleen tussen seeds of epochs.
2. Empirische Bevindingen over Families: De studie stelt vast dat low-rank methoden niet equivalent zijn aan full-rank training noch aan elkaar. Belangrijke bevindingen zijn:
   • Geometrische Scheiding: Methodes convergeren naar geometrisch onderscheiden basins, waarbij IMBH monotoon groeit naarmate de training vordert.
   • Spectrale Signatuur: GaLore komt de spectrale structuur van full-rank updates nauwkeurig overeen, terwijl CoLA dit niet doet.
   • Divergentie van Representaties: Low-rank activaties divergeren in latere lagen steeds meer van full-rank. CoLA vertoont orthogonale activaties, terwijl Fira alleen in de laatste laag alignering herstelt.
3. Verbeterde Downstream Voorspelling: Validatie-perplexiteit alleen is een onvolledige voorspeller van downstream prestaties. Het aanvullen van perplexiteit met geometrische en spectrale kenmerken verbetert de voorspellingnauwkeurigheid aanzienlijk (LOSO Pearson-correlatie stijgt van 0,873 naar 0,913).
4. Gefedereerde Codebase: De auteurs geven een open-source pipeline vrij die pre-training, metriekberekening en downstream evaluatie integreert om verder onderzoek te faciliteren.

Resultaten

• Geometrie van het Verlieslandschap: Full-rank training vestigt zich in scherpere basins langs willekeurige richtingen, maar in vlakkere basins langs de top-1 PCA-richting. Omgekeerd convergeren low-rank methoden zoals Fira en SLTrain naar vlakkere basins langs willekeurige richtingen, terwijl CoLA en ReLoRA de scherpste basins bereiken. GaLore vertoont een gemengd profiel: matig vlak langs willekeurige richtingen maar scherp langs zijn top-1 PCA-richting.
• Inter-Method Barrières: Alle methoden bezetten geometrisch gescheiden basins. Op de 60M-schaal is de barrière tussen full-rank en elke low-rank methode hoger dan barrières tussen low-rank methoden. Fira convergeert naar de basin die geometrisch het dichtst bij full-rank ligt, terwijl SLTrain het meest geïsoleerd blijft.
• Spectrale Analyse: Full-rank domineert alle rank-metrieken. GaLore volgt full-rank nauwkeurig in effectieve en stabiele rank, maar vertoont een kleinere spectrale gap. Fira en GaLore reproduceren de spectrale verdeling van full-rank weight updates het meest trouw.
• Activatie Alignering: ReLoRA volgt full-rank representaties het nauwst over alle schalen. CoLA vertoont de grootste afwijking, waarbij activaties in latere lagen orthogonaal worden aan de baseline. Fira toont een hoge drift gemiddeld over lagen, maar verrassend lage afwijking in de laatste laag.
• Downstream Prestaties: Validatie-perplexiteit voorspelt de downstream ranking niet consequent. Bijvoorbeeld, Fira bereikt de laagste perplexiteit op alle schalen, terwijl CoLA leidt in meerdere downstream categorieën bij 130M en 350M, ondanks dat het niet in de top twee staat voor perplexiteit.

Betekenis
Het paper concludeert dat low-rank pre-training methoden niet uitwisselbare benaderingen zijn van full-rank training; ze veranderen fundamenteel de bereikte oplossingen. De studie demonstreert dat validatie-perplexiteit een onvolledig beeld van modelkwaliteit biedt. Door het opnemen van geometrische en spectrale metrieken kunnen onderzoekers downstream prestaties beter voorspellen en de onderscheidende inductieve biases begrijpen die door verschillende low-rank strategieën worden geïntroduceerd. Het werk benadrukt dat hoewel sommige methoden (zoals Fira) convergeren naar basins die geometrisch dicht bij full-rank liggen, anderen (zoals CoLA) unieke valleien bezetten, wat suggereert dat de keuze van de low-rank methode de specifieke regio van het verlieslandschap dicteert waarin het model zich bevindt.