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

1. Il Problema

Il pre-training continuo su dati specifici per un compito (ad esempio, ragionamento matematico o generazione di codice) è un metodo efficace per adattare i Large Language Models (LLM) a nuovi campi. Tuttavia, un approccio diretto che utilizza solo dati del campo target comporta il rischio di dimenticanza catastrofica delle capacità originali del modello (campo sorgente), a causa del significativo spostamento della distribuzione dei dati.

La soluzione comune consiste nel creare un mix di dati tra il campo sorgente e quello target per bilanciare le prestazioni. Le strategie esistenti per il ripesaggio dei domini (domain reweighting) si basano su:

• Euristiche manuali definite da esperti umani.
• Risultati empirici o intuizioni soggettive.
• Algoritmi statici o dinamici che non apprendono strategie generalizzabili.

Questi metodi spesso non riescono a scoprire e sfruttare euristiche più generali e complesse che potrebbero migliorare le prestazioni in modo più efficiente.

2. Metodologia: Data Mixing Agent

Gli autori propongono il Data Mixing Agent, il primo framework basato su modello, end-to-end, che impara a ripesare i domini attraverso l'apprendimento per rinforzo (RL).

Formulazione come Processo Decisionale di Markov (MDP)

Il problema del ripesaggio dei domini è formalizzato come un MDP $(S, A, f, r, \rho_s, \rho_e)$:

• Stato ($S$): La distribuzione dei dati dai passaggi di ripesaggio precedenti.
• Azione ($A$): La distribuzione di probabilità dei dati da allocare tra i vari domini nel passo corrente.
• Politica ($f$): L'agente (Data Mixing Agent) che decide l'azione.
• Ricompensa ($r$): Il feedback fornito da un ambiente di valutazione leggero che misura le capacità del modello su diversi benchmark.

Pipeline di Addestramento

1. Campionamento delle Traiettorie: Vengono generate casualmente migliaia di traiettorie di mixing dei dati. Per ogni traiettoria, piccoli modelli "proxy" vengono addestrati da zero.
2. Raccolta del Feedback: I checkpoint di questi modelli proxy vengono valutati in un ambiente leggero (basato su benchmark come MMLU per le capacità generali e MATH per il ragionamento matematico) a ogni passo di ripesaggio. Questo genera un set di dati di traiettorie e feedback.
3. Ottimizzazione con RL (CQL): L'agente viene addestrato su queste traiettorie utilizzando l'algoritmo Conservative Q-Learning (CQL).
   • Viene utilizzato un approccio off-policy per evitare la sovrastima dei valori Q per azioni fuori distribuzione.
   • L'agente impara a massimizzare le ricompense (prestazioni bilanciate) evitando azioni che degradano le prestazioni, imparando così euristiche generali.
4. Architettura: L'agente è un modello Transformer leggero (2.1M parametri) che prende in input la storia delle distribuzioni e i feedback ambientali per prevedere la distribuzione ottimale per il passo successivo.

Applicazione in Tempo Reale

Durante il pre-training continuo del modello target, l'agente predice "on-the-fly" la distribuzione dei domini per il passo successivo, adattandosi dinamicamente senza bisogno di riaddestramento.

3. Contributi Chiave

• Primo Agente di Mixing Dati: Introduzione del primo metodo basato su modello per il ripesaggio dei domini nel pre-training continuo, che apprende euristiche generalizzabili invece di affidarsi a regole fisse.
• Generalizzazione Senza Riaddestramento: L'agente dimostra la capacità di generalizzare a:
  • Campi sorgente non visti (diversi dataset di pre-training).
  • Modelli target diversi (diverse architetture o dimensioni).
  • Spazi di dominio diversi (definizioni diverse dei domini).
  • Nuovi campi target (es. applicare un agente addestrato su matematica alla generazione di codice).
• Efficienza dei Dati: Il metodo riesce a ottenere prestazioni superiori riducendo la dipendenza dai dati del campo sorgente, risparmiando token di addestramento.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su compiti di ragionamento matematico e generazione di codice, utilizzando modelli LLaMA da 3B parametri e Pythia.

• Prestazioni Bilanciate: Il Data Mixing Agent (DataAgentRL) supera significativamente le baseline (RegMix, DBL, Naive Training).
  • Nel ragionamento matematico, ha ottenuto un miglioramento medio del 3.02% rispetto a RegMix su 12 benchmark (8 generali + 4 matematici).
  • Ha ridotto la varianza delle prestazioni tra i diversi benchmark, garantendo una stabilità superiore.
• Generalizzazione Cross-Domain:
  • Un agente addestrato su dati matematici è stato applicato direttamente alla generazione di codice, ottenendo prestazioni migliori rispetto alle baseline ottimizzate specificamente per il codice (miglioramento del 1.45% su RegMix).
  • L'agente ha funzionato bene su modelli non visti (LLaMA-FWE, LLaMA-Nemo) e su spazi di dominio a 52 dimensioni senza riaddestramento.
• Efficienza Computazionale:
  • Nonostante l'overhead dell'agente, il metodo richiede meno ore GPU rispetto a metodi come RegMix (che richiede l'addestramento di centinaia di modelli proxy per il fitting della regressione).
  • L'agente raggiunge prestazioni superiori utilizzando meno dati del campo sorgente (fino a 2.14B token in meno nel budget totale).
• Allineamento con l'Intuizione Umana: L'analisi delle traiettorie mostra che l'agente impara strategie coerenti con le intuizioni umane (es. aumentare i domini "Science" e "Health" per MMLU, ridurre "Fashion") e scopre nuove euristiche (es. ridurre i dati su "Pets&Animals").

5. Significato e Impatto

Il lavoro rappresenta un passo avanti fondamentale nella gestione dei dati per l'adattamento dei LLM:

1. Automazione delle Strategie: Sostituisce la progettazione manuale e l'euristica statica con un agente che apprende dinamicamente le migliori strategie di mixing.
2. Riduzione dei Costi: Dimostra che è possibile ottenere prestazioni di stato dell'arte con meno dati e meno tempo di calcolo, rendendo il pre-training continuo più accessibile.
3. Robustezza: La capacità di generalizzare a nuovi domini e modelli senza riaddestramento rende il metodo scalabile per futuri scenari di adattamento dei modelli, riducendo la necessità di costose esplorazioni di iperparametri per ogni nuovo compito.

In sintesi, il Data Mixing Agent trasforma il processo di selezione dei dati da un compito di ingegneria manuale a un problema di apprendimento automatico, offrendo una soluzione più robusta, efficiente e generalizzabile al problema della dimenticanza catastrofica nel pre-training continuo.