Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis

1. Il Problema

L'Apprendimento in Contesto (In-Context Learning - ICL) è una proprietà fondamentale dei Large Language Models (LLM) che permette loro di risolvere nuovi task basandosi su esempi forniti nel prompt, senza aggiornamenti dei pesi del modello. Nonostante il successo empirico, il meccanismo sottostante rimane poco compreso. La ricerca attuale è divisa in due paradigmi principali che non riescono a fornire una visione unificata:

1. Paradigma Introspezione (Component-level): Si concentra su singoli componenti del modello, come le Induction Heads (IH), analizzando come l'ablazione di questi influenzi le prestazioni. Tuttavia, questo approccio spesso non spiega come questi head realizzino l'ICL o come si comportino sotto input variabili.
2. Paradigma Olistico (Functional-level): Analizza l'ICL come un processo funzionale, decomponendolo in due componenti: Riconoscimento del Task (Task Recognition - TR), ovvero identificare lo spazio delle etichette possibili, e Apprendimento del Task (Task Learning - TL), ovvero imparare la mappatura specifica tra testo ed etichetta dagli esempi. Questo approccio, però, non riesce a tracciare queste funzioni fino a componenti specifici del modello.

Manca un quadro teorico che unisca la precisione meccanicistica (identificazione degli head) con la chiarezza funzionale (separazione tra TR e TL).

2. Metodologia

Gli autori propongono un nuovo framework basato sull'analisi degli attention heads per localizzare e separare le funzioni di TR e TL.

Task Subspace Logit Attribution (TSLA)

Il cuore della metodologia è un nuovo metodo chiamato TSLA, progettato per identificare quali attention head contribuiscono specificamente al TR e al TL.

• Per il Task Recognition (TR): Gli autori calcolano un punteggio basato sulla proiezione dell'output dell'head sullo spazio dei task (il sottospazio generato dai vettori di unembedding delle etichette di dimostrazione). Un head è considerato un "TR head" se il suo output è fortemente allineato a questo sottospazio, indicando che sta raccogliendo informazioni sullo spazio delle etichette possibili.
• Per il Task Learning (TL): Il punteggio TL misura l'allineamento dell'output dell'head con la direzione media tra l'etichetta corretta e quelle errate all'interno dello spazio del task. Un "TL head" è quello che ruota lo stato nascosto verso l'etichetta corretta, massimizzando il divario (logit gap) rispetto alle etichette concorrenti.

Analisi Geometrica e Steering

Oltre all'identificazione, gli autori utilizzano:

• Ablazione: Rimozione selettiva degli head identificati per osservare l'impatto su accuratezza e sul "Task Recognition Ratio" (la proporzione di predizioni che rientrano nell'insieme delle etichette di contesto).
• Steering (Guidaggio): Iniezione degli output degli head identificati (come vettori di task) negli stati nascosti di input a zero-shot per vedere se è possibile recuperare le prestazioni dell'ICL.
• Analisi Geometrica: Studio dell'evoluzione degli stati nascosti per verificare come TR e TL influenzino la geometria dello spazio latente (allineamento al sottospazio vs. rotazione interna).

3. Contributi Chiave

1. Decomposizione TR/TL a livello di Attention Head: Gli autori derivano TSLA per mappare i componenti funzionali dell'ICL (TR e TL) su specifici attention head, risolvendo il divario tra analisi olistica e meccanicistica.
2. Separabilità Causale: Dimostrano attraverso ablation studies e perturbazioni di input che TR e TL sono controllabili in modo indipendente. La rimozione degli head TR distrugge la capacità di riconoscere lo spazio delle etichette, mentre la rimozione degli head TL riduce l'accuratezza senza necessariamente distruggere la riconoscibilità del task.
3. Ruolo Geometrico Distinto:
   • TR Heads: Allineano gli stati nascosti allo spazio delle etichette del task (Task Subspace), sopprimendo token irrilevanti e permettendo il riconoscimento del task.
   • TL Heads: Eseguono rotazioni all'interno di questo sottospazio per orientare lo stato nascosto verso l'etichetta corretta, affinando la previsione.
4. Riconciliazione con le Induction Heads (IH): Dimostrano che le famose Induction Heads sono in realtà un sottoinsieme degli TR heads. La loro importanza risiede nel riconoscere lo spazio delle etichette, non nell'apprendimento specifico della mappatura testo-etichetta.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su diversi modelli (Llama3, Qwen, Yi) e dataset (SST-2, SUBJ, TREC, ecc.).

• Correlazione e Sovrapposizione: Gli head TR mostrano una forte sovrapposizione e correlazione con le Induction Heads, specialmente negli strati più profondi della rete. Al contrario, gli head TL sono più variabili tra i dataset e si trovano in strati leggermente più superficiali.
• Distribuzione dell'Attenzione:
  • Gli TR heads prestano attenzione principalmente ai token delle etichette di dimostrazione (es. "positive", "negative").
  • Gli TL heads prestano attenzione principalmente ai token della query (il testo da classificare) per estrarne il significato semantico.
• Ablazione:
  • Ablare i TR heads riduce drasticamente sia l'accuratezza che il Task Recognition Ratio (il modello non sa più quali etichette sono possibili).
  • Ablare i TL heads riduce l'accuratezza ma lascia il Task Recognition Ratio quasi intatto (il modello sa quali etichette sono possibili, ma sbaglia la mappatura).
• Steering e Zero-Shot:
  • In task di classificazione con etichette fisse, i vettori di task costruiti dai TR heads sono essenziali per recuperare le prestazioni dell'ICL partendo da zero-shot.
  • In task di generazione aperta (es. scrivere una recensione), i vettori costruiti dai TL heads sono più efficaci, poiché il task richiede di apprendere la mappatura specifica piuttosto che riconoscere un insieme chiuso di etichette.
• Geometria: L'analisi conferma che l'aggiunta degli output TR aumenta l'allineamento degli stati nascosti al sottospazio del task, mentre l'aggiunta degli output TL aumenta la discriminazione (logit difference) verso l'etichetta corretta all'interno di quel sottospazio.

5. Significato e Implicazioni

Questo lavoro fornisce una spiegazione unificata e interpretabile del funzionamento dell'ICL nei LLM:

• Unificazione Teorica: Collega le osservazioni sulle Induction Heads e sui Task Vectors a un quadro meccanicistico coerente, mostrando che l'ICL è un processo a due stadi: prima si riconosce il task (TR), poi si impara la mappatura (TL).
• Interpretabilità: Offre strumenti pratici (TSLA) per identificare e manipolare specifici componenti del modello per controllare il comportamento dell'ICL.
• Progettazione di Modelli: Suggerisce che per migliorare l'ICL in scenari diversi (classificazione vs. generazione), è necessario ottimizzare diverse componenti: la capacità di allineamento al sottospazio (TR) per task strutturati e la capacità di rotazione discriminativa (TL) per task generativi complessi.
• Spiegazione del Fallimento Zero-Shot: Il fallimento delle prestazioni zero-shot è attribuito principalmente alla mancanza di allineamento al sottospazio del task (carenza di TR), piuttosto che all'incapacità di apprendere la mappatura.

In sintesi, il paper dimostra che l'ICL non è un fenomeno monolitico, ma è implementato da circuiti neurali distinti e specializzati che possono essere localizzati, analizzati e manipolati con precisione.