Sintesi Tecnica: Modelli di Linguaggio Prototipici (PRISM)

Definizione del Problema

Gli attuali Large Language Models (LLM) affrontano una sfida fondamentale nell'Attribuzione dei Dati di Addestramento (TDA): determinare quali specifici esempi di addestramento guidino l'output di un modello. Nelle architetture dense standard, i dati di addestramento influenzano le predizioni indirettamente attraverso effetti accumulati su miliardi di parametri all'interno di percorsi densi ed entrelacciati. Ciò rende l'influenza di un singolo esempio distribuita e non tracciabile senza analisi post-hoc costose, approssimative e spesso statisticamente errate.

I metodi di TDA basati sul gradiente esistenti (ad es., le Funzioni di Influenza) affrontano due problemi composti su scala:

1. Inefficienza Computazionale: Il calcolo esatto del Prodotto Hessiano-Vettore inverso (iHVP) è impraticabile per modelli con miliardi di parametri, rendendo necessarie approssimazioni che scambiano la fedeltà con la scalabilità.
2. Errata Specificazione Statistica: Questi metodi spesso assumono obiettivi deterministici e convessi ottimizzati verso la convergenza, condizioni che non si verificano nell'addestramento stocastico degli LLM. Di conseguenza, le stime di attribuzione possono non corrispondere al ri-addestramento "leave-one-out" e sono sensibili agli iperparametri della pipeline di addestramento (profondità, dimensione del batch, ordine dei dati).

Il documento sostiene che la TDA non debba essere trattata meramente come un problema di stima post-hoc, ma come un problema di progettazione architettonica.

Metodologia: Architettura PRISM

Gli autori introducono i Prototypes for Interpretable Sequence Modeling (PRISM), un'architettura di modello linguistico che genera predizioni tramite un mix sparso e non negativo di prototipi appresi, piuttosto che tramite un head di output denso.

Architettura Core

PRISM mantiene un backbone Transformer decoder standard, ma sostituisce l'ultimo percorso di output denso con una ricostruzione prototipica sparsa:

1. Prototype Bank: Viene mantenuto un insieme di $K$ vettori prototipici appresi $P = [p_1, \dots, p_K] \in \mathbb{R}^{d \times K}$.
2. Attivazione Sparsa: Per un dato stato nascosto $z_t$, il modello calcola le similitudini del coseno con tutti i prototipi, applica ReLU per imporre la non-negatività e seleziona i top-$k$ prototipi attivi ($k \ll K$).
3. Ricostruzione: Lo stato nascosto viene decomposto in una ricostruzione prototipica $\hat{z}_t$ e un residuo $r_t$:
   $$z_t = \hat{z}_t + r_t = \sum_{i \in \mathcal{K}_t} a_{t,i} p_i + r_t$$
   dove $\mathcal{K}_t$ è l'insieme dei prototipi attivi e $a_{t,i}$ sono le loro attivazioni.
4. Formazione dei Logit: I logit del token successivo sono formati additivamente:
   $$\ell_t = W z_t = W r_t + \sum_{i \in \mathcal{K}_t} a_{t,i} (W p_i)$$
   Qui $W p_i$ rappresenta una "firma token-logit" fissa per ogni prototipo.

Obiettivi di Addestramento

PRISM è addestrato end-to-end con quattro termini di perdita per bilanciare la fedeltà del modello linguistico con l'interpretabilità dei prototipi:

• Cross-Entropy ($L_{CE}$): Garantisce la qualità standard della predizione del token successivo.
• Loss di Ricostruzione ($L_{REC}$): Penalizza il residuo, incoraggiando i prototipi a spiegare la maggior parte dello stato nascosto.
• Loss di Clustering ($L_{R1}, L_{R2}$): Questi obiettivi simmetrici ancorano i prototipi a vicinati coerenti dei dati di addestramento. $L_{R1}$ attira i prototipi verso le rappresentazioni dei token più vicine nel batch, mentre $L_{R2}$ attira gli stati dei token verso i prototipi più vicini. Ciò assicura che i prototipi giacciano sul manifold dei dati piuttosto che derivare in direzioni arbitrarie di un dizionario sparso.

Intuizione Teorica: Localizzazione della Curvatura

Un contributo teorico chiave è l'analisi dell'Hessiano della perdita nello spazio dei prototipi. Gli autori dimostrano che gli obiettivi di clustering localizzano la curvatura:

• L'Hessiano si separa in un termine di co-utilizzo globale (che cattura come i prototipi vengono usati insieme) e termini locali a blocchi diagonali (che catturano la geometria attorno a ciascun prototipo).
• Un clustering più forte riduce l'accoppiamento tra i prototipi, rendendo l'Hessiano meglio condizionato.
• Questa struttura abilita una TDA scalabile: poiché la curvatura è localizzata, i punteggi di influenza possono essere calcolati rapidamente utilizzando statistiche sparse memorizzate (cached), evitando la necessità di prodotti di gradiente completi nello spazio dei parametri.

