Riepilogo Tecnico: REAR: Riallineamento delle Preferenze al Tempo di Test attraverso la Decomposizione della Ricompensa

Definizione del Problema

Allineare i Large Language Models (LLM) con le diverse e personalizzate preferenze degli utenti è una sfida critica. Sebbene i metodi di post-training come il Reinforcement Learning from Human Feedback (RLHF) siano efficaci, richiedono una costosa cura dei dati e un ulteriore addestramento, rendendoli poco flessibili per le esigenze dinamiche o specifiche dell'utente. Il Test-Time Scaling (TTS) offre un'alternativa che non richiede addestramento, sfruttando una computazione aggiuntiva durante l'inferenza per migliorare la qualità dell'output. Tuttavia, gli attuali metodi di TTS sono stati ampiamente limitati a domini verificabili (ad esempio, matematica, programmazione) dove la correttezza è facilmente giudicabile. Applicare il TTS all'allineamento delle preferenze è difficile perché la qualità della risposta è olistica e non riducibile a una semplice risposta verificabile. Gli attuali approcci di allineamento al tempo di test impongono spesso carichi computazionali fissi o richiedono l'addestramento di head di valore, limitandone la scalabilità e l'efficienza.

Metodologia

Il paper introduce REAR (REAlignment Reward), un nuovo framework che inquadra l'allineamento delle preferenze al tempo di test come un problema di riallineamento. L'intuizione centrale è che l'obiettivo di addestramento di un modello base potrebbe non soddisfare in modo ottimale la specifica preferenza di un utente al tempo di inferenza. REAR affronta questo problema scomponendo la funzione di ricompensa sottostante in due componenti e riscalando dinamicamente la loro importanza.

1. Decomposizione della Ricompensa

Gli autori formalizzano la struttura della ricompensa di un LLM pre-addestrato all'interno del framework Maximum Entropy RL. Essi postulano che la ricompensa $r(s \oplus x_p, a)$ per uno stato $s$ (domanda) e una preferenza $x_p$ possa essere decomposta in:

• Ricompensa relativa alla domanda ($r_0$): La ricompensa implicita che il modello base massimizza per rispondere correttamente alla domanda.
• Ricompensa relativa alla preferenza ($r_p$): La componente di ricompensa che cattura quanto bene la risposta soddisfi la preferenza $x_p$.

La ricompensa composita è modellata come $r(s \oplus x_p, a) = r_0(s, a) + \alpha r_p(s \oplus x_p, a)$, dove $\alpha$ è un coefficiente di bilanciamento inerente.

2. Ricompensa di Riallineamento (REAR)

Per riallineare il modello per un utente specifico, REAR introduce un coefficiente controllabile $\hat{\alpha}$ (o $\lambda = \hat{\alpha}/\alpha$) per riscalare la componente di preferenza durante l'inferenza:
$$r_{REAR}(s \oplus x_p, a) = r_0(s, a) + \hat{\alpha} r_p(s \oplus x_p, a)$$

Fondamentalmente, gli autori derivano una formulazione computazionalmente efficiente per questa ricompensa. Sfruttando la relazione tra politiche ottimali e funzioni di valore nel framework MaxEnt RL, dimostrano che $r_{REAR}$ può essere espressa come una combinazione lineare di log-probabilità di policy a livello di token:
$$S_{REAR}(\tau) = \sum_{t=0}^{T} \gamma^t \left( (1-\lambda) \log \pi(a_t|s_t) + \lambda \log \pi(a_t|s_t \oplus x_p) \right)$$
Questa funzione di punteggio dipende solo dalle log-probabilità del modello base sotto due contesti: il prompt relativo alla sola domanda e il prompt completo includendo le preferenze. Non richiede l'addestramento di un modello di ricompensa separato né la generazione di token aggiuntivi per la valutazione.

3. Integrazione con il Test-Time Scaling

Il punteggio REAR viene integrato in due algoritmi di TTS:

• Best-of-N (BoN) Sampling: Genera $N$ risposte e seleziona quella con il punteggio $S_{REAR}$ più alto.
• Diverse Verifier Tree Search (DVTS): Un algoritmo di ricerca ad albero che utilizza i punteggi REAR per guidare l'espansione e la selezione dei rami, promuovendo la diversità pur ottimizzando l'allineamento alle preferenze.

