GenARM: 시험 시간 정렬을 위한 자기 회귀 보상 모델을 활용한 보상 지도 생성

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

October 10, 2024

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

cs.AI

초록

대형 언어 모델(Large Language Models, LLMs)은 인상적인 능력을 보여주지만 인간의 선호와 조화를 이루어야 합니다. 기존의 훈련 시간 방법은 인간 선호 데이터셋을 사용하여 LLMs를 세밀하게 조정하지만 상당한 훈련 비용이 발생하며 다양한 사용자 선호를 처리하기 위해 반복적인 훈련이 필요합니다. 시험 시간 조정 방법은 훈련 없이 동결된 LLMs를 안내하기 위해 보상 모델(Reward Models, RMs)을 사용하여 이 문제에 대처합니다. 그러나 기존의 시험 시간 접근 방식은 완전한 응답을 평가하는 데 사용되는 경로 수준의 RMs에 의존하므로 부분 응답에서 다음 토큰 보상을 계산해야 하는 자기회귀 텍스트 생성에는 적합하지 않습니다. 이를 해결하기 위해 우리는 Autoregressive Reward Model을 활용하는 GenARM이라는 시험 시간 조정 방법을 소개합니다. 이는 효율적이고 효과적인 자기회귀 생성을 위해 설계된 새로운 보상 매개변수화인 Autoregressive Reward Model을 활용합니다. 이론적으로 이 매개변수화가 KL-정규화 강화 학습 프레임워크 내에서 기존 RMs가 달성 가능한 모든 분포로 동결된 LLMs를 안내할 수 있음을 증명합니다. 실험 결과는 GenARM이 이전의 시험 시간 조정 기준을 크게 능가하고 훈련 시간 방법의 성능과 일치함을 보여줍니다. 게다가 GenARM은 대형 LLMs를 작은 RMs와 조화시키는 효율적인 약한-강한 안내를 가능하게 하며, 더 큰 모델을 훈련하는 높은 비용 없이 다양한 사용자 선호에 대응할 수 있는 다중 목적 조정을 지원합니다.

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.