Value Augmented Sampling for Language Model Alignment and Personalization

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Han, Seungwook, Shenfeld, Idan, Srivastava, Akash, Kim, Yoon, Agrawal, Pulkit
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866929339465990144
author Han, Seungwook
Shenfeld, Idan
Srivastava, Akash
Kim, Yoon
Agrawal, Pulkit
author_facet Han, Seungwook
Shenfeld, Idan
Srivastava, Akash
Kim, Yoon
Agrawal, Pulkit
contents Aligning Large Language Models (LLMs) to cater to different human preferences, learning new skills, and unlearning harmful behavior is an important problem. Search-based methods, such as Best-of-N or Monte-Carlo Tree Search, are performant, but impractical for LLM adaptation due to their high inference cost. On the other hand, using Reinforcement Learning (RL) for adaptation is computationally efficient, but performs worse due to the optimization challenges in co-training the value function and the policy. We present a new framework for reward optimization, Value Augmented Sampling (VAS), that can maximize different reward functions using data sampled from only the initial, frozen LLM. VAS solves for the optimal reward-maximizing policy without co-training the policy and the value function, making the optimization stable, outperforming established baselines, such as PPO and DPO, on standard benchmarks, and achieving comparable results to Best-of-128 with lower inference cost. Unlike existing RL methods that require changing the weights of the LLM, VAS does not require access to the weights of the pre-trained LLM. Thus, it can even adapt LLMs (e.g., ChatGPT), which are available only as APIs. In addition, our algorithm unlocks the new capability of composing several rewards and controlling the extent of each one during deployment time, paving the road ahead for the future of aligned, personalized LLMs.
format Preprint
id arxiv_https___arxiv_org_abs_2405_06639
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Value Augmented Sampling for Language Model Alignment and Personalization
Han, Seungwook
Shenfeld, Idan
Srivastava, Akash
Kim, Yoon
Agrawal, Pulkit
Machine Learning
Artificial Intelligence
Computation and Language
Aligning Large Language Models (LLMs) to cater to different human preferences, learning new skills, and unlearning harmful behavior is an important problem. Search-based methods, such as Best-of-N or Monte-Carlo Tree Search, are performant, but impractical for LLM adaptation due to their high inference cost. On the other hand, using Reinforcement Learning (RL) for adaptation is computationally efficient, but performs worse due to the optimization challenges in co-training the value function and the policy. We present a new framework for reward optimization, Value Augmented Sampling (VAS), that can maximize different reward functions using data sampled from only the initial, frozen LLM. VAS solves for the optimal reward-maximizing policy without co-training the policy and the value function, making the optimization stable, outperforming established baselines, such as PPO and DPO, on standard benchmarks, and achieving comparable results to Best-of-128 with lower inference cost. Unlike existing RL methods that require changing the weights of the LLM, VAS does not require access to the weights of the pre-trained LLM. Thus, it can even adapt LLMs (e.g., ChatGPT), which are available only as APIs. In addition, our algorithm unlocks the new capability of composing several rewards and controlling the extent of each one during deployment time, paving the road ahead for the future of aligned, personalized LLMs.
title Value Augmented Sampling for Language Model Alignment and Personalization
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2405.06639