Technische Samenvatting: REAR: Test-time Preference Realignment door Reward Decomposition

Probleemstelling

Het afstemmen van Large Language Models (LLMs) op diverse en gepersonaliseerde gebruikersvoorkeuren is een cruciale uitdaging. Hoewel post-training methoden zoals Reinforcement Learning from Human Feedback (RLHF) effectief zijn, vereisen ze kostbare datacuratie en extra training, wat ze inflexibel maakt voor dynamische of specifieke gebruikersbehoeften. Test-Time Scaling (TTS) biedt een training-vrije alternatief door extra rekenkracht tijdens de inferentie te benutten om de outputkwaliteit te verbeteren. Bestaande TTS-methoden zijn echter grotendeels beperkt tot verifieerbare domeinen (bijv. wiskunde, programmeren) waar de juistheid gemakkelijk kan worden beoordeeld. Het toepassen van TTS op voorkeursafstemming (preference alignment) is moeilijk omdat de kwaliteit van een respons holistisch is en niet reduceerbaar tot een eenvoudig verifieerbaar antwoord. Huidige benaderingen voor afstemming tijdens de testtijd leggen vaak vaste computationele overhead op of vereisen het trainen van value heads, wat de schaalbaarheid en efficiëntie beperkt.

Methodologie

Het artikel introduceert REAR (REAlignment Reward), een nieuw framework dat afstemming van voorkeuren tijdens de testtijd frameert als een realignment-probleem. De kerninzicht is dat het trainingsdoel van een basismodel de voorkeur van een specifieke gebruiker tijdens de inferentie mogelijk niet optimaal bevredigt. REAR pakt dit aan door de onderliggende beloningsfunctie (reward function) te decomponeren in twee componenten en hun belang dynamisch te herschalen.

1. Reward Decomposition

De auteurs formaliseren de beloningsstructuur van een voorgetraind LLM binnen het Maximum Entropy RL-framework. Zij stellen dat de beloning $r(s \oplus x_p, a)$ voor een staat $s$ (vraag) en voorkeur $x_p$ gedecomposeerd kan worden in:

• Vraag-gerelateerde beloning ($r_0$): De impliciete beloning die het basismodel maximaliseert om de vraag correct te beantwoorden.
• Voorkeur-gerelateerde beloning ($r_p$): De beloningscomponent die vastlegt hoe goed de respons voldoet aan de voorkeur $x_p$.

De samengestelde beloning wordt gemodelleerd als $r(s \oplus x_p, a) = r_0(s, a) + \alpha r_p(s \oplus x_p, a)$, waarbij $\alpha$ een inherente balancerende coëfficiënt is.

2. Realignment Reward (REAR)

Om het model te herafstemmen (realign) voor een specifieke gebruiker, introduceert REAR een controleerbare coëfficiënt $\hat{\alpha}$ (of $\lambda = \hat{\alpha}/\alpha$) om de voorkeurscomponent tijdens de inferentie te herschalen:
$$r_{REAR}(s \oplus x_p, a) = r_0(s, a) + \hat{\alpha} r_p(s \oplus x_p, a)$$

Cruciaal is dat de auteurs een computationeel efficiënte formulering voor deze beloning afleiden. Door gebruik te maken van de relatie tussen optimale policies en value functions in MaxEnt RL, tonen zij aan dat $r_{REAR}$ uitgedrukt kan worden als een lineaire combinatie van token-niveau policy log-waarschijnlijkheden:
$$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)$$
Deze scorefunctie is afhankelijk van enkel de log-waarschijnlijkheden van het basismodel onder twee contexten: de prompt met alleen de vraag en de volledige prompt inclusief voorkeuren. Dit vereist geen training van een aparte reward model of het genereren van extra tokens voor scoring.

3. Integratie met Test-Time Scaling

De REAR-score wordt geïntegreerd in twee TTS-algoritmen:

• Best-of-N (BoN) Sampling: Genereer $N$ responsen en selecteer de respons met de hoogste $S_{REAR}$-score.
• Diverse Verifier Tree Search (DVTS): Een tree search algoritme dat REAR-scores gebruikt om tak-expansie en selectie te sturen, waarbij diversiteit wordt bevorderd terwijl geoptimaliseerd wordt voor voorkeursafstemming.

Belangrijkste Bijdragen

