Beyond Linear Steering: Unified Multi-Attribute Control for Language Models

1. Il Problema

Il controllo di più attributi comportamentali (es. tono, fattualità, tossicità, stili di dibattito) nei Large Language Models (LLM) durante l'inferenza è una sfida complessa. I metodi di "steering" (guida) esistenti si basano prevalentemente su interventi lineari nello spazio delle attivazioni. Questi approcci presentano due limitazioni fondamentali:

1. Assunzione di additività: Presumono che i vettori di steering siano additivi e che le combinazioni di comportamenti siano lineari. Tuttavia, combinare semplicemente vettori per attributi diversi spesso porta a interferenze, degradazione della fluidità del testo o fallimento nel raggiungere l'obiettivo desiderato.
2. Complessità di tuning: Richiedono vettori separati per ogni attributo e un tuning iperparametrico specifico per ciascuno, rendendo difficile la composizione dinamica di comportamenti multipli senza riaddestramento.

2. Metodologia: K-Steering

Gli autori introducono K-Steering, un approccio unificato e flessibile che supera le assunzioni di linearità. Il metodo si basa su un classificatore non lineare addestrato sulle attivazioni nascoste del modello.

Componenti Chiave:

• Classificatore MLP: Viene addestrato un Multi-Layer Perceptron (MLP) per la classificazione multi-etichetta sulle attivazioni finali di uno strato specifico del modello. L'input è la dimensione dell'attivazione ($d_{model}$) e l'output è la probabilità su $K$ categorie (es. toni o stili di dibattito).
• Intervento basato sui Gradienti: A differenza dei metodi che aggiungono vettori statici, K-Steering calcola le direzioni di intervento dinamicamente durante l'inferenza utilizzando i gradienti del classificatore.
  • La formula di base per l'aggiornamento dell'attivazione $a_i$ è:
    $$a'_i = a_i - \alpha \nabla_{a_i} \mathcal{L}(g_\phi(a_i))$$
    Dove $\mathcal{L}$ è una funzione di perdita che massimizza la probabilità delle classi target e minimizza quella delle classi da evitare, e $\alpha$ è un fattore di scala.
• Algoritmi Proposti:
  1. Steering Iterativo (Algorithm 1): Aggiorna le attivazioni in modo iterativo basandosi sul gradiente della perdita, permettendo di navigare in spazi non lineari complessi.
  2. Rimozione per Proiezione (Algorithm 2): Una tecnica più efficiente per rimuovere attributi indesiderati. Utilizza una riflessione di Householder per invertire la componente dell'attivazione lungo la direzione del gradiente associata all'attributo da evitare, senza bisogno di iterazioni multiple.

3. Contributi Chiave

1. Nuovi Dataset (Benchmark):
   • TONEBANK: Un dataset di domande progettato per essere risposto in 6 toni conversazionali diversi (es. empatico, esperto, conciso).
   • DEBATEMIX: Un dataset di domande di dibattito progettato per essere affrontato in 10 stili retorici diversi (es. Reductio ad Absurdum, Burden of Proof Shift, Empirical Grounding).
   • Questi dataset sono specifici per testare la componibilità degli attributi, ovvero la capacità di un modello di gestire combinazioni di stili in una singola risposta.
2. Metodologia Unificata: Dimostrazione che un singolo classificatore non lineare può gestire il controllo simultaneo di più direzioni, superando i limiti dei metodi lineari come CAA (Contrastive Activation Addition) e DCT (Deep Causal Transcoding).
3. Framework di Valutazione: Sviluppo di un protocollo semi-automatizzato che utilizza la ricerca binaria per calibrare l'intensità dello steering ($\alpha$), massimizzando l'efficacia del controllo minimizzando lo spostamento della distribuzione (preservando la coerenza del testo).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su tre famiglie di modelli (Llama-3.2-3B, Mistral-7B, OLMo-2-7B).

• Prestazioni Superiori: K-Steering supera i baseline (CAA e DCT) in tutti i casi di steering multi-attributo (K=1, 2, 3) su entrambi i dataset.
• Robustezza alla Diluizione: Mentre i metodi lineari tendono a diluire gli effetti quando si combinano più vettori, K-Steering mantiene una forte presenza degli attributi target, come confermato sia da classificatori basati sulle attivazioni che da giudici basati su LLM.
• Analisi Multi-Layer e Multi-Step:
  • L'applicazione dello steering su più strati simultaneamente migliora le prestazioni per compiti complessi.
  • L'uso di più passi gradiente (multi-step) con passi di apprendimento più piccoli permette di esplorare meglio il paesaggio della perdita non lineare, migliorando l'efficacia, sebbene con un costo computazionale maggiore.
• Preservazione delle Capacità: Su benchmark generali come MMLU, K-Steering mantiene le prestazioni del modello base, dimostrando di non degradare significativamente le capacità linguistiche generali.

5. Significato e Implicazioni

Il lavoro di K-Steering rappresenta un passo avanti significativo verso il controllo comportamentale compositivo negli LLM.

• Superamento della Linearità: Dimostra che le relazioni tra attributi comportamentali nello spazio latente dei modelli sono intrinsecamente non lineari e che approcci basati su gradienti di classificatori sono necessari per gestirle efficacemente.
• Flessibilità: Permette di combinare dinamicamente comportamenti (es. "rispondi in modo empatico ma conciso") senza dover addestrare nuovi vettori o riaddestrare il modello.
• Efficienza: Sebbene il costo computazionale sia superiore rispetto all'aggiunta di vettori statici (CAA), il metodo elimina la necessità di tuning manuale per ogni combinazione di attributi e offre un controllo molto più preciso.

In sintesi, K-Steering offre un framework robusto per guidare i modelli linguistici verso comportamenti complessi e multi-dimensionali, aprendo la strada a interazioni più sofisticate e controllabili con l'IA.