Contributi Chiave

1. Architettura (PRISM): La prima famiglia di modelli linguistici prototipici su scala LLM per la predizione del token successivo. Estende il ragionamento basato su casi al modeling autoregressivo, esponendo un mix sparso e ancorato ai dati di addestramento di prototipi per ogni predizione.
2. Teoria e Attribuzione: Dimostra che le strutture prototipiche sparse e gli obiettivi di clustering localizzano la curvatura nel panorama della perdita. Ciò produce un Hessiano strutturato che consente l'attribuzione dei dati di addestramento $\sim$500$\times$ più velocemente rispetto alle baseline dense basate sul gradiente (come EK-FAC), consumando una memoria equivalente.
3. Performance Empirica e Controllo:
   • Performance: Su modelli che vanno da 130M a 1.6B di parametri addestrati su fino a 50B di token, PRISM eguaglia o rimane entro 2.5 punti percentuali dei baseline GPT densi su benchmark downstream.
   • Controllabilità: I controller prototipici possono migliorare l'accuratezza downstream di circa 3 punti ripesando i prototipi attivi, con correzioni riconducibili a specifici vicinati di addestramento.
   • Allineamento: La soppressione mirata dei prototipi può rimuovere comportamenti indesiderati (ad es., contenuti NSFW) al momento dell'inferenza senza finetuning o perdite misurabili nella qualità della generazione.

Risultati Sperimentali

Scalabilità e Performance

• Scala: PRISM è stato scalato da 130M a 1.6B di parametri su corpora FineWeb e Nemotron-CC (fino a 50B di token).
• Accuratezza: A 1.6B di parametri, PRISM addestrato su 10B di token ha raggiunto una performance comparabile ai baseline GPT densi. Con 20B–50B di token, ha colmato il divario con i baseline densi, raggiungendo un'accuratezza media downstream entro 2.5 punti.
• Efficienza: Lo strato prototipico aggiunge un overhead minimo (ad es., un aumento del 1.7% dei parametri alla scala XL) e raggiunge il 97.3% del throughput di addestramento dei modelli GPT densi.
• Grounding: Le metriche di copertura dei prototipi ($\bar{R}_1, R_2$) sono rimaste elevate (0.94–1.00), indicando che i prototipi sono rimasti radicati nei contesti reali di addestramento piuttosto che derivare.

Efficienza dell'Attribuzione

• Velocità: L'attribuzione dei prototipi tramite cache ha ottenuto un incremento di velocità di 470$\times$ rispetto a EK-FAC su un modello da 124M sotto vincoli di memoria GPU equivalenti.
• Preservazione del Segnale: I punteggi in cache hanno preservato circa l'80% delle sequenze di addestramento con il ranking più alto rispetto ai gradienti completi dello spazio prototipico, dimostrando che la geometria sparsa e localizzata mantiene il necessario segnale di attribuzione.

Controllo e Steering

• Controller: Un controller lineare sparso sui prototipi attivi ha migliorato l'accuratezza su quattro task di scelta multipla non binari di +2.72 punti in media, superando significativamente la baseline di un controller basato su SAE (+0.32 punti).
• Allineamento della Sicurezza: Sopprimendo 86 prototipi identificati come associati a NSFW (lo 0.52% del bank) al momento dell'inferenza, il punteggio medio NSFW è sceso da 0.668 a 0.154 su 50.000 generazioni, con la qualità del testo invariata. Nei casi più difficili (punteggio NSFW base = 2), il punteggio è sceso a 0.144.

Significato e Rivendicazioni

Il documento sostiene che PRISM dimostri un cambiamento nel modo in cui i modelli linguistici interpretabili dovrebbero essere progettati. Invece di fare affidamento su difficili analisi post-hoc per recuperare tracce dei dati di addestramento dai parametri densi, PRISM espone direttamente coordinate sparse e ancorate all'addestramento all'interno del percorso di predizione.

Gli autori sostengono che, quando i modelli raggiungono prestazioni predittive simili, le architetture le cui computazioni sono più facili da attribuire, auditare e controllare siano preferibili. PRISM ottiene questo:

1. Rendendo l'attribuzione dei dati di addestramento una proprietà nativa e computazionalmente economica dell'architettura del modello.
2. Permettendo l'intervento diretto (steering, soppressione) sul comportamento del modello senza finetuning.
3. Fornendo una "scienza della generalizzazione del modello" dove comportamenti specifici al tempo del test possono essere ricondotti a specifici vicinati di addestramento.

Il lavoro suggerisce che il modeling "interpretabile per design" è fattibile su larga scala, offrendo un'alternativa pratica all'entrelacciamento degli effetti tipico dei normali LLM densi.