Technische Samenvatting: SPOKES – Optimaliseren voor diverse pretraining-data selectie

1. Probleemstelling

Data-diversiteit is een cruciale factor bij het construeren van pretraining-corpora, omdat het redundantie vermindert en de prestaties bij downstream-taken verbetert onder een vastgesteld databudget. Het optimaliseren van diversiteit is echter inherent moeilijk omdat het een set-niveau eigenschap is die afhankelijk is van interacties tussen datapunten in plaats van individuele voorbeelden. Directe optimalisatie van diversiteit is computationeel onuitvoerbaar vanwege de combinatorische aard van subsetselectie, met name op de schaal van pretraining.

Bestaande benaderingen vertrouwen doorgaans op proxies of benaderingen, zoals grofmazige clustering (op basis van onderwerpen, vaardigheden of ongesuperviseerde clusters) of semantische deduplicatie. Hoewel deze methoden representatie van clusters waarborgen, vertrouwen ze op laag-resolutie partities en falen ze vaak om informatie in de 'long-tail' of ondervertegenwoordigde kennis te vangen. Bov�iment zijn huidige methoden vaak gericht op het behandelen van diversiteitsmetrieken (zoals de G-Vendi score) als post-hoc evaluatietools in plaats van als optimalisatie-objectieven. De centrale uitdaging die door dit werk wordt aangepakt, is hoe men betrouwbaar diverse subsets kan extraheren uit grootschalige pretraining-corpora om downstream-resultaten te verbeteren zonder te vertrouwen op heuristieken die niet schalen of fijnmazige diversiteit missen.

2. Methodologie: SPOKES

De auteurs introduceren SPOKES (Spokes: Optimizing for Diverse Pretraining Data Selection), een schaalbaar framework dat direct optimaliseert voor diversiteit met behulp van de G-Vendi score binnen een probabilistisch diversificatie-framework.

2.1 Kernobjectief

De methode formuleert data-selectie als een globaal optimalisatieprobleem dat gezamenlijk kwaliteit en diversiteit maximaliseert. Gegeven een dataset van $n$ voorbeelden, is het doel om een subset $S$ van grootte $k$ te selecteren die het volgende maximaliseert:
$$\max_{S \subseteq [n], |S|=k} \alpha \ln \left( \frac{1}{k} \sum_{i \in S} q_i \right) + (1 - \alpha) \ln \text{Vendi}(K_S)$$
Waarbij:

• $q_i$ de kwaliteitsscore van voorbeeld $i$ vertegenwoordigt (bijv. van een FineWeb-Edu classifier).
• $K_S$ de gelijkenis-kernel is die beperkt is tot de geselecteerde subset.
• $\alpha \in [0, 1]$ een instelbare parameter is die de afweging tussen kwaliteit en diversiteit controleert.

2.2 Gradiënt-gebaseerde Diversiteit (G-Vendi)

In tegenstelling tot traditionele methoden die statische embeddings gebruiken, opereert SPOKES in de gradiëntruimte. Het berekent gradiënten $\nabla_\theta \ell(x_i)$ met betrekking tot modelparameters $\theta$ voor elk datapunt met behulp van een proxy-model. De G-Vendi score meet diversiteit door de entropie van het spectrum van de gelijkenis-matrix te onderzoeken die is geconstrueerd uit deze gradiënten. Deze aanpak stimuleert orthogonaliteit tussen representatievectoren, wat redundantie in update-richtingen vermindert en meer onafhankelijke, informatieve optimalisatiestappen bevordert.

2.3 Optimalisatie via Exponentiated Gradient Descent

Om het combinatorische subsetselectieprobleem op te lossen, gebruikt SPOKES een relatieve optimalisatiestrategie:

1. Continue Relaxatie: In plaats van discrete selectie, wordt een continue niet-negatieve gewichtsvector $w$ toegewezen aan elk datapunt, beperkt tot de simplex ($w \in \Delta_n$).
2. Kernel Constructie: Een gewichtsafhankelijke kernel $K(w)$ wordt gedefinieerd waarbij de entries worden geschaald door $\sqrt{w_i w_j}$.
3. Optimalisatie: Het objectief wordt geoptimaliseerd met behulp van exponentiated gradient descent. In elke iteratie worden gewichten bijgewerkt via $w_i^{(t+1)} \propto w_i^{(t)} \exp(-\eta g_i^{(t)})$ en genormaliseerd om instorting op enkele punten met een hoog gewicht te voorkomen.
4. Discrete Recovering: Na convergentie wordt de uiteindelijke subset teruggewonnen door de top-$k$ entries van de geleerde gewichtsvector te selecteren.

2.4 Schaalbaarheidsverbeteringen

Om deze aanpak haalbaar te maken op de schaal van pretraining, implementeren de auteurs verschillende efficiëntiestrategieën:

