Контролируемое декодирование языковых моделей

Controlled Decoding from Language Models

October 25, 2023

Авторы: Sidharth Mudgal, Jong Lee, Harish Ganapathy, YaGuang Li, Tao Wang, Yanping Huang, Zhifeng Chen, Heng-Tze Cheng, Michael Collins, Trevor Strohman, Jilin Chen, Alex Beutel, Ahmad Beirami

cs.AI

Аннотация

Мы предлагаем управляемое декодирование (Controlled Decoding, CD) — новый метод обучения с подкреплением вне политики, предназначенный для управления авторегрессионной генерацией языковых моделей в направлении достижения высоких значений награды. CD решает задачу обучения с подкреплением вне политики с помощью функции ценности для награды, которую мы называем оценщиком префиксов. Оценщик префиксов используется на этапе вывода для направления генерации к результатам с более высокой наградой. Мы показываем, что оценщик префиксов может быть обучен на (возможно) данных вне политики для предсказания ожидаемой награды при продолжении декодирования из частично декодированного ответа. Эмпирически мы демонстрируем, что CD эффективен как механизм управления на корпусе диалогов Reddit. Мы также показываем, что модульность конструкции CD позволяет управлять несколькими наградами, эффективно решая задачу многокритериального обучения с подкреплением без дополнительной сложности. Наконец, мы показываем, что CD может быть применен в новом блочном режиме на этапе вывода, снова без необходимости изменений на этапе обучения, что фактически устраняет разрыв между популярной стратегией best-of-K и обучением с подкреплением на уровне токенов. Это делает CD перспективным подходом для согласования языковых моделей.

English

We propose controlled decoding (CD), a novel off-policy reinforcement learning method to control the autoregressive generation from language models towards high reward outcomes. CD solves an off-policy reinforcement learning problem through a value function for the reward, which we call a prefix scorer. The prefix scorer is used at inference time to steer the generation towards higher reward outcomes. We show that the prefix scorer may be trained on (possibly) off-policy data to predict the expected reward when decoding is continued from a partially decoded response. We empirically demonstrate that CD is effective as a control mechanism on Reddit conversations corpus. We also show that the modularity of the design of CD makes it possible to control for multiple rewards, effectively solving a multi-objective reinforcement learning problem with no additional complexity. Finally, we show that CD can be applied in a novel blockwise fashion at inference-time, again without the need for any training-time changes, essentially bridging the gap between the popular best-of-K strategy and token-level reinforcement learning. This makes CD a promising approach for alignment of language models.