Value Augmented Sampling

Amostragem aumentada por valor

Este repositório contém a implementação oficial do nosso artigo, amostragem aumentada por valor para alinhar e personalizar modelos de linguagem.

Resumo

Alinhar grandes modelos de linguagem (LLMS) para atender a diferentes preferências humanas, aprender novas habilidades e desaprender o comportamento prejudicial é um problema importante. Métodos baseados em pesquisa, como a Best of-N ou Monte-Carlo Tree Search, são executivos, mas impraticáveis ​​para a adaptação do LLM devido ao seu alto custo de inferência. Por outro lado, o uso de aprendizado de reforço (RL) para adaptação é computacionalmente eficiente, mas tem um desempenho pior devido aos desafios de otimização no treinamento da função de valor e da política. Apresentamos uma nova estrutura para otimização de recompensa, amostragem aumentada por valor (VAS), que pode maximizar diferentes funções de recompensa usando dados amostrados apenas do LLM inicial e congelado. O VAS resolve a política ideal de maximização de recompensa sem conter a política e a função do valor, tornando a otimização estável, superando linhas de base estabelecidas, como PPO e DPO, em benchmarks padrão e alcançando resultados comparáveis ​​para o melhor dos 128 com o melhor de 128 com o Best of 128 com menor custo de inferência. Ao contrário dos métodos RL existentes que exigem alterar os pesos do LLM, o VAS não requer acesso aos pesos do LLM pré-treinado. Assim, ele pode até adaptar o LLMS (por exemplo, ChatGPT), que estão disponíveis apenas como APIs. Além disso, nosso algoritmo desbloqueia a nova capacidade de compor várias recompensas e controlar a extensão de cada um durante o tempo de implantação, abrindo a estrada à frente para o futuro do LLMS alinhado e personalizado.

Instalação

Instale nossa versão personalizada do trl :

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

Clone e instale a base de código:

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

Como usar

Fornecemos um script para treinar um modelo Tinyllama-1b como estimador de valor de um modelo LLAMA-2 7B no conjunto de dados HH da Anthropic.

Para seguir o oleoduto descrito no artigo, fornecemos uma versão fina e afinada desses 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

Citação

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