• Afgekante Gradiënten (Truncated Gradients): In plaats van volledige modelgradiënten te berekenen, benadert de methode deze door alleen de laatste $n$ transformer-lagen (meestal 2 lagen) te gebruiken. Dit vermindert de computationele kosten aanzienlijk terwijl de structuur van de gelijkenis-kernel behouden blijft.
• Johnson-Lindenstrauss (JL) Projectie: Om hoog-dimensionale gradiëntvectoren (bijv. 200M parameters) te verwerken, passen de auteurs willekeurige projecties toe (Rademacher-matrices) om de dimensionaliteit te reduceren. Empirische resultaten tonen aan dat zelfs laag-dimensionale projecties (bijv. $k=1024$) een hoge subset-stabiliteit behouden (95,5% overlap met hoog-dimensionale referenties).
• Parallellisatie: De dataset kan worden opgedeeld in disjuncte subsets die onafhankelijk worden geoptimaliseerd, wat een lineaire versnelling oplevert.

3. Belangrijkste Bijdragen

1. SPOKES Framework: Introductie van een schaalbare, principiële techniek voor diversiteitsoptimalisatie die efficiënt werkt op de schaal van moderne pretraining-datasets, waarmee verder wordt gegaan dan heuristische clustering.
2. Directe G-Vendi Optimalisatie: Demonstratie dat de G-Vendi score direct geoptimaliseerd kan worden via exponentiated gradient descent om zeer diverse subsets te extraheren, in plaats van enkel te dienen als een post-hoc metriek.
3. Gezamenlijke Kwaliteit-Diversiteit Optimalisatie: Bewijs dat het gezamenlijk optimaliseren voor kwaliteit en diversiteit superieure resultaten oplevert vergeleken met het optimaliseren voor slechts één objectief of het gebruik van standaard baselines.
4. Efficiëntieverbeteringen: Validatie dat afgekante gradiëntberekening en dimensionaliteitsreductie via JL-projectie de methode in staat stellen om te schalen naar massieve corpora zonder buitensporige computationele kosten.

4. Experimentele Resultaten

De methode werd geëvalueerd op twee grootschalige pretraining-corpora: DCLM (Dolmino subset, agressief gefilterd) en FineWeb (minder agressief gefilterd). Modellen werden vanaf nul getraind op een LLaMA-1B architectuur met 100 miljard tokens.

4.1 Diversiteitsmetrieken

• G-Vendi Score Toename: Op een subset van 500k samples behaalde SPOKES een toename van +489 punten in de G-Vendi score vergeleken met random sampling. In contrast hiermee behaalde semantische deduplicatie (SemDeDup) slechts een toename van +7 punten.
• Batch Diversiteit: Tijdens de training behielden door SPOKES geselecteerde batches consequent hogere gemiddelde diversiteitsscores (stijgend van 315 naar 520 op DCLM) vergeleken met random sampling.

4.2 Downstream Prestaties

Evaluatie werd uitgevoerd op 10 Engelse benchmarks uit de OLMES-suite.

• Alleen Diversiteit: SPOKES (alleen diversiteit) verbeterde de gemiddelde downstream prestaties met +0,4 (DCLM) en +0,5 (FineWeb) punten ten opzichte van random sampling.
• Gezamenlijke Optimalisatie: De sterkste resultaten werden behaald door kwaliteit en diversiteit gezamenlijk te optimaliseren. SPOKES behaalde winsten van +1,5 (DCLM) en +1,4 (FineWeb) punten ten opzichte van random sampling, waarmee het alle baselines versloeg, inclusief semantische deduplicatie en kwaliteit-alleen filtering.
• Complementaire Sterktes: De resultaten wijzen erop dat kwaliteit en diversiteit verschillende maar complementaire eigenschappen vatten. Terwijl kwaliteit-filtering uitblinkt in kennisintensieve benchmarks (bijv. MMLU), verbetert diversiteitsoptimalisatie het commonsense reasoning (bijv. HellaSwag). Gezamenlijke optimalisatie overbrugt deze kloof en presteert competitief over alle taken.

5. Betekenis en Claims

Het artikel claimt dat directe optimalisatie van diversiteit zowel hanteerbaar op schaal als praktisch voordelig is. Door weg te bewegen van grofmazige benaderingen en heuristieken, demonstreert SPOKES dat:

• Diversiteit een kritieke as is: Zelfs zonder expliciete kwaliteits-supervisie levert het optimaliseren voor diversiteit alleen consistente winsten op, wat suggereert dat een diverse gradiënt-representatie van de datadistributie de generalisatie verbetert.
• Kwaliteitsmetrieken onvoldoende zijn: Standaard kwaliteitsmetrieken vatten de bruikbaarheid van data niet volledig. Het gezamenlijk optimaliseren voor diversiteit kan complementaire trainingssignalen bieden die leiden tot betere downstream-prestaties, zelfs wanneer de geselecteerde subset een lagere gemiddelde kwaliteit heeft dan een kwaliteit-alleen baseline.
• Schaalbaarheid: De voorgestelde methode extraheert succesvol dichte, zeer diverse subsets uit massieve pretraining-corpora, wat een levensvatbaar alternatief biedt voor bestaande selectie-pipelines die vertrouwen op semantische deduplicatie of random sampling.

De auteurs positioneren SPOKES als een instrument voor het construeren van hoogwaardige en diverse subsets, waarbij zij benadrukken dat de voordelen van diversiteit betrouwbaar vertalen naar datasets met verschillende filtering-regimes en ruisprofielen.