Data Mixing Agent: Learning to Re-weight Domains for Continual Pre-training

Probleemstelling

Grote Taalmodellen (LLMs) worden vaak getraind op grote, algemene datasets om generalisatievermogen te bereiken. Wanneer deze modellen moeten worden aangepast aan nieuwe, specifieke domeinen (zoals wiskundig redeneren of codegeneratie) via voortdurende vooropleiding (continual pre-training), ontstaat er een groot risico op catastrofaal vergeten (catastrophic forgetting). Dit betekent dat het model zijn oorspronkelijke vaardigheden verliest terwijl het nieuwe kennis leert.

De gebruikelijke oplossing is het mengen van data uit het brondomein (oorspronkelijke kennis) en het doeldomein (nieuwe kennis). Echter, het bepalen van de optimale verhouding (re-weighting) tussen deze domeinen is complex. Bestaande methoden vertrouwen vaak op:

1. Handmatige heuristieken: Gebaseerd op menselijke intuïtie of empirische resultaten.
2. Statiche of simpele dynamische strategieën: Die niet volledig in staat zijn om complexe patronen in de data-distributie te leren.

Er is een behoefte aan een methode die automatisch en adaptief de data-mix kan bepalen om een evenwicht te vinden tussen het behouden van oude vaardigheden en het leren van nieuwe, zonder handmatige tuning.

Methodologie: Data Mixing Agent

De auteurs introduceren de Data Mixing Agent, het eerste model-based, end-to-end framework dat leert om domeinen te herwegen voor voortdurende vooropleiding. De aanpak wordt als volgt opgebouwd:

1. Formuleren als Markov Beslissingsproces (MDP):
Het probleem van domein-herweging wordt gemodelleerd als een MDP:

• State (S): De geschiedenis van data-distributies van vorige stappen.
• Action (A): De huidige data-distributie (de verhouding van domeinen) voor de volgende trainingsstap.
• Reward (r): Feedback uit een lichtgewicht evaluatie-omgeving die de prestaties van het model meet op zowel algemene als doel-taken.

2. Traject Sampling en Proxy Models:
Om het ruimte van heuristieken te verkennen, worden grote hoeveelheden willekeurige data-mixing trajecten gegenereerd. Voor elk traject wordt een klein proxy-model (bijv. 50M parameters) vanaf nul getraind. Op elke stap van het traject wordt dit proxy-model geëvalueerd in de omgeving om feedback (rewards) te verzamelen.

3. Agent Architectuur:
De agent zelf is een lichtgewicht Transformer-decoder (slechts 2,1 miljoen parameters) die is ontworpen om tijdsreeksen te modelleren.

• Input: De vorige trajecten en de bijbehorende omgevingsfeedback.
• Output: De voorspelde domeindistributie voor de volgende stap.

4. Training met Reinforcement Learning (RL):
De agent wordt geoptimaliseerd met Conservative Q-Learning (CQL), een off-policy RL-algoritme.

• SFT Warm-up: Eerst wordt de agent gesuperviseerd gefinetuned (SFT) op hoogwaardige trajecten om de zoekruimte te verkleinen.
• CQL Optimalisatie: Vervolgens wordt CQL gebruikt om de agent te leren van zowel goede als slechte trajecten. CQL voorkomt dat de agent overoptimaliseert op acties die buiten de verdeling van de trainingsdata liggen (out-of-distribution), wat cruciaal is voor stabiel leren.

5. Toepassing:
Tijdens de daadwerkelijke voortdurende vooropleiding van het doelmodel, voorspelt de agent "on-the-fly" de optimale data-distributie voor de volgende stap, gebaseerd op de huidige staat en eerdere feedback.

Belangrijkste Bijdragen

1. Eerste Model-based Aanpak: De Data Mixing Agent is het eerste systeem dat leert om domeinen te herwegen via een end-to-end model in plaats van handmatige regels.
2. Generalisatie zonder Hertraining: De agent leert algemene heuristieken die generaliseren over ongezette bron-domeinen, doelmodellen en domeinruimtes zonder opnieuw getraind te hoeven worden.
3. Efficiëntie: Het framework reduceert de afhankelijkheid van bron-domein data, wat leidt tot lagere rekentijd en token-kosten.
4. Menselijke Alignering: De geleerde heuristieken blijken sterk overeen te komen met menselijke intuïtie over welke domeinen gunstig of schadelijk zijn voor specifieke taken.

Resultaten

De methode is geëvalueerd op wiskundig redeneren en codegeneratie met verschillende doelmodellen (o.a. LLaMA-varianten en Pythia).

• Prestaties: De Data Mixing Agent (DataAgentRL) overtreft state-of-the-art baselines zoals RegMix (static regression) en DBL (dynamic batch loading) significant.
  • Op wiskundige taken behaalde de agent een gemiddelde verbetering van 3,02% ten opzichte van RegMix over 12 benchmarks.
  • Het bereikte een evenwichtige prestatie: het verbeterde de doel-taken (wiskunde) terwijl het de algemene vaardigheden (MMLU, etc.) beter behield dan naieve training.
• Generalisatie: De agent, getraind op wiskunde, kon direct worden toegepast op codegeneratie zonder hertraining en presteerde nog steeds beter dan gespecialiseerde baselines. Ook generaliseerde het naar ongezette doelmodellen en domeinruimtes (2D vs. 52D).
• Efficiëntie: De agent bereikte superieure prestaties met minder tokens uit het bron-domein, wat resulteerde in een besparing van ongeveer 2,14 miljard tokens in het trainingsbudget vergeleken met RegMix.
• Strategie: De agent ontwikkelde een verfijnde strategie: een vroege fase met veel bron-data, een middenfase met een snelle toename van doel-data, en een late fase waarin bron-data weer wordt geïntroduceerd om stabiliteit te garanderen.

Significantie

Dit werk markeert een verschuiving van handmatige, heuristische data-mixing naar geautomatiseerd, leerbaar data-curatie.

• Het lost het probleem van catastrofaal vergeten op door dynamisch de balans tussen oude en nieuwe kennis te vinden.
• Het demonstreert dat kleine agent-modellen complexe optimalisatieproblemen in de data-distributie kunnen oplossen, wat de weg vrijmaakt voor efficiëntere en robuustere voortdurende vooropleiding van LLMs.
• De methode is schaalbaar en kosteneffectief, aangezien de agent zelf zeer lichtgewicht is en slechts één keer getraind hoeft te worden om op verschillende scenario's toegepast te worden.

Kortom, de Data Mixing Agent biedt een krachtig, generaliserend framework om de prestaties van taalmodellen in nieuwe domeinen te maximaliseren zonder hun bestaande kennis te verliezen.