1. Formalisering van Realignment: Het artikel formaliseert afstemming van voorkeuren tijdens de testtijd als een realignment-probleem en stelt REAR voor als een computationeel efficiënte beloning afgeleid van een gedecomposeerd doel.
2. Training-vrije Scoring: REAR is uitsluitend afgeleid van de interne waarschijnlijkheden van het basismodel, waardoor er geen aanvullende training, training van een reward model of speculatieve generatiestappen nodig zijn.
3. Schaalbare Algoritmen: De auteurs ontwikkelen twee schaalbare test-time algoritmen (BoN en DVTS) die geleid worden door REAR en die de generatieprestaties kunnen verbeteren met variërende sample-budgetten.
4. Generaliseerbaarheid: De methode blijkt te generaliseren buiten expliciete voorkeursafstemming naar wiskundige redeneringen en visuele taken door taakspecifieke voorkeurstekst te injecteren.

Experimentele Resultaten

De auteurs evalueerden REAR op verschillende benchmarks met behulp van Qwen en Llama modellen:

• Benchmarks voor Voorkeursafstemming: Op PrefEval (dekking van expliciete/impliciete voorkeuren en keuzes), Multifaceted Bench, en Ping-Pong (role-playing), presteerden REAR-gestuurde methoden (zowel BoN als DVTS) consequent beter dan baselines inclusief greedy decoding, Amulet, Linear Alignment (LA), en Best-of-N met Generative Reward Models (GenRM). Bijvoorbeeld, DVTS met REAR verbeterde de prestaties met gemiddeld 11,6% ten opzichte van greedy decoding over verschillende metrieken.
• Robuustheid bij Lange Context: REAR demonstreerde superieure robuustheid in scenario's met lange context (tot 16k tokens) vergeleken met GenRM, dat significant achteruitging naarmate de contextlengte toenam.
• Wiskundig Redeneren: Bij toepassing op wiskundige taken (MATH500, AIME, AMC) met passende voorkeurstekst (bijv. "Acteer als een rigoureuze wiskundige"), presteerde de REAR-gestuurde selectie beter dan majority voting, wat consistente schaalvoordelen liet zien.
• Visuele Taken: Op MMHal-Bench voor visuele hallucinatie-detectie, verbeterde REAR de algemene scores en verminderde het de hallucinatiecijfers vergeleken met greedy decoding en GenRM.
• Cross-Model Transfer: De methode generaliseerde naar Llama-3.1-8B-Instruct zonder per-model tuning, wat bevestigt dat de beloningsdecompositie niet specifiek is voor de Qwen-familie.
• Efficiëntie: REAR-gestuurde methoden waren aanzienlijk sneller dan baselines die externe reward modellen of iteratieve optimalisatie vereisen (zoals Amulet), omdat ze geen aparte modellen hoeven te laden en vertrouwen op een enkele forward pass voor scoring.

Betekenis en Claims

Het artikel claimt dat REAR een schaalbare en algemene oplossing biedt voor het afstemmen van LLM's op gebruikersvoorkeuren tijdens de testtijd zonder de overhead van training of externe reward modellen. Door de beloning te decomponeren in vraag- en voorkeurscomponenten, maakt REAR dynamische controle mogelijk over de afweging tussen het beantwoorden van de vraag en het naleven van de voorkeuren.

De auteurs benadrukken dat REAR test-time scaling mogelijk maakt voor subjectieve, open eind-domeinen (zoals voorkeursafstemming en role-playing) waar traditionele verifiers falen. Zij merken op dat hoewel een standaard $\lambda = 20$ robuust werkt over veel taken, de methode de mogelijkheid biedt om de sterkte van de realignment af te stemmen. Het werk positioneert zichzelf als een stap naar het toegankelijker en efficiënter maken van test-time alignment, waarmee de kloof tussen het succes van TTS in verifieerbare domeinen en de complexiteit van subjectieve voorkeursafstemming wordt overbrugd.

Erkende Beperkingen: De auteurs merken op dat REAR een aanpasbare hyperparameter ($\lambda$) aan de gebruiker blootstelt, wat kleine validatie-sweeps kan vereisen voor taken met significant andere voorkeursverdelingen. Daarnaast blijft het identificeren van het optimale computationele budget een richting voor toekomstig werk, hoewel TTS lichter is dan fine-tuning.