Q-Probe: A Lightweight Approach to Reward Maximization for Language Models
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2024
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866911897944588288 |
|---|---|
| author | Li, Kenneth Jelassi, Samy Zhang, Hugh Kakade, Sham Wattenberg, Martin Brandfonbrener, David |
| author_facet | Li, Kenneth Jelassi, Samy Zhang, Hugh Kakade, Sham Wattenberg, Martin Brandfonbrener, David |
| contents | We present an approach called Q-probing to adapt a pre-trained language model to maximize a task-specific reward function. At a high level, Q-probing sits between heavier approaches such as finetuning and lighter approaches such as few shot prompting, but can also be combined with either. The idea is to learn a simple linear function on a model's embedding space that can be used to reweight candidate completions. We theoretically show that this sampling procedure is equivalent to a KL-constrained maximization of the Q-probe as the number of samples increases. To train the Q-probes we consider either reward modeling or a class of novel direct policy learning objectives based on importance weighted policy gradients. With this technique, we see gains in domains with ground-truth rewards (code generation) as well as implicit rewards defined by preference data, even outperforming finetuning in data-limited regimes. Moreover, a Q-probe can be trained on top of an API since it only assumes access to sampling and embeddings. Code: https://github.com/likenneth/q_probe . |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2402_14688 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Q-Probe: A Lightweight Approach to Reward Maximization for Language Models Li, Kenneth Jelassi, Samy Zhang, Hugh Kakade, Sham Wattenberg, Martin Brandfonbrener, David Machine Learning We present an approach called Q-probing to adapt a pre-trained language model to maximize a task-specific reward function. At a high level, Q-probing sits between heavier approaches such as finetuning and lighter approaches such as few shot prompting, but can also be combined with either. The idea is to learn a simple linear function on a model's embedding space that can be used to reweight candidate completions. We theoretically show that this sampling procedure is equivalent to a KL-constrained maximization of the Q-probe as the number of samples increases. To train the Q-probes we consider either reward modeling or a class of novel direct policy learning objectives based on importance weighted policy gradients. With this technique, we see gains in domains with ground-truth rewards (code generation) as well as implicit rewards defined by preference data, even outperforming finetuning in data-limited regimes. Moreover, a Q-probe can be trained on top of an API since it only assumes access to sampling and embeddings. Code: https://github.com/likenneth/q_probe . |
| title | Q-Probe: A Lightweight Approach to Reward Maximization for Language Models |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2402.14688 |