Riassunto Tecnico: IG-Lens

Definizione del Problema

I transformer decoder-only sono spesso concettualizzati come processi di inferenza iterativi in cui lo stream residuo affina una stima corrente del prossimo token. Una questione critica di interpretabilità sorge: in quale specifico strato viene effettivamente deciso il valore della probabilità di un token predetto?

Gli strumenti esistenti per la lettura (readout) per strato non riescono a rispondere con precisiono nello spazio delle probabilità:

• Logit Lens & Tuned Lens: Questi riportano il livello di probabilità ad ogni strato ma non forniscono una decomposizione additiva. Le loro stime sono distorte e non monotone attraverso la profondità, il che significa che la somma delle probabilità per strato non ha un significato matematico.
• Direct Logit Attribution (DLA) & Metodi di Embedding: Questi sono additivi ma operano nello spazio dei logit. Poiché la funzione softmax è altamente non lineare, una decomposizione additiva dei logit non si traduce in una decomposizione additiva delle probabilità, che è la quantità di interesse pratico.
• Layer Conductance & Layer Integrated Gradients (IG): Questi metodi integrano i gradienti per strato, ma attribuiscono ogni strato al proprio baseline indipendente. Di conseguenza, i termini per strato non "telescopano" (non si sommano) alla variazione totale della probabilità di predizione.

Metodologia: IG-Lens

Gli autori introducono IG-Lens, un metodo che fornisce una decomposizione additiva esatta della probabilità attraverso gli strati del transformer.

Concetto Centrale

IG-Lens applica un'applicazione telescopica di Integrated Gradients (IG) lungo un singolo percorso attraverso gli stati nascosti da un baseline all'ultimo strato.

1. Definizione del Percorso: Il metodo seleziona una sequenza di strati $0 = L_0 < L_1 < \dots < L_k = n$ e definisce un percorso di stati nascosti $h_{L_0} \to h_{L_1} \to \dots \to h_{final}$.
2. Readout "Head-Once": Fondamentalmente, la funzione di readout $f(h)$ è definita come l'unembedding finale applicato una sola volta a un singolo vettore nascosto: $f(h) = \text{softmax}(\text{head}(\text{norm}(h)))_{y_t}$. Questa funzione non contiene attenzione o MLP; tutto il mixing tra i token è assunto essere avvenuto nel passaggio forward che ha prodotto gli stati nascosti distaccati.
3. Attribuzione Telescopica: Il contributo dello strato $L_j$ è definito come il segmento IG di $f$ lungo la linea retta da $h_{L_{j-1}}$ a $h_{L_j}$. Per il teorema del gradiente, l'integrale del gradiente su un segmento è uguale alla differenza dei valori degli estremi. Sommando questi segmenti si ottiene una somma telescopica:
   $$\sum_{j=1}^k IGL_j = f(h_{final}) - f(h_{base}) = \Delta p$$
   Ciò garantisce che la somma delle attribuzioni sia esattamente la variazione della probabilità target, con la softmax e la LayerNorm integrate nel percorso piuttosto che linearizzate o congelate.

Stimatore Sensibile alla Predizione (Prediction-Aware Estimator)

Sebbene l'integrale sia esatto, gli stimatori a passi finiti solitamente introducono errori di discretizzazione. IG-Lens adotta un estimatore sensibile alla predizione (ispirato a IDGI) come default:

• Invece di accreditare i passi basandosi sui gradienti grezzi, accredita ogni passo in base alla sua variazione osservata nella probabilità target ($f(h^{(s)}) - f(h^{(s-1)})$).
• Per una readout di probabilità scalare, ciò collassa ogni segmento in una somma telescopica di valori di estremi.
• Risultato: La completezza si mantiene alla precisione del floating-point a qualsiasi numero di passi ($m \ge 1$). I passi in cui il modello mostra sensibilità al gradiente ma nessuna variazione in uscita vengono naturalmente filtrati, eliminando l'errore di discretizzazione di Riemann.

Implementazione

Il metodo utilizza un algoritmo batch a singolo passaggio che calcola la mappa completa token-per-strato senza richiedere un passaggio backward (autograd).

