Riassunto Tecnico: TaskPGM – Apprendimento di Miscele di Task tramite Affinità di Task

Definizione del Problema

Il Fine-Tuning Supervisionato (SFT) per i Large Language Models (LLM) dipende fortemente dalla distribuzione del budget di addestramento attraverso un insieme eterogeneo di task di istruzione. Le pratiche attuali si affidano spesso a semplici euristiche, come il campionamento uniforme o la ponderazione proporzionale alla dimensione del dataset. Questi metodi non tengono conto delle interazioni tra i task, portando a due principali inefficienze:

1. Ridondanza: Sovrapesatura di task che inducono aggiornamenti comportamentali quasi identici, sprecando il budget computazionale.
2. Opportunità di Trasferimento Mancate: Sottoutilizzo di task che offrono informazioni complementari o benefici di trasferimento cross-domain.

Gli approoli automatizzati esistenti (ad esempio, il recupero basato sull'influenza come LESS, o la selezione submodulare come SMART) ottimizzano spesso sotto-insiemi discreti piuttosto che proporzioni di miscela continue, si affidano alla somiglianza semantica che ignora il comportamento del modello a valle, o richiedono un costoso addestramento iterativo. La sfida centrale è determinare automaticamente una miscela di task continua che sia simultaneamente rappresentativa (alta utilità) e diversificata (bassa ridondanza) per massimizzare le prestazioni a valle sotto un budget fisso.

Metodologia: TaskPGM

Gli autori propongono TaskPGM (Task Mixtures via Probabilistic Graphical Modelling), un framework che inquadra la miscelazione dei dati SFT come un problema di minimizzazione dell'energia su un campo casuale di Markov (MRF) denso.

1. Formulazione del Modello Grafico Probabilistico

I task sono modellati come nodi in un MRF. L'obiettivo è trovare una distribuzione di campionamento ottimale $p^* \in \Delta_n$ (dove $\Delta_n$ è il simplesso di probabilità) che massimizzi una funzione di energia $E(p)$. La funzione di energia consiste in due componenti:

• Potenziali Unari ($\Psi_{un}$): Codificano l'utilità o la rappresentatività dei singoli task. Questo è derivato dalla massa totale di somiglianza di un task rispetto all'intera collezione.
• Potenziali Pari a Coppia ($\Psi_{pair}$): Codificano ridondanza e diversità. Questi termini penalizzano la co-occorrenza di task altamente simili nella miscela.

L'obiettivo di ottimizzazione è formulato come:
$$\max_{p \in \Delta_n} E(p) = \Psi_{un}^\top p - \frac{1}{2} p^\top \Psi_{pair} p$$
Questo è equivalente alla minimizzazione di una funzione quadratica convessa (sotto specifiche condizioni) soggetta a vincoli di simplesso.

2. Metriche di Affinità Comportamentale dei Task

Una chiave innovazione è il metodo per calcolare le affinità di coppia ($S_{ij}$). Invece di affidarsi alla somiglianza semantica (embedding) o ai metadati, TaskPGM calcola le affinità basandosi su divergenze comportamentali derivate dalle distribuzioni predittive di modelli sottoposti a fine-tuning su singolo task.

• Informazione Mutua Punto per Punto (PMI): Misura l'informazione mutua tra le distribuzioni predittive di due task quando valutate l'una sugli altri dataset.
• Divergenza di Jensen–Shannon (JSD): Una misura simmetrica della discrepanza tra le distribuzioni predittive di due modelli su esempi cross-task.

Queste metriche catturano come i task modellano il comportamento del modello su esempi fuori dal task, affrontando direttamente il trasferimento e la ridondanza tra i task.

3. Ottimizzazione e Solver

• Miscela Continua: Gli autori derivano una soluzione in forma chiusa per le pesature della miscela continua utilizzando le condizioni di Karush-Kuhn-Tucker (KKT).
  • Se la matrice di somiglianza a coppia non è Semidefinita Positiva (PSD), applicano lo shifting dello spettro (aggiungendo una costante sulla diagonale) per garantire la convessità, il che introduce un termine di regolarizzazione che orienta la soluzione verso l'uniformità se la struttura non-PSD è severa.
  • Quando l'inversa è mal condizionata o i vincoli di non-negatività vengono violati, utilizzano la Discesa del Gradiente Proiettato (PGD) sul simplesso di probabilità.
• Variante Discreta/Budgetizzata: Per scenari che richiedono la selezione di un sottoinsieme di task (top-$k$), gli autori elevano l'ottimizzazione continua a una funzione di insieme discreta. Essi dimostrano che questa funzione è monotona e debolmente submodulare sotto standard condizioni di regolarità (Ristretta Forza Convessa e Smoothness), consentendo una costruzione greedy efficiente con garanzie di approssimazione.

Contributi Chiave

1. Framework Probabilistico per il Design della Miscela: Una nuova formulazione della miscelazione dei dati SFT come minimizzazione dell'energia su un MRF, che scambia esplicitamente rappresentatività (termini unari) e diversità (termini a coppia).
2. Metriche di Affinità Basate sul Comportamento: La proposta di calcolare le relazioni tra i task utilizzando le divergenze delle distribuzioni predittive (JSD/PMI) da modelli sottoposti a fine-tuning su singolo task, andando oltre la statica somiglianza semantica per catturare le dipendenze funzionali tra i task.
3. Ottimizzazione Supportata Teoricamente:
   • Derivazione di una soluzione in forma chiusa per le miscele continue tramite condizioni KKT.
   • Stabilimento della debole submodularità per la variante discreta con budget, supportando la selezione greedy con garanzie teoriche.
4. Validazione Empirica: Dimostrazione di guadagni di performance costanti su molteplici famiglie di modelli (Qwen2-7B, Llama-2-7B) e benchmark (BIG-Bench Hard, MMLU, ecc.) rispetto alle euristiche e ai baseline di selezione allo stato dell'arte.

Risultati Sperimentali

Il framework è stato valutato su 309 task di istruzione utilizzando i modelli Qwen2-7B e Llama-2-7B con budget di 25K e 50K istanze.

• Guadagni di Performance: TaskPGM (utilizzando sia le varianti PMI che JSD) ha superato costantemente il campionamento uniforme, la proporzionalità dimensionale (EPM) e i baseline casuali.
  • Su Qwen2-7B (budget 25K), la variante PMI ha migliorato le prestazioni di quasi il 3% rispetto ai baseline euristici su benchmark di ragionamento complesso.
  • Su Llama-2-7B, la miscela derivata da PMI ha superato Uniform e EPM del 2% al 4% o più su BBH.
• Confronto con SOTA: TaskPGM è rimasto competitivo o superiore a LESS (un forte baseline basato su funzioni di influenza). Significativamente, TaskPGM raggiunge questo risultato con un singolo checkpoint general-purpose, mentre LESS costruisce modelli distinti e target-aware per ogni benchmark.
• Efficienza: TaskPG-M è computazionalmente più efficiente di LESS. Il calcolo della somiglianza è un costo ammortizzato una tantum (~16–34 ore GPU), dopo il quale le miscele per qualsiasi budget possono essere derivate in secondi tramite la soluzione in forma chiusa. Al contrario, LESS richiede ~320 ore GPU per costruire i datastore di gradiente per ogni target.
• Scalabilità: L'aumento del budget da 25K a 50K ha prodotto guadagni consistenti per TaskPGM, mentre i metodi euristici (Uniform, EPM) hanno mostrato un comportamento frammentato e non robusto all'aumentare dei budget.

Significato e Rivendicazioni

Il paper sostiene che TaskPGM fornisce un approccio principato, interpretabile ed efficiente alla miscelazione dei dati SFT.

• Interpretabilità: I pesi appresi e le mappe di calore della somiglianza rivelano cluster di task, ridondanza e task "attrattori", offrendo intuizioni sul perché certe fonti siano utili.
• Generalizzazione: A differenza dei metodi che ottimizzano per target specifici, TaskPGM produce una singola miscela che migliora le prestazioni in un ampio spettro di task a valle (ragionamento, comprensione del linguaggio, conoscenza).
• Robustezza: I guadagni sono guidati dalla selezione funzionale dei task basata sul comportamento del modello piuttosto che da una favorevole varianza di campionamento, come dimostrato dai miglioramenti costanti attraverso diversi modelli base e metriche di somiglianza.

Gli autori concludono che una miscelazione basata sulla somiglianza e sui principi è essenziale per l'instruction tuning scalabile, spesso eguagliando o superando le miscele euristiche con una quantità di dati sostanzialmente inferiore. Riconoscono i limiti riguardanti il pool fisso di task e il costo computazionale della stima dell'affinità, suggerendo approssimazioni scalabili come direzione futura per il lavoro.