Value Augmented Sampling

Échantillonnage augmenté de valeur

Ce référentiel contient l'implémentation officielle de notre article, l'échantillonnage augmenté de valeur pour l'alignement et la personnalisation des modèles de langage.

Abstrait

L'alignement de grands modèles de langue (LLM) pour répondre à différentes préférences humaines, apprendre de nouvelles compétences et désapprendre un comportement nocif est un problème important. Les méthodes basées sur la recherche, telles que Best-of-N ou la recherche d'arbres Monte-Carlo, sont performantes, mais peu pratiques pour l'adaptation LLM en raison de leur coût d'inférence élevé. D'un autre côté, l'utilisation de l'apprentissage par renforcement (RL) pour l'adaptation est efficace sur le calcul, mais fonctionne moins bien en raison des défis d'optimisation dans le co-entraînement de la fonction de valeur et de la politique. Nous présentons un nouveau cadre pour l'optimisation des récompenses, l'échantillonnage augmenté de valeur (EVA), qui peut maximiser différentes fonctions de récompense à l'aide de données échantillonnées à partir de la LLM gelée initiale. VAS résout pour la politique optimale de maximisation de récompense sans co-entraîner la politique et la fonction de valeur, ce qui rend l'optimisation stable, surprenant les lignes de base établies, telles que PPO et DPO, sur des repères standard et obtenant des résultats comparables au meilleur de 128 avec Coût d'inférence inférieure. Contrairement aux méthodes RL existantes qui nécessitent de modifier les poids du LLM, l'EVA ne nécessite pas d'accès aux poids du LLM pré-formé. Ainsi, il peut même adapter les LLM (par exemple, Chatgpt), qui sont disponibles uniquement sous forme d'API. De plus, notre algorithme débloque la nouvelle capacité de composer plusieurs récompenses et de contrôler l'étendue de chacun pendant le temps de déploiement, en pavage de la route à venir pour l'avenir des LLM personnalisés alignés et personnalisés.

Installation

Installez notre version personnalisée de trl :

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

Cloner et installer la base de code:

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

Comment utiliser

Nous fournissons un script pour la formation d'un modèle TinyLlama-1b comme estimateur de valeur d'un modèle LLAMA-2 7B sur l'ensemble de données HH d'Anthropic.

Pour suivre le pipeline décrit dans l'article, nous fournissons une version ajustée fin supervisée de ces modèles:

 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

Citation

@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}
}