• Raggruppando tutti i punti di interpolazione in un singolo tensore ed valutando la funzione puntuale $f$ in un unico passaggio forward, il metodo evita l'overhead del calcolo del gradiente.
• Questa efficienza è possibile specificamente grazie al design "head-once", che garantisce l'indipendenza tra i punti dei token nel batch.

Contributi Chiave

1. Decomposizione Additiva Esatta: IG-Lens è il primo metodo che fornisce una decomposizione additiva esatta della probabilità (non dei logit) attraverso gli strati, dove la somma delle attribuzioni è uguale a $\Delta p$.
2. Identità Telescopica: Il paper fornisce una prova formale (Proposizione 1) che gli integrali di segmento telescopano alla variazione totale della predizione, mantenendo la softmax all'interno del percorso di integrazione.
3. Stimatore Sensibile alla Predizione: L'estimatore di default garantisce la completezza alla precisione macchina indipendentemente dal numero di passi e filtra i passi di "sensibilità senza risposta".
4. Implementazione Efficiente: Una valutazione forward a singolo passaggio calcola l'intera mappa di attribuzione senza chiamate backward, rendendolo computazionalmente efficiente.
5. Posizionamento: Gli autori delineano chiaramente perché i metodi esistenti falliscono nello spazio delle probabilità (additività dello spazio dei logit, bias o mancanza di telescopia) e posizionano IG-Lens come il contrappeso nello spazio delle probabilità ai metodi nello spazio dei logit.

Risultati e Piano di Valutazione

• Verifica Numerica: Su test controllati e su un modello reale (Llama-3.2-1B-Instruct), l'estimatore di default riproduce le differenze tra gli estremi dei segmenti a $\sim 10^{-7}$ (precisione float32). La somma telescopica si mantiene esatta rispetto al floating-point, indipendentemente dal numero di passi $m$.
• Output Illustrativo: Negli esperimenti con il prompt "What is the capital of Vietnam?", il metodo riesce a concentrare la massa di attribuzione nei segmenti di strato specifici in cui il valore del logit-lens transita da basso ad alto. La somma delle attribuzioni per ogni token corrisponde a $\Delta p$ con precisione di floating-point.
• Strategia di Valutazione: Gli autori riconoscono che l' "onset" (l'inizio) non ha una verità fondamentale intrinseca. Pertanto, propongono di validare il metodo tramite ablazione causale: se IG-Lens identifica lo strato $L$ come l'onset, l'ablazione di $L$ dovrebbe deprimere la probabilità target più di quanto faccia l'ablazione di strati con punteggi di onset inferiori. Questo test è progettato per essere economico da eseguire e capace di falsificare la tesi.

Significato e Rivendicazioni

Il paper posiziona IG-Lens come un raffinamento quantitativo della visione dell'inferenza iterativa. A differenza del Logit Lens, che chiede "qual è la predizione corrente a questo strato?", IG-Lens chiede "di quanta parte della probabilità finale è responsabile ciascuno strato?" e garantisce che queste responsabilità sommate costituiscano il tutto.

Rivendicazioni Modeste e Limitazioni:

• Marginale Condizionale: Gli autori dichiarano esplicitamente che IG-Lens misura un marginale condizionale: la probabilità aggiuntiva che la norm+head possa leggere da $h_L$ dato $h_{L-1}$. Il metodo scarta l'effetto di $h_L$ instradato attraverso blocchi di attenzione e MLP superiori.
• Compromesso: Il metodo sceglie "somma = $\Delta p$" invece dell' "effetto causale totale tramite l'intero passaggio forward". Per gli effetti causali totali, la decomposizione dei componenti (come DLA) rimane lo strumento appropriato.
• Non-Monotonicità: Il metodo gestisce le traiettorie di probabilità non monotone (dove la probabilità viene costruita e parzialmente annullata) utilizzando attribuzioni con segno, offrendo una visione più fedele dei contributi degli strati rispetto ai metodi basati su valori assoluti.

La domanda centrale aperta rimane empirica: se l'onset di IG-Lens predice lo strato in cui l'ablazione causale deprime maggiormente la predizione, particolarmente per i token con traiettorie non monotone.