GenARM: Beloningsgestuurde Generatie met Autoregressief Beloningsmodel voor Uitlijning op Testtijd

GenARM: Reward Guided Generation with Autoregressive Reward Model for Test-time Alignment

October 10, 2024

Auteurs: Yuancheng Xu, Udari Madhushani Sehwag, Alec Koppel, Sicheng Zhu, Bang An, Furong Huang, Sumitra Ganesh

cs.AI

Samenvatting

Grote Taalmodellen (LLM's) vertonen indrukwekkende mogelijkheden, maar vereisen zorgvuldige afstemming op menselijke voorkeuren. Traditionele trainingsmethoden finetunen LLM's met behulp van datasets met menselijke voorkeuren, maar brengen aanzienlijke trainingskosten met zich mee en vereisen herhaalde training om diverse gebruikersvoorkeuren aan te kunnen. Testtijd-afstemmingsmethoden pakken dit aan door beloningsmodellen (RMs) te gebruiken om bevroren LLM's te begeleiden zonder opnieuw te trainen. Echter, bestaande testtijdbenaderingen vertrouwen op trajectniveau RMs die zijn ontworpen om complete reacties te evalueren, waardoor ze ongeschikt zijn voor autoregressieve tekstgeneratie die het berekenen van beloningen voor het volgende token vereist vanuit gedeeltelijke reacties. Om dit aan te pakken, introduceren we GenARM, een testtijd-afstemmingsbenadering die gebruikmaakt van het Autoregressieve Beloningsmodel - een nieuw beloningsparametrisatie ontworpen om beloningen voor het volgende token efficiënt en effectief te voorspellen voor autoregressieve generatie. Theoretisch tonen we aan dat deze parametrisatie bevroren LLM's aantoonbaar kan begeleiden naar elke distributie die haalbaar is met traditionele RMs binnen het KL-geregulariseerde reinforcement learning kader. Experimentele resultaten tonen aan dat GenARM aanzienlijk beter presteert dan eerdere testtijd-afstemmingsbaselines en overeenkomt met de prestaties van trainingsmethoden. Bovendien maakt GenARM efficiënte zwak-naar-sterk begeleiding mogelijk, waarbij grotere LLM's worden afgestemd op kleinere RMs zonder de hoge kosten van het trainen van grotere modellen. Verder ondersteunt GenARM multi-objectieve afstemming, waardoor real-time afwegingen tussen voorkeursdimensies mogelijk zijn en tegemoet wordt gekomen aan diverse gebruikersvoorkeuren zonder opnieuw te trainen.

English

Large Language Models (LLMs) exhibit impressive capabilities but require careful alignment with human preferences. Traditional training-time methods finetune LLMs using human preference datasets but incur significant training costs and require repeated training to handle diverse user preferences. Test-time alignment methods address this by using reward models (RMs) to guide frozen LLMs without retraining. However, existing test-time approaches rely on trajectory-level RMs which are designed to evaluate complete responses, making them unsuitable for autoregressive text generation that requires computing next-token rewards from partial responses. To address this, we introduce GenARM, a test-time alignment approach that leverages the Autoregressive Reward Model--a novel reward parametrization designed to predict next-token rewards for efficient and effective autoregressive generation. Theoretically, we demonstrate that this parametrization can provably guide frozen LLMs toward any distribution achievable by traditional RMs within the KL-regularized reinforcement learning framework. Experimental results show that GenARM significantly outperforms prior test-time alignment baselines and matches the performance of training-time methods. Additionally, GenARM enables efficient weak-to-strong guidance, aligning larger LLMs with smaller RMs without the high costs of training larger models. Furthermore, GenARM supports multi-objective alignment, allowing real-time trade-offs between preference dimensions and catering to diverse user preferences without retraining.