Advancing General-Purpose Reasoning Models with Modular Gradient Surgery

Riassunto Tecnico: Modular Gradient Surgery (MGS)

1. Il Problema: Interferenza tra Domini nel Reinforcement Learning (RL)

L'obiettivo recente dello sviluppo di Large Language Models (LLM) è la creazione di Large Reasoning Models (LRM) capaci di ragionamento complesso (Chain-of-Thought). Sebbene il Reinforcement Learning (RL) sia estremamente efficace per ottimizzare capacità specifiche (come la matematica o il codice tramite premi verificabili), addestrare un unico modello general-purpose che eccella in domini eterogenei (es. matematica, chat aperta e istruzioni) presenta sfide critiche.

Il paper identifica due strategie comuni che falliscono a causa dell'interferenza cross-domain:

• SEQUENTIAL RL (RL Sequenziale): Addestrare il modello su un dominio dopo l'altro. Questo causa due fenomeni negativi:
  • Forgetting (Dimenticanza): L'ottimizzazione su un nuovo dominio degrada le prestazioni su quello precedente.
  • Rigidity (Rigidità): L'addestramento su un dominio "rigido" (come la matematica) limita la capacità del modello di esplorare comportamenti flessibili necessari per altri compiti (come la chat).
• MIXED RL (RL Misto/Multi-task): Combinare dati di diversi domini nello stesso batch. Questo causa conflitti di gradiente: l'aggiornamento dei parametri richiesto per migliorare un compito può essere in direzione opposta a quello richiesto per un altro, annullando i progressi di entrambi.

2. Metodologia: Modular Gradient Surgery (MGS)

Per risolvere i conflitti nel Mixed RL, gli autori introducono la Modular Gradient Surgery (MGS).

L'intuizione fondamentale è che i Transformer sono architetture altamente modulari: i livelli MLP sono associati alla memoria e alla conoscenza, mentre i livelli di Attention gestiscono il routing delle informazioni e il ragionamento. I conflitti di gradiente non sono distribuiti uniformemente, ma sono spesso localizzati in moduli specifici.

Il meccanismo di MGS:
Invece di applicare una "chirurgia del gradiente" globale (come il metodo PCGrad, che proietta l'intero vettore del gradiente), MGS opera a livello di singolo modulo ($\phi \in \Phi$):

1. Si calcolano i gradienti per ogni compito all'interno del batch.
2. Per ogni modulo (es. un blocco Attention o MLP), si verifica se il prodotto scalare tra i gradienti di due compiti $i$ e $j$ è negativo ($v_i \cdot v_j < 0$), indicando un conflitto.
3. Se esiste un conflitto, il gradiente del compito $i$ viene proiettato sul piano ortogonale al gradiente del compito $j$, eliminando la componente distruttiva.
4. Questo processo permette di preservare gli aggiornamenti compatibili e mitigare solo quelli che causano interferenza, evitando di essere eccessivamente conservativi (problema tipico della chirurgia globale).

3. Contributi Chiave

• Analisi Sistematica: Una classificazione rigorosa dei fallimenti del RL multi-dominio (Mode Interference: Forgetting vs. Rigidity).
• Nuovo Algoritmo (MGS): Un approccio modulare alla gestione dei gradienti che si adatta alla struttura interna dei Transformer.
• Efficienza Computazionale: Dimostrazione che l'overhead di memoria e tempo è trascurabile, specialmente quando si utilizzano framework di training parallelo come FSDP.

4. Risultati Sperimentali

Gli autori hanno testato MGS su modelli Llama-3.1-8B e Qwen-2.5-7B su tre domini: Matematica, Chat Generale e Instruction Following.

• Miglioramento delle prestazioni: MGS ha ottenuto incrementi medi di 4.3 punti (16.6%) su Llama e 4.5 punti (11.1%) su Qwen rispetto al Mixed RL standard.
• Superamento dei baselines: MGS supera costantemente le strategie di Model Merging, Normalized Advantage e la Global Gradient Surgery (GGS).
• Scalabilità: MGS dimostra di essere efficace anche aumentando il numero di compiti (aggiungendo l'Instruction Following) e aumentando il numero di step di addestramento, dove le strategie sequenziali tendono invece a degradare.
• Bilanciamento: A differenza del training sequenziale, che crea modelli "specialisti" ma sbilanciati, MGS produce modelli con capacità elevate e bilanciate in tutti i domini testati.

5. Significato e Impatto

Il lavoro è significativo perché fornisce una soluzione pratica e strutturale al problema della generalizzazione nel post-training tramite RL. Dimostra che per costruire modelli di ragionamento universali non è sufficiente aumentare i dati o i parametri, ma è necessario gestire la dinamica dell'ottimizzazione in modo che la specializzazione in un dominio non avvenga a scapito della versatilità cognitiva del modello. MGS rappresenta un passo avanti verso la creazione di agenti AI capaci di ragionare correttamente in contesti diversi senza perdere la fluidità comunicativa.