Guiding Giants: Lightweight Controllers for Weighted Activation Steering in LLMs

1. Il Problema

I Large Language Models (LLM) sono vulnerabili alla generazione di contenuti dannosi, tossici o non allineati ai valori umani. Le attuali strategie di sicurezza presentano due limiti principali:

• Fine-tuning costoso: L'addestramento estensivo (es. RLHF, filtraggio dati) è oneroso in termini di risorse computazionali e può causare "dimenticanza catastrofica" o degradazione delle prestazioni su compiti utili.
• Metodi di steering (guida) a inferenza insufficienti: Le tecniche esistenti di activation steering (come ITI, CAST, ACT) spesso applicano modifiche fisse e uniformi su tutti i livelli del modello o mancano di adattabilità fine. Se l'intervento è troppo forte, il modello rifiuta indiscriminatamente anche prompt innocui; se è troppo debole, fallisce nel bloccare richieste dannose. Manca un controllo granulare che adatti la forza e la posizione dell'intervento in base al singolo prompt.

2. Metodologia: Weighted Activation Steering (WAS)

Gli autori propongono WAS, un meccanismo di controllo a tempo di inferenza che utilizza una rete controller leggera e trainabile per modulare dinamicamente le attivazioni del modello.

Architettura e Funzionamento

Il processo si articola in quattro fasi:

1. Caching delle Attivazioni: Un LLM congelato (frozen) elabora prompt dannosi e benigni, salvando le attivazioni intermedie.
2. Pre-computazione del Vettore di Rifiuto ($d_{steer}$): Viene calcolato un vettore di direzione fisso che rappresenta la differenza tra le medie degli embedding di token di rifiuto (es. "Mi dispiace", "Non posso") e token di risposta (es. "Certo", "Ecco").
   $$d_{steer} = \frac{\bar{e}_{refuse} - \bar{e}_{answer}}{||\bar{e}_{refuse} - \bar{e}_{answer}||_2}$$
3. Training del Controller ($f_c$): Una rete MLP leggera (Multi-Layer Perceptron) viene addestrata in modo discriminativo.
   • Input: Le attivazioni concatenate degli strati finali del modello (tipicamente l'ultimo terzo) per un dato prompt.
   • Output: La rete predice due valori:
     • Uno scalare globale ($s$): indica l'intensità dell'intervento (vicino a 1.0 per prompt dannosi, vicino a 0.0 per quelli benigni).
     • Un vettore di pesi per strato ($w$): una serie di pesi (tramite funzione sigmoide) che determina quanto applicare la correzione in ciascuno strato del modello.
   • Loss Function: Mean Squared Error (MSE) per far sì che $s \approx 1$ per i prompt dannosi e $s \approx 0$ per quelli benigni. I pesi $w$ vengono appresi implicitamente attraverso la condivisione dei layer nascosti, imparando quali strati sono più critici per diversi tipi di minacce.
4. Applicazione Dinamica a Inferenza: Durante la generazione, il controller osserva le attivazioni, calcola $s$ e $w$, e applica una patch alle attivazioni nascoste ($h_l$) di ogni strato $l$:
   $$\Delta h_{l} = s \cdot w_l \cdot \alpha_{global} \cdot d_{steer}$$
   Dove $\alpha_{global}$ è un iperparametro fisso che scala l'intensità complessiva.

3. Contributi Chiave

• Controllo Granulare e Adattivo: A differenza dei metodi precedenti che usano vettori fissi, WAS adatta sia l'intensità globale ($s$) sia la distribuzione dell'intervento tra gli strati ($w$) in base al contenuto specifico del prompt.
• Efficienza Computazionale: Il controller è una rete MLP estremamente leggera che non modifica i pesi del modello base, permettendo un'intervento a basso overhead.
• Apprendimento di Pattern Strutturati: L'analisi mostra che il controller impara a identificare "hotspot" specifici negli strati del modello per diverse categorie di pericolo (es. cybercrime vs. molestie), dimostrando che l'intervento mirato è più efficace di quello uniforme.
• Implementazione Pratica: L'approccio è implementato tramite hooks PyTorch, rendendolo facilmente integrabile senza riscrivere il codice del modello base.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su Llama-3.1-8B, Llama-3.2-1B e Mistral-7B.

• Sicurezza (Refusal Rate):
  • Su ToxicChat, WAS ha raggiunto un tasso di rifiuto del 93.0% (vs 32% del modello base e <50% di altri metodi come CAST o ITI).
  • Su AdvBench, ha ottenuto fino al 98.8% di rifiuto.
  • Supera costantemente CAST, ACT e ITI su tutti i benchmark di sicurezza testati.
• Utilità e Capacità Generali:
  • Su AlpacaEval (valutazione di utilità su prompt benigni), WAS mantiene prestazioni statisticamente indistinguibili dal modello base (win rate ~49.8% vs 50%).
  • Su benchmark di ragionamento e conoscenza (MMLU, GSM8K, HellaSwag), il degrado delle prestazioni è minimo o nullo (es. MMLU scende solo da 63.0% a 60.8%).
• Efficienza: WAS offre un ottimo compromesso tra latenza e sicurezza, essendo significativamente più veloce di ITI e più sicuro di CAST.

5. Significato e Implicazioni

Il lavoro dimostra che è possibile ottenere un controllo comportamentale robusto e mirato negli LLM senza il costo del fine-tuning completo.

• Paradigma Shift: Sposta l'attenzione da modifiche globali e statiche a interventi dinamici e specifici per istanza e per strato.
• Interpretabilità: L'analisi dei pesi appresi ($w$) rivela che il modello apprende strategie differenziate per diverse categorie di rischio, fornendo intuizioni su come le rappresentazioni di sicurezza siano distribuite nell'architettura del Transformer.
• Limitazioni: Il metodo dipende dalla qualità del vettore di steering pre-calcolato e potrebbe essere vulnerabile ad attacchi avversari progettati per ingannare il controller (es. prompt che generano attivazioni simili a quelle benignhe pur essendo dannosi). Inoltre, la generalizzazione a nuove categorie di minacce non presenti nei dati di addestramento rimane una sfida.

In sintesi, WAS rappresenta un approccio promettente per implementare strati di sicurezza leggeri e adattivi nei modelli linguistici già distribuiti, bilanciando efficacemente sicurezza e utilità.