GenARM — метод потокенного реворда

Сегодня разберём простую, но интересную статью. Авторы сделали потокенный реворд, чтобы использовать его в тест-тайме для генерации ответов.

Попыток использовать реворд для генерации ответов предпринималось немало. Скажем, можно использовать обученный на полных ответах реворд на частях генерации. А можно считать награду, генерируя полный ответ для каждого следующего токена. У таких подходов есть минусы. В первом случае при генерации могут возникать неточности из-за того, что реворд обучался только на полных ответах, во втором случае — существенно возрастает «стоимость» инференса.

Решением проблем, по мнению авторов, могло бы стать использование суммы авторегрессионного RM-скоринга для каждого токена-кандидата и LLM-скоринга. На основе полученных результатов и должен выбираться ответ. Инженеры задались целью создать именно такой реворд.

Авторы взяли SFT-модель и данные предпочтений (preference data) в виде пар. Это инструкция и два ответа — победный и проигрышный, — размеченные людьми. Реворд-модель обучается на этих парах с использованием негативного лосса.

Идея авторов статьи заключается в том, чтобы представить итоговый реворд как сумму потокенных вероятностей для каждого токена. То есть каждый следующий токен получает какую-то оценку, эти оценки складываются для получения итоговой награды. Эту параметризацию подставляют в лосс, чтобы обучить реворд-модель.

В тест-тайме авторы получают скоры для каждого токена по формуле, которая учитывает предсказания базовой и реворд-моделей. Это намного эффективнее, чем генерировать целые тексты и прогонять их через реворд.

Можно также использовать несколько ревордов — например, когда ответ должен быть одновременно и полезным, и этичным, или когда нужно склонить генерацию в какую-либо сторону. Для обоих показателей нужно натренировать отдельную реворд-модель.

Эксперименты показали, что метод, предложенный авторами, оказывается лучше, чем другие известные бейзлайны — например, ARGS и Transfer Q — по качеству и скорости инференса. Однако он уступает DPO, который намного более сложен и дорог в исполнении, чем GenARM.

Ещё из интересного: авторы заметили, что маленькие модели могут выступать хорошим ревордом у крупных при использовании GenARM. Эксперименты проводили на Tulu2 с числом параметров 7B, 12B и 70B. И в этом случае метод из статьи превзошёл всё, кроме DPO.

Разбор подготовил ❣ Илья Черемушкин

Душный NLP