Rethinking Data Mixing from the Perspective of Large Language Models

Probleemstelling

Het trainen van Large Language Models (LLM's) is sterk afhankelijk van de strategie voor het mengen van trainingsdata (data mixing). Hoewel empirisch bewijs aantoont dat ongeschikte strategieën de generalisatie van modellen aanzienlijk kunnen verminderen, blijven er fundamentele vragen onbeantwoord:

1. Definitie van een domein: Wat vormt precies een "domein" voor een LLM? Bestaande methoden vertrouwen vaak op menselijke intuïtie (bijv. Wikipedia, boeken, code), maar deze menselijke definities komen mogelijk niet overeen met hoe het model de data intern waarneemt.
2. Dynamische misalignement: Tijdens het trainingsproces verandert de perceptie van het model. Initieel vormen data uit verschillende domeinen duidelijke clusters, maar naarmate het trainen vordert, smelten deze samen tot een isotrope verdeling. Bestaande methoden die op statische, menselijk gedefinieerde domeinen zijn gebaseerd, houden hier geen rekening mee en worden daardoor vertekend.
3. Efficiëntie en schaalbaarheid: Het exhaustief zoeken naar de optimale sampling-strategie is onmogelijk vanwege de hoge kosten van LLM-training. Bestaande oplossingen zijn ofwel gebaseerd op schaalwetten van kleine modellen (die niet altijd vertaalbaar zijn) of vereisen onrealistische aannames en hoge rekenkosten.

Methodologie: DoGraph

De auteurs introduceren DoGraph, een herschikking van data-mixing die gebaseerd is op de dynamiek van gradients in plaats van menselijke labels.

1. Theoretisch Kader: Domeinen via Gradienten
De kern van de methode is de herdefinitie van een "domein" vanuit het perspectief van het leerproces.

• Gradient Geometry: De auteurs bewijzen dat verschillen in data-distributies worden gecodeerd in de geometrie van de gradients. Volgens Stelling 2.2 kan het verschil in verwachte gradients tussen twee distributies ($P_1$ en $P_2$) worden uitgedrukt als een Kernelized Mean Embedding (MMD) in de gradient-ruimte.
• Conclusie: Een domein wordt dus gedefinieerd door de verwachte stroom van gradients, niet door de token-frequenties. Dit stelt het model in staat om domeinen te "leren" die dynamisch evolueren tijdens het trainen.

2. Het DoGraph Framework
DoGraph formuleert het plannen van data als een graf-geconstrueerd optimalisatieprobleem. Het proces verloopt als volgt per epoch:

• Gradient Extractie & Projectie: Per-sample gradients worden geëxtraheerd en geprojecteerd naar een lagere dimensie via een willekeurige projectie (Random Projection), gebaseerd op de Johnson-Lindenstrauss-lemma. Dit behoudt de geometrische relaties terwijl de rekentijd wordt verlaagd.
• Clustering: In deze geprojecteerde ruimte wordt K-means clustering toegepast om de data te partitioneren in $m$ groepen. Deze groepen vertegenwoordigen model-gecentreerde domeinen die dynamisch veranderen naarmate de gradienten evolueren.
• Optimalisatie van Gewichten: Voor elke cluster (domein) wordt de gemiddelde gradient berekend. Vervolgens worden adaptieve gewichten ($w$) voor elk domein bepaald door een optimalisatieprobleem op te lossen:
  $$\min_{w \in \Delta_{m-1}} \mathcal{L}_{opt}\left(\sum_{j=1}^m w_j \bar{g}_j\right)$$
  Waarbij $\mathcal{L}_{opt}$ een doelwit is (bijv. minimalisatie van gradient-variatie of robuustheid) en $\Delta_{m-1}$ het waarschijnlijkheids-simplex is.
• Update: De modelparameters worden bijgewerkt met de gewogen som van de cluster-gemiddelde gradients.

Belangrijkste Bijdragen

1. Theoretische Link: De auteurs leggen een formele connectie tussen domeindistributies en gradient-dynamiek, en valideren empirisch dat domeinrepresentaties tijdens het trainen dynamisch worden gecorrigeerd.
2. DoGraph Framework: Ze presenteren DoGraph, een theoretisch onderbouwde, graf-gebaseerde herschikking die data-scheduling transformeert in een optimalisatieprobleem dat zich aanpast aan de interne perceptie van het model.
3. Empirische Validatie: Uitgebreide experimenten tonen consistente verbeteringen in prestaties en domein-balans, wat de superioriteit van de aanpak bevestigt ten opzichte van bestaande baselines.

Resultaten

De methode werd getest op GPT-2 modellen (varierend van 210M tot 300M parameters) en LLaMA-modellen (1.1B en 3.2B), getraind op datasets zoals SlimPajama en The Pile.

• Benchmark Prestaties: DoGraph presteerde consistent beter dan state-of-the-art baselines (zoals DoReMi, DOGE, RegMix, en Uniform sampling) op negen downstream benchmarks, waaronder redenering (Reasoning), leesbegrip (Reading Comprehension) en taalmodelleren.
  • Op de SlimPajama-dataset (GPT-2 Medium) behaalde DoGraph een gemiddelde nauwkeurigheid van 37.9%, vergeleken met 37.2% voor de beste concurrent (Dynamic Loss-Based) en 35.4% voor Uniform sampling.
  • De grootste winsten werden geboekt op redenerings-taken (bijv. HellaSwag, PiQA, OBQA), wat suggereert dat DoGraph beter in staat is logische en common-sense afhankelijkheden over domeinen heen te vangen.
• Perplexiteit: DoGraph behaalde de laagste validatie-perplexiteit (3.09 op SlimPajama), wat wijst op een betere generalisatie en een evenwichtigere integratie van domeinen. Andere methoden zoals RegMix vertoonden een hogere perplexiteit (4.51).
• Stabiliteit: De methode bleek robuust over verschillende modelgroottes, waarbij de prestaties verbeterden naarmate het model groeide, in tegenstelling tot uniforme of statische weging.
• Efficiëntie: Hoewel DoGraph extra stappen (projectie en clustering) introduceert, is de rekentijd slechts 4.51% hoger dan RegMix, wat binnen het aanvaardbare budget voor grote schaal training valt.

Betekenis en Impact

Dit paper biedt een paradigmaverschuiving in hoe we data voor LLM's selecteren en wegen.

• Van Mens naar Model: Het benadrukt dat domeinen niet statisch door mensen moeten worden gedefinieerd, maar dynamisch door het model zelf moeten worden geïdentificeerd op basis van gradienten.
• Theoretische Fundament: Door de link te leggen tussen gradient-geometrie en data-distributie, biedt het een wiskundig onderbouwd kader voor data-mixing, in plaats van puur empirische heuristieken.
• Toekomstgericht: DoGraph is schaalbaar en vereist geen proxy-modellen of dure validatie-loops, waardoor het direct toepasbaar is op toekomstige, grotere LLM's. Het opent de weg voor meer geautomatiseerde en adaptieve data-strategieën die de interne representaties van AI-modellen beter benutten.