Efficient Estimation of Kernel Surrogate Models for Task Attribution

Riepilogo Tecnico: Stima Efficiente di Modelli Surrogati a Kernel per l'Attribuzione dei Compiti

Enunciato del Problema
I sistemi di intelligenza artificiale moderni, in particolare i grandi modelli linguistici (LLM) e gli apprendisti multi-compito, sono addestrati su insiemi diversificati di compiti (ad esempio, traduzione, generazione di codice, ragionamento). Una sfida fondamentale nell'interpretabilità è l'attribuzione dei compiti: quantificare l'influenza di ciascun compito di addestramento individuale sulle prestazioni del modello su un compito target specifico. La metrica di verità fondamentale per questo è il retraining Leave-One-Out (LOO), che misura la variazione delle prestazioni quando un compito specifico viene escluso dall'addestramento. Tuttavia, il LOO è computazionalmente proibitivo su larga scala, richiedendo $K+1$ esecuzioni complete di addestramento per $K$ compiti.

Le alternative efficienti esistenti includono le funzioni di influenza (che stimano l'effetto di perturbazioni infinitesimali dei dati tramite prodotti vettore-Hessiana) e i modelli surrogati lineari (che apprendono una mappatura lineare dagli indicatori di inclusione del compito alle prestazioni). Sebbene i surrogati lineari siano scalabili, sono limitati alla cattura di relazioni additive del primo ordine. Non riescono a modellare le interazioni non lineari tra i compiti, come sinergia, antagonismo o effetti di tipo XOR, che sono comuni in paesaggi di apprendimento complessi.

Metodologia
Il documento propone KERNELSM (Modelli Surrogati a Kernel), un framework che estende l'attribuzione dei compiti per catturare interazioni non lineari di secondo ordine e di ordine superiore tra i compiti di addestramento.

1. Analisi Teorica dei Surrogati Lineari:
   Gli autori stabiliscono innanzitutto un framework unificato di ponderazione dei compiti. Attraverso uno sviluppo in serie di Taylor del secondo ordine della funzione di prestazione $F(s)$ attorno al vettore di pesi uniforme $s^* = [1]^K$, analizzano i modelli surrogati lineari.

   • Dimostrano che i coefficienti del surrogato lineare ($\hat{\beta}$) approssimano il gradiente del paesaggio delle prestazioni ($\nabla_s F(s^*)$) fino a termini indotti dalla distribuzione di campionamento e dagli effetti di interazione del secondo ordine.
   • Crucialmente, dimostrano che, sotto l'assunzione di interazioni del secondo ordine piccole (cioè un paesaggio quasi lineare), i modelli surrogati lineari sono approssimativamente equivalenti alle funzioni di influenza. Ciò spiega perché i metodi lineari funzionano bene in alcuni contesti ma falliscono quando le interazioni tra i compiti non sono additive.

2. Modellazione Surrogata a Kernel:
   Per affrontare la non additività, gli autori sostituiscono la mappatura lineare con una funzione non lineare appresa tramite Regressione a Pettine a Kernel (KRR) in uno Spazio di Hilbert a Kernel Riproducente (RKHS).

   • Scelta del Kernel: Utilizzano il kernel a funzione di base radiale (RBF) ($k(s^{(a)}, s^{(b)}) = \exp(-\gamma \|s^{(a)} - s^{(b)}\|^2)$). Questa scelta è motivata dalla geometria dello spazio delle sottoinsiemi binari $\{0, 1\}^K$, dove il kernel RBF codifica naturalmente la località basata sulla distanza di Hamming, permettendo al modello di catturare dipendenze complesse tra le combinazioni di compiti.
   • Obiettivo: La funzione surrogata $g_\theta$ minimizza la perdita empirica tra le prestazioni osservate $F(s^{(i)})$ e le previsioni, regolarizzata dalla norma RKHS.

3. Stima Efficiente Basata su Gradienti:
   Un collo di bottiglia maggiore nella modellazione surrogata è la necessità di addestrare $m$ modelli separati per ottenere i valori di prestazione $F(s^{(i)})$ per $m$ sottoinsiemi campionati. Gli autori introducono una procedura di stima efficiente per evitare l'addestramento ripetuto:

   • Approssimazione del Primo Ordine: Sfruttano uno sviluppo in serie di Taylor del primo ordine dell'output del modello $f_W(x)$ attorno a un'inizializzazione pre-addestrata $W_0$: $f_W(x) \approx f_{W_0}(x) + \langle \nabla f_{W_0}(x), W - W_0 \rangle$.
   • Caratteristiche dei Gradienti: Invece di re-addestrare, stimano la perturbazione ottimale dei parametri $Z^*_s$ per un sottoinsieme $s$ risolvendo un problema di regressione regolarizzato in cui le caratteristiche sono i gradienti del modello pre-addestrato ($\nabla f_{W_0}$).
   • Riduzione della Dimensionalità: Per gestire gradienti ad alta dimensionalità, applicano proiezioni casuali gaussiane (lemma di Johnson-Lindenstrauss) per comprimere lo spazio delle caratteristiche, consentendo calcoli rapidi sulle CPU.
   • Accuratezza: I risultati empirici mostrano che questa approssimazione del primo ordine comporta un errore relativo inferiore al 2% su dataset e dimensioni di modello diversi (fino a 34 miliardi di parametri).

Contributi Chiave

• Unificazione Teorica: Il documento stabilisce una connessione formale tra i modelli surrogati lineari e le funzioni di influenza, mostrando che sono equivalenti in condizioni di interazioni tra compiti del secondo ordine trascurabili.
• Modelli Surrogati a Kernel (KERNELSM): L'introduzione di un approccio basato su kernel (specificamente RBF) che cattura efficacemente le interazioni non lineari tra i compiti (sinergia/antagonismo) che i modelli lineari non riescono a cogliere.
• Algoritmo di Stima Scalabile: Una procedura di stima basata su gradienti che approssima le prestazioni di qualsiasi sottoinsieme di compiti utilizzando i gradienti da un singolo modello pre-addestrato, eliminando la necessità di re-addestramento ripetuto mantenendo alta accuratezza.
• Validazione Empirica: Valutazione completa su ragionamento aritmetico modulare, apprendimento in contesto (ICL) e apprendimento per rinforzo multi-obiettivo (RL).

Risultati Sperimentali
Gli autori valutano KERNELSM rispetto a baseline che includono Funzioni di Influenza, TracIn, TRAK, SOURCE e Modelli Surrogati Lineari.

• Accuratezza dell'Attribuzione: KERNELSM supera costantemente le baseline. Nei compiti di aritmetica modulare (che coinvolgono strutture XOR e quadratiche), migliora l'accuratezza dell'attribuzione fino al 42% rispetto ai metodi esistenti. In tutti i domini, raggiunge una correlazione del 25% superiore alla verità fondamentale LOO rispetto ai surrogati lineari e alle funzioni di influenza.
• Selezione del Compito a valle: Quando utilizzato per selezionare sottoinsiemi ottimali per compiti a valle:
  • Apprendimento in Contesto: KERNELSM produce una perdita del 40% inferiore nella selezione delle dimostrazioni rispetto ai surrogati lineari.
  • Apprendimento per Rinforzo Multi-Obiettivo: Migliora le ricompense target del 15% nel benchmark Meta-World MT10.
• Scalabilità: Il metodo scala a un gran numero di compiti (ad esempio, 500 prompt) senza degradazione delle prestazioni ed esegue in modo efficiente sulle CPU.
• Robustezza: L'approccio a kernel rimane robusto sotto diversi ottimizzatori (SGD, SAM, NSO) e regolarizzazioni dell'Hessiana, mentre i surrogati lineari degradano significativamente quando il paesaggio della perdita è altamente curvo (alta traccia dell'Hessiana).

Significato e Affermazioni
Il documento afferma che KERNELSM fornisce una soluzione scalabile e flessibile per l'attribuzione dei compiti in sistemi black-box. Superando le ipotesi additive dei modelli lineari, consente una comprensione più accurata di come le interazioni complesse e non lineari tra i compiti di addestramento plasmino il comportamento del modello. Gli autori sottolineano che il loro approccio non si basa su ipotesi specifiche riguardanti il processo di addestramento (oltre all'approssimazione del primo ordine del modello pre-addestrato) ed è ampiamente applicabile a paradigmi emergenti come l'apprendimento in contesto e l'apprendimento per rinforzo. Il lavoro colma il divario tra le funzioni di influenza teoriche e l'attribuzione dei dati pratica e scalabile per i moderni sistemi di intelligenza artificiale su larga scala.