InfAlign: Inference-aware language model alignment
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , , , , , , , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2024
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866913999830908928 |
|---|---|
| author | Balashankar, Ananth Sun, Ziteng Berant, Jonathan Eisenstein, Jacob Collins, Michael Hutter, Adrian Lee, Jong Nagpal, Chirag Prost, Flavien Sinha, Aradhana Suresh, Ananda Theertha Beirami, Ahmad |
| author_facet | Balashankar, Ananth Sun, Ziteng Berant, Jonathan Eisenstein, Jacob Collins, Michael Hutter, Adrian Lee, Jong Nagpal, Chirag Prost, Flavien Sinha, Aradhana Suresh, Ananda Theertha Beirami, Ahmad |
| contents | Language model alignment is a critical step in training modern generative language models. Alignment targets to improve win rate of a sample from the aligned model against the base model. Today, we are increasingly using inference-time algorithms (e.g., Best-of-N, controlled decoding, tree search) to decode from language models rather than standard sampling. We show that this train/test mismatch makes standard RLHF framework sub-optimal in view of such inference-time methods. To this end, we propose a framework for inference-aware alignment (InfAlign), which aims to optimize inference-time win rate of the aligned policy against the base model. We prove that for any inference-time decoding procedure, the optimal aligned policy is the solution to the standard RLHF problem with a transformation of the reward. This motivates us to provide the calibrate-and-transform RL (InfAlign-CTRL) algorithm to solve this problem, which involves a reward calibration step and a KL-regularized reward maximization step with a transformation of the calibrated reward. For best-of-N sampling and best-of-N jailbreaking, we propose specific transformations offering up to 3-8% improvement on inference-time win rates. Finally, we also show that our proposed reward calibration method is a strong baseline for optimizing standard win rate. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2412_19792 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | InfAlign: Inference-aware language model alignment Balashankar, Ananth Sun, Ziteng Berant, Jonathan Eisenstein, Jacob Collins, Michael Hutter, Adrian Lee, Jong Nagpal, Chirag Prost, Flavien Sinha, Aradhana Suresh, Ananda Theertha Beirami, Ahmad Machine Learning Computation and Language Information Theory Language model alignment is a critical step in training modern generative language models. Alignment targets to improve win rate of a sample from the aligned model against the base model. Today, we are increasingly using inference-time algorithms (e.g., Best-of-N, controlled decoding, tree search) to decode from language models rather than standard sampling. We show that this train/test mismatch makes standard RLHF framework sub-optimal in view of such inference-time methods. To this end, we propose a framework for inference-aware alignment (InfAlign), which aims to optimize inference-time win rate of the aligned policy against the base model. We prove that for any inference-time decoding procedure, the optimal aligned policy is the solution to the standard RLHF problem with a transformation of the reward. This motivates us to provide the calibrate-and-transform RL (InfAlign-CTRL) algorithm to solve this problem, which involves a reward calibration step and a KL-regularized reward maximization step with a transformation of the calibrated reward. For best-of-N sampling and best-of-N jailbreaking, we propose specific transformations offering up to 3-8% improvement on inference-time win rates. Finally, we also show that our proposed reward calibration method is a strong baseline for optimizing standard win rate. |
| title | InfAlign: Inference-aware language model alignment |
| topic | Machine Learning Computation and Language Information Theory |
| url | https://arxiv.org/abs/2412.19792 |