Pref-CTRL: Preference Driven LLM Alignment using Representation Editing

Riassunto Tecnico: Pref-CTRL

1. Il Problema (Problem Statement)

L'allineamento dei Large Language Models (LLM) ai valori umani è fondamentale per prevenire risposte dannose, distorte o indesiderate. Sebbene metodi come il Reinforcement Learning from Human Feedback (RLHF) o il Direct Preference Optimization (DPO) siano efficaci, essi richiedono un fine-tuning costoso in termini di risorse computazionali e scalabilità.

Una soluzione emergente è l'allineamento al tempo di test (test-time alignment), che interviene sulle rappresentazioni interne del modello durante l'inferenza senza aggiornare i pesi. Il framework di riferimento in questo ambito è RE-Control, che utilizza una funzione di valore (value function) per guidare la generazione tramite l'editing dei gradienti degli stati latenti. Tuttavia, RE-Control presenta un limite critico: addestra la sua funzione di valore basandosi solo su punteggi di allineamento individuali, ignorando la natura comparativa (pairwise) dei dati di preferenza umana, che sono tipicamente strutturati come "risposta A è migliore di risposta B".

2. Metodologia (Methodology)

Gli autori propongono Pref-CTRL, un framework di addestramento multi-obiettivo che estende RE-Control integrando la struttura delle preferenze nei dati.

Architettura di base:
Il modello tratta la generazione autoregressiva come un sistema dinamico stocastico. Durante l'inferenza, un segnale di controllo $u_t$ viene aggiunto agli stati nascosti ($h_t$) e ai logit ($o_t$) per orientare la generazione verso stati con reward più elevati, calcolati tramite una funzione di valore $V_\phi$ (un MLP leggero).

Il contributo di Pref-CTRL (Nuova Funzione di Perdita):
A differenza di RE-Control, che usa solo una perdita di regressione ($L_{Regression}$), Pref-CTRL introduce due nuovi termini nella funzione di perdita totale:

1. Margin Loss ($L_{Margin}$): Forza la funzione di valore a mantenere un divario (margine) tra lo score della risposta preferita e quello della risposta rifiutata. Questo permette alla funzione di catturare meglio le gerarchie di preferenza relative.
   $$L_{Margin} = -\log \sigma(V_\phi(s_{pref}) - V_\phi(s_{rej}))$$
2. Regularization Loss ($L_{Regularizer}$): Per evitare l'over-optimization (o "reward hacking") e preservare la fluidità del linguaggio, questo termine assicura che lo score della rappresentazione generata rimanga vicino allo score della rappresentazione preferita.
   $$L_{Regularizer} = (V_\phi(s_{gen}) - V_\phi(s_{pref}))^2$$

La perdita finale è la combinazione di questi tre componenti: $L_{Total} = L_{Regression} + L_{Margin} + L_{Regularizer}$.

3. Contributi Chiave (Key Contributions)

• Allineamento basato sulle preferenze al tempo di test: È il primo approccio di representation editing che sfrutta esplicitamente la struttura pairwise dei dataset di allineamento.
• Obiettivo di addestramento multi-obiettivo: L'introduzione della Margin Loss e della Regularization Loss risolve i problemi di distinzione tra risposte e di degradazione della qualità generativa.
• Efficienza: Fornisce un metodo di allineamento leggero che non richiede il fine-tuning del modello base, rendendolo estremamente flessibile.

4. Risultati (Results)

Il metodo è stato testato su dataset standard come Stanford SHP e Anthropic HH-RLHF, confrontandolo con RE-Control e il metodo di fine-tuning DPO.

• Performance di Allineamento: Pref-CTRL ha superato costantemente RE-Control in termini di Win Rate (percentuale di risposte preferite dai giudici LLM) e Average Reward. Ad esempio, su SHP (Vicuna-7B), l'aggiunta di Margin + Regularizer ha portato il Win Rate dal 66.80% al 73.50%.
• Confronto con DPO: Sebbene DPO sia un metodo di fine-tuning (più potente), Pref-CTRL ottiene risultati molto vicini, dimostrando l'efficacia dell'intervento al tempo di test.
• Generalizzazione Out-of-Domain: Il modello ha mostrato una robustezza superiore quando testato su dataset non visti durante l'addestramento (PKU-SafeRLHF e Nectar), confermando che la funzione di valore appresa è in grado di generalizzare i concetti di preferenza.
• Qualità della generazione: I test su metriche di Diversity e Coherence indicano che Pref-CTRL mantiene la qualità del linguaggio, evitando le risposte ripetitive o incoerenti tipiche di un allineamento troppo aggressivo.

5. Significato (Significance)

Il lavoro di Pref-CTRL è significativo perché colma il divario tra i metodi di allineamento basati su preferenze (solitamente limitati al fine-tuning) e le tecniche di controllo dell'inferenza. Dimostra che è possibile "insegnare" a un modello a seguire le preferenze umane agendo direttamente sulle sue attivazioni interne, offrendo una via di mezzo efficiente tra la rigidità dei modelli pre-addestrati e l'alto costo computazionale del RLHF tradizionale.