Value Augmented Sampling

Valor de muestreo aumentado

Este repositorio contiene la implementación oficial para nuestro documento, muestreo aumentado de valor para alinear y personalizar modelos de idiomas.

Abstracto

Alinear modelos de idiomas grandes (LLM) para atender a diferentes preferencias humanas, aprender nuevas habilidades y desaprender un comportamiento dañino es un problema importante. Los métodos basados ​​en la búsqueda, como la búsqueda de árboles Best-of-N o Monte-Carlo, son desempeñados, pero son poco prácticos para la adaptación de LLM debido a su alto costo de inferencia. Por otro lado, el uso del aprendizaje de refuerzo (RL) para la adaptación es computacionalmente eficiente, pero funciona peor debido a los desafíos de optimización en la formación conjunta de la función de valor y la política. Presentamos un nuevo marco para la optimización de recompensas, el muestreo aumentado de valor (VAS), que puede maximizar las diferentes funciones de recompensa utilizando datos muestreados de solo la LLM inicial y congelada. VAS se resuelve para la política óptima de maximización de recompensas sin capacitar la política y la función de valor, lo que hace que la optimización estable, superan las líneas de base establecidas, como PPO y DPO, en puntos de referencia estándar, y logre resultados comparables a Best Of 128 con menor costo de inferencia. A diferencia de los métodos RL existentes que requieren cambiar los pesos de la LLM, VAS no requiere acceso a los pesos de la LLM previamente capacitada. Por lo tanto, incluso puede adaptar LLM (por ejemplo, ChatGPT), que están disponibles solo como API. Además, nuestro algoritmo desbloquea la nueva capacidad de componer varias recompensas y controlar el alcance de cada uno durante el tiempo de despliegue, allanando el camino por delante para el futuro de los LLM alineados y personalizados.

Instalación

Instale nuestra versión personalizada de trl :

 git clone [email protected]:idanshen/trl.git
cd trl
python setup.py install

Clon e instale la base de código:

 git clone [email protected]:idanshen/Value-Augmented-Sampling.git
cd Value-Augmented-Sampling
pip install -e .

Cómo usar

Proporcionamos un script para capacitar a un modelo Tinyllama-1B como el estimador de valor de un modelo LLAMA-2 7B en el conjunto de datos HH de Anthrope.

Para seguir la tubería descrita en el documento, proporcionamos una versión ajustada supervisada de estos modelos:

 python tinyllama_hh.py --log_with=wandb --ref_model_name hanseungwook/vas-llama-2-7b-hh-sft --model_name hanseungwook/vas-tiny-llama-1.1b-hh-sft

Citación

@inproceedings{
  han2024value,
  title={Value Augmented Sampling for Language Model Alignment and Personalization},
  author={Seungwook Han and Idan Shenfeld and Akash Srivastava and Yoon Kim and Pulkit Agrawal},
  booktitle={ICLR 2024 Workshop on Reliable and Responsible Foundation Models},
  year={2024},
  url={https://arxiv.org/abs/2405.06639}
}