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

Аннотация

Большие языковые модели (LLM) обладают впечатляющими возможностями, но требуют тщательного согласования с предпочтениями человека. Традиционные методы обучения настраивают LLM с использованием наборов данных предпочтений людей, но при этом возникают значительные затраты на обучение и требуется повторное обучение для учета разнообразных пользовательских предпочтений. Методы выравнивания на этапе тестирования решают эту проблему, используя модели вознаграждения (RMs) для направления замороженных LLM без повторного обучения. Однако существующие подходы на этапе тестирования опираются на модели вознаграждения на уровне траектории, которые предназначены для оценки полных ответов, что делает их непригодными для авторегрессивной генерации текста, требующей вычисления вознаграждений за следующий токен из частичных ответов. Для решения этой проблемы мы представляем GenARM, подход к выравниванию на этапе тестирования, который использует Авторегрессивную Модель Вознаграждения - новую параметризацию вознаграждения, разработанную для прогнозирования вознаграждений за следующий токен для эффективной авторегрессивной генерации. Теоретически мы демонстрируем, что эта параметризация может доказуемо направлять замороженные LLM к любому распределению, достижимому с использованием традиционных RMs в рамках обучения с подкреплением с регуляризацией KL. Экспериментальные результаты показывают, что GenARM значительно превосходит предыдущие базовые линии выравнивания на этапе тестирования и соответствует результатам методов обучения на этапе обучения. Кроме того, GenARM обеспечивает эффективное руководство от слабого к сильному, выравнивая более крупные LLM с более маленькими RMs без высоких затрат на обучение более крупных моделей. Более того, GenARM поддерживает многокритериальное выравнивание, позволяя в реальном времени находить компромиссы между измерениями предпочтений и учитывая разнообразные пользовательские предпочтения без повторного обучения.

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.