Soft Head Selection for Injecting ICL-Derived Task Embeddings

1. Il Problema

I Large Language Models (LLM) vengono adattati ai compiti specifici principalmente attraverso due paradigmi:

• Fine-tuning efficiente dei parametri (PEFT): Come LoRA o (IA)³. Offre prestazioni elevate ma richiede l'aggiornamento di parametri (anche se pochi) e comporta costi di addestramento.
• Apprendimento nel contesto (ICL): Utilizza esempi nel prompt per guidare il modello senza addestramento. Tuttavia, richiede prompt lunghi, aumenta l'overhead di inferenza e le prestazioni possono essere instabili.

Recentemente è emerso un terzo approccio: l'adattamento basato su embedding guidato dall'ICL. Questo metodo estrae "vettori di compito" (task embeddings) dalle attivazioni intermedie del modello durante un'ICL a pochi shot (few-shot) e li inietta durante l'inferenza zero-shot.
Il limite attuale: Sebbene concettualmente promettente, questo approccio non ha finora dimostrato vantaggi consistenti rispetto al PEFT o all'ICL standard. Le prestazioni sono spesso limitate e i metodi esistenti non riescono a identificare in modo efficace dove e come iniettare queste informazioni nel modello, spesso basandosi su euristiche o spazi di ricerca ristretti.

2. Metodologia: SITE (Soft head-selection for ICL-derived Task Embeddings)

Gli autori propongono SITE, un metodo che identifica dinamicamente le "testine di attenzione" (attention heads) rilevanti per un compito specifico e inietta gli embedding derivati dall'ICL solo in quelle testine.

Il processo si articola in tre fasi:

1. Costruzione dell'Embedding del Compito (Stage 1):

   • Si utilizzano prompt ICL a pochi shot (es. 10 esempi) per generare le attivazioni del modello.
   • Vengono estratte le attivazioni dell'ultimo token per ogni testina di attenzione di ogni layer.
   • Queste attivazioni vengono mediate su tutti i prompt di few-shot per creare un vettore di embedding specifico per ogni testina di attenzione ($t^{(l,h)}$).

2. Ottimizzazione della Selezione Soft delle Testine (Stage 2):

   • Invece di selezionare rigidamente le testine, SITE introduce un parametro continuo e apprendibile $\alpha^{(l,h)} \in [0, 1]$ per ogni testina di attenzione.
   • Questo parametro controlla il grado di iniezione dell'embedding del compito rispetto all'attivazione originale:
     $$o^{(l,h)} \leftarrow (1 - \alpha^{(l,h)}) \cdot o^{((l,h)} + \alpha^{(l,h)} \cdot t^{(l,h)}$$
   • Il modello LLM rimane congelato. Vengono ottimizzati solo i parametri $\alpha$ tramite discesa del gradiente (usando la cross-entropy su prompt zero-shot) per massimizzare le prestazioni del compito.
   • Questo trasforma la selezione delle testine in un problema di ottimizzazione continua, permettendo al modello di "imparare" quali testine sono cruciali per quel compito specifico.

3. Inferenza Zero-shot (Stage 3):

   • Durante l'inferenza, si utilizzano gli embedding del compito pre-calcolati e i parametri $\alpha$ ottimizzati.
   • L'iniezione avviene una sola volta, all'ultimo token del prompt di input iniziale (prima della generazione autoregressiva), sfruttando la cache KV. Questo mantiene l'efficienza dell'inferenza zero-shot senza la necessità di ripetere l'iniezione ad ogni passo di generazione.

3. Contributi Chiave

• Selezione Soft delle Testine: Un metodo basato sul gradiente che identifica in modo differenziabile le testine di attenzione rilevanti, superando i limiti delle selezioni rigide o euristica.
• Efficienza dei Parametri: SITE richiede un numero di parametri addestrabili estremamente ridotto (solo i parametri di selezione $\alpha$, circa 1K per un modello da 8B) rispetto al PEFT (milioni di parametri), pur mantenendo il modello base congelato.
• Analisi Meccanicistica: Gli autori estendono l'analisi tramite activation patching per dimostrare che:
  • Le testine selezionate catturano effettivamente informazioni rilevanti per il compito (patching su testine ad alto $\alpha$ migliora drasticamente le prestazioni).
  • Esiste una forte dipendenza dal compito: compiti simili condividono le stesse testine importanti, mentre compiti diversi ne utilizzano di diverse. Questo sposta la prospettiva da una visione "agnostica rispetto al compito" a una "specificità per compito" delle funzionalità delle testine.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su 12 LLM diversi (famiglie Llama, Mistral, Qwen, Gemma) con dimensioni da 4B a 70B parametri, su 4 benchmark principali (FV, ANLI, MMLU-Pro, Big-Bench Hard).

• Prestazioni Superiori: SITE supera significativamente i metodi di adattamento basati su embedding precedenti (come FV, TV, MTV) e l'ICL a pochi shot (10-shot).
  • Su FV Benchmark, SITE supera l'ICL a 10-shot di circa il 13.26%.
  • Le prestazioni sono paragonabili o superiori al PEFT (LoRA, (IA)³) su compiti di generazione aperta e comprensione del linguaggio, avvicinandosi alle prestazioni del PEFT su compiti complessi di ragionamento.
• Generalizzazione: Il metodo funziona coerentemente su modelli di diverse dimensioni e famiglie, dimostrando una forte generalità.
• Efficienza:
  • Parametri: Richiede ~1.02K parametri addestrabili contro i ~3.4M di LoRA.
  • Tempo: Il tempo di inferenza è simile a quello zero-shot e inferiore all'ICL a molti shot, poiché l'iniezione avviene una sola volta all'inizio.
  • Memoria: Overhead di memoria minimo (~0.5 MB per Llama-3.1-8B).

5. Significato e Implicazioni

Il lavoro di SITE rappresenta un passo avanti significativo nell'adattamento efficiente dei LLM:

1. Nuovo Paradigma di Adattamento: Dimostra che è possibile ottenere prestazioni vicine al fine-tuning completo senza modificare i pesi del modello, utilizzando solo l'iniezione mirata di informazioni estratte dal contesto.
2. Interpretabilità: Fornisce nuove intuizioni meccanicistiche, confermando che le funzioni delle testine di attenzione non sono fisse ma altamente specifiche per il compito. Questo apre la strada a future ricerche su come sfruttare queste strutture per un'inferenza controllabile e adattiva.
3. Praticità: Offre una soluzione pratica per scenari dove l'addestramento è costoso o impossibile, ma si dispone di alcuni esempi di addestramento (few-shot) per derivare embedding di alta qualità.

In sintesi, SITE risolve il problema dell'inefficienza e della scarsa performance dei metodi precedenti di "steering" delle attivazioni, introducendo un meccanismo di selezione delle testine basato sull'ottimizzazione continua che bilancia perfettamente efficienza computazionale e prestazioni di task.