Contributi Chiave

1. Formalizzazione del Riallineamento: Il paper formalizza l'allineamento delle preferenze al tempo di test come un problema di riallineamento, proponendo REAR come una ricompensa computazionalmente efficiente derivata da un obiettivo decomposto.
2. Scoring Senza Addestramento: REAR è derivato esclusivamente dalle probabilità interne del modello base, non richiedendo addestramento aggiuntivo, addestramento di modelli di ricompensa o passaggi di generazione speculativa.
3. Algoritmi Scalabili: Gli autori sviluppano due algoritmi di test-time scalabili (BoN e DVTS) guidati da REAR che possono migliorare le prestazioni di generazione con budget di campionamento variabili.
4. Generalizzabilità: Il metodo si dimostra capace di generalizzare oltre l'allineamento esplicito delle preferenze verso il ragionamento matematico e i compiti visivi iniettando testo di preferenza specifico per il compito.

Risultati Sperimentali

Gli autori hanno valutato REAR su diversi benchmark utilizzando i modelli Qwen e Llama:

• Benchmark di Allineamento delle Preferenze: Su PrefEval (che copre preferenze esplicite/implicite e scelte), Multifaceted Bench e Ping-Pong (role-playing), i metodi guidati da REAR (sia BoN che DVTS) hanno superato costantemente i baseline, inclusi greedy decoding, Amulet, Linear Alignment (LA) e Best-of-N con Generative Reward Models (GenRM). Ad esempio, DVTS con REAR ha migliorato mediamente il greedy decoding dell'11,6% attraverso le metriche.
• Robustezza nel Long-Context: REAR ha dimostrato una robustezza superiore negli scenari a lungo contesto (fino a 16k token) rispetto a GenRM, che degradava significativamente all'aumentare della lunghezza del contesto.
• Ragionamento Matematico: Applicato ai compiti matematici (MATH500, AIME, AMC) con il testo di preferenza appropriato (ad esempio, "Agisci come un matematico rigoroso"), la selezione guidata da REAR ha superato il majority voting, mostrando guadagni di scalabilità costanti.
• Compiti Visivi: Su MMHal-Bench per il rilevamento delle allucinazioni visive, REAR ha migliorato i punteggi complessivi e ridotto i tassi di allucinazione rispetto al greedy decoding e a GenRM.
• Trasferimento Cross-Modello: Il metodo si è generalizzato a Llama-3.1-8B-Instruct senza tuning per modello, confermando che la decomposizione della ricompensa non è specifica della famiglia Qwen.
• Efficienza: I metodi guidati da REAR sono stati significativamente più veloci dei baseline che richiedono modelli di ricompensa esterni o ottimizzazione iterativa (come Amulet), poiché evitano il caricamento di modelli separati e si affidano a un singolo passaggio in avanti (forward pass) per lo scoring.

Significato e Claim

Il paper sostiene che REAR fornisca una soluzione scalabile e generale per allineare gli LLM con le preferenze degli utenti al tempo di test senza l'overhead di addestramento o modelli di ricompensa esterni. Decomponendo la ricompensa in componenti di domanda e preferenza, REAR consente il controllo dinamico del compromesso tra il rispondere alla domanda e l'aderire alle preferenze.

Gli autori sottolineano che REAR permette al test-time scaling di estendersi a domini soggettivi e aperti (come l'allineamento delle preferenze e il role-playing) dove i verifier tradizionali falliscono. Notano che, sebbene un valore predefinito di $\lambda = 20$ funzioni robustamente in molti compiti, il metodo consente di regolare la forza del riallineamento. Il lavoro si posiziona come un passo verso il rendere l'allineamento al tempo di test più accessibile ed efficiente, colmando il divario tra il successo del TTS nei domini verificabili e la complessità dell'allineamento delle preferenze soggettive.

Limitazioni Riconosciute: Gli autori osservano che REAR espone all'utente un iperparametro regolabile ($\lambda$), che potrebbe richiedere piccoli sweep di validazione per compiti con distribuzioni di preferenza significativamente diverse. Inoltre, sebbene il TTS sia più leggero del fine-tuning, l'identificazione del budget computazionale ottimale rimane una direzione per il lavoro futuro.