KTO Trainer

Visão geral

O TRL oferece suporte ao Treinador de Otimização Kahneman-Tversky (KTO) para o treinamento de modelos de linguagem, conforme descrito no artigo KTO: Model Alignment as Prospect Theoretic Optimization de Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, Douwe Kiela.

O resumo do artigo é o seguinte:

A teoria da perspectiva de Kahneman & Tversky nos diz que os humanos percebem variáveis aleatórias de uma maneira enviesada, porém bem definida; por exemplo, os humanos são notoriamente avessos a perdas. Mostramos que os objetivos para alinhar LLMs com feedback humano incorporam implicitamente muitos desses vieses — o sucesso desses objetivos (por exemplo, DPO) em relação à minimização da entropia cruzada pode ser atribuído em parte ao fato de serem funções de perda conscientes dos humanos (HALOs). No entanto, as funções de utilidade que esses métodos atribuem aos humanos ainda diferem daquelas na literatura da teoria da perspectiva. Usando um modelo de utilidade humana de Kahneman-Tversky, propomos um HALO que maximiza diretamente a utilidade das gerações, em vez de maximizar a log-verossimilhança das preferências, como os métodos atuais fazem. Chamamos essa abordagem de Otimização de Kahneman-Tversky (KTO), e ela iguala ou supera o desempenho de métodos baseados em preferências em escalas de 1B a 30B. Crucialmente, o KTO não precisa de preferências — apenas de um sinal binário de se uma saída é desejável ou indesejável para uma entrada dada. Isso o torna muito mais fácil de usar no mundo real, onde dados de preferência são escassos e caros.

O código oficial pode ser encontrado em ContextualAI/HALOs.

Este método de pós-treinamento foi contribuído por Kashif Rasul, Younes Belkada, Lewis Tunstall, Pablo Vicente e, posteriormente, refatorado por Albert Villanova del Moral.

Início rápido

Este exemplo demonstra como treinar um modelo usando o método KTO. Usamos o modelo Qwen 0.5B como modelo base. Usamos os dados de preferência do KTO Mix 14k. Você pode visualizar os dados no dataset aqui:

Abaixo está o script para treinar o modelo:

# train_kto.py
from datasets import load_dataset
from trl import KTOConfig, KTOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
train_dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

training_args = KTOConfig(output_dir="Qwen2-0.5B-KTO")
trainer = KTOTrainer(model=model, args=training_args, processing_class=tokenizer, train_dataset=train_dataset)
trainer.train()

Execute o script usando o seguinte comando:

accelerate launch train_kto.py

Distribuído em 8 x H100 GPUs, o treinamento leva aproximadamente 30 minutos. Você pode verificar o progresso do treinamento analisando o gráfico de recompensa. Uma tendência de aumento na margem de recompensa indica que o modelo está melhorando e gerando respostas melhores ao longo do tempo.

Para ver como o modelo treinado se comporta, você pode usar o Transformers Chat CLI.

$ transformers chat trl-lib/Qwen2-0.5B-KTO
<quentin_gallouedec>:
What is the best programming language?

<trl-lib/Qwen2-0.5B-KTO>:
The best programming language can vary depending on individual preferences, industry-specific requirements, technical skills, and familiarity with the specific use case or task. Here are some widely-used programming languages that have been noted as popular and widely used:

Here are some other factors to consider when choosing a programming language for a project:

 1 JavaScript: JavaScript is at the heart of the web and can be used for building web applications, APIs, and interactive front-end applications like frameworks like React and Angular. It's similar to C, C++, and F# in syntax structure and is accessible and easy to learn, making it a popular choice for beginners and professionals alike.
 2 Java: Known for its object-oriented programming (OOP) and support for Java 8 and .NET, Java is used for developing enterprise-level software applications, high-performance games, as well as mobile apps, game development, and desktop applications.
 3 C++: Known for its flexibility and scalability, C++ offers comprehensive object-oriented programming and is a popular choice for high-performance computing and other technical fields. It's a powerful platform for building real-world applications and games at scale.
 4 Python: Developed by Guido van Rossum in 1991, Python is a high-level, interpreted, and dynamically typed language known for its simplicity, readability, and versatility.

Tipo e formato de dataset esperados

O KTO requer um conjunto de dados de preferência não pareada. Alternativamente, você pode fornecer um conjunto de dados de preferência pareada (também conhecido simplesmente como conjunto de dados de preferências). Nesse caso, o treinador o converterá automaticamente para o formato não pareado, separando as respostas escolhidas e rejeitadas, atribuindo label = True às conclusões escolhidas e label = False às rejeitadas.

O KTOTrainer é compatível com os formatos de dataset padrão e conversacional. Quando fornecido com um dataset conversacional, o trainer aplica automaticamente o chat template ao dataset.

# Standard format
unpaired_preference_example = {"prompt": "The sky is", "completion": " blue.", "label": True}

# Conversational format
unpaired_preference_example = {"prompt": [{"role": "user", "content": "What color is the sky?"}],
                               "completion": [{"role": "assistant", "content": "It is blue."}],
                               "label": True}

Em teoria, o conjunto de dados deve conter pelo menos uma conclusão escolhida e uma rejeitada. No entanto, alguns usuários executaram o KTO com sucesso usando apenas dados escolhidos ou apenas rejeitados. Ao usar apenas dados rejeitados, é aconselhável adotar uma taxa de aprendizado conservadora.

Aprofundando-se no método KTO

A Otimização Kahneman-Tversky (KTO) é um método de treinamento projetado para alinhar um modelo de linguagem usando apenas um sinal binário de se uma saída é desejável ou indesejável para uma entrada dada, em vez de pares de conclusões preferidas/despreferidas. Baseando-se na teoria da perspectiva de Kahneman e Tversky, ela define uma perda consciente do humano (HALO) que maximiza diretamente a utilidade das gerações, ponderando exemplos desejáveis e indesejáveis de forma assimétrica para refletir a aversão à perda humana.

Esta seção detalha como o KTO funciona na prática, abordando as etapas principais: pré-processamento e cálculo da perda.

Pré-processamento e tokenização

Durante o treinamento, cada exemplo deve conter um prompt, uma completion e um label booleano indicando se a completion é desejável (True) ou indesejável (False). Para mais detalhes sobre os formatos esperados, consulte Dataset formats. O KTOTrainer tokeniza cada entrada usando o tokenizer do modelo.

Calculando a perda

O termo de divergência KL usado pela perda KTO é estimado ao emparelhar cada prompt com uma conclusão não correspondente retirada de outra parte do lote. Por esse motivo, os tipos de perda que estimam o termo KL (todos, exceto apo_zero_unpaired) exigem um tamanho de lote de treinamento por dispositivo maior que 1 e uma estratégia de amostragem sequencial, para que os pares não correspondentes sejam estáveis ao longo de um lote.

A perda usada no KTO (Eq. 7 do artigo) é definida da seguinte forma: $$\mathcal{L}_{\mathrm{KTO}}(\theta) = \mathbb{E}_{(x,y)}\!\left[ w(y)\Big(1 - v(x, y)\Big) \right]$$

onde o valor $v(x, y)$ é $$v(x, y) = \begin{cases} \sigma\!\left(\beta\big(\log\frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)} - \mathrm{KL}\big)\right) & \text{if } y \text{ is desirable} \\ \sigma\!\left(\beta\big(\mathrm{KL} - \log\frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}\big)\right) & \text{if } y \text{ is undesirable} \end{cases}$$

Aqui $x$ é o prompt, $y$ é a conclusão, $\pi_{\theta}$ é o modelo de política em treinamento, $\pi_{\mathrm{ref}}$ é o modelo de referência, $\sigma$ é a função sigmoide, $\beta > 0$ controla o desvio em relação ao modelo de referência, e $\mathrm{KL}$ é o termo de divergência KL estimado. O peso $w(y)$ é desirable_weight para conclusões desejáveis e undesirable_weight para as indesejáveis, usado para compensar um desequilíbrio entre o número de exemplos desejáveis e indesejáveis.

Tipos de Perda

Métricas registradas

Durante o treinamento e a avaliação, registramos as seguintes métricas:

• global_step: O número total de etapas do otimizador realizadas até o momento.
• epoch: O número do epoch atual, com base na iteração do conjunto de dados.
• num_tokens: O número total de tokens processados até o momento.
• loss: A perda média de KTO no intervalo de registro atual.
• entropy: A entropia média da distribuição de tokens prevista pelo modelo para tokens não mascarados.
• aux_loss: A perda auxiliar de balanceamento de carga de um modelo Mixture-of-Experts, antes de ser escalada por router_aux_loss_coef e adicionada à perda. Registrada apenas quando o modelo é um modelo MoE e router_aux_loss_coef é diferente de zero.
• kl: A divergência de KL média estimada entre o modelo de política e o modelo de referência, usada como ponto de referência na perda KTO.
• learning_rate: a taxa de aprendizado atual, que pode mudar dinamicamente se um agendador for usado.
• grad_norm: A norma L2 dos gradientes, calculada antes do recorte de gradientes.
• logits/chosen: Os valores médios de logit atribuídos pelo modelo aos tokens na conclusão escolhida (desejável).
• logits/rejected: Os valores médios de logit atribuídos pelo modelo aos tokens na conclusão rejeitada (indesejável).
• logps/chosen: A média da log-probabilidade atribuída pelo modelo aos tokens na conclusão escolhida (desejável).
• logps/rejected: A média da log-probabilidade atribuída pelo modelo aos tokens na conclusão rejeitada (indesejável).
• rewards/chosen: A recompensa implícita média calculada para a conclusão escolhida (desejável), calculada como $\beta \log \frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}$.
• rewards/rejected: A recompensa implícita média calculada para a conclusão rejeitada (indesejável), calculada como $\beta \log \frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}$.
• rewards/margins: A margem média de recompensa implícita entre as conclusões escolhidas e rejeitadas.

Personalização

Compatibilidade e restrições

Algumas combinações de argumentos são intencionalmente restritas na implementação atual do KTOTrainer:

• Adaptadores PEFT que visam lm_head não são suportados (use modules_to_save=["lm_head"] em vez disso).
• sync_ref_model=True não é suportado ao treinar com modelos PEFT que não mantêm um ref_model independente.
• sync_ref_model=True não pode ser combinado com precompute_ref_log_probs=True.
• precompute_ref_log_probs=True não é suportado com IterableDataset (treino ou avaliação) nem com conjuntos de dados de visão.
• Os tipos de perda que estimam o termo de divergência KL (todos, exceto "apo_zero_unpaired") exigem train_sampling_strategy="sequential" e um tamanho de lote de treinamento por dispositivo maior que 1.

Inicialização do modelo

Você pode passar diretamente os kwargs do método from_pretrained() para o KTOConfig. Por exemplo, se você quiser carregar um modelo em uma precisão diferente, análogo a

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct", dtype=torch.bfloat16)

você pode fazer isso passando o argumento model_init_kwargs={"dtype": torch.bfloat16} para o KTOConfig.

from trl import KTOConfig

training_args = KTOConfig(
    model_init_kwargs={"dtype": torch.bfloat16},
)

Observe que todos os argumentos de palavra-chave de from_pretrained() são suportados.

Treine adaptadores com PEFT

Suportamos integração estreita com a biblioteca 🤗 PEFT, permitindo que qualquer usuário treine adaptadores de forma conveniente e os compartilhe no Hub, em vez de treinar o modelo inteiro.

from datasets import load_dataset
from trl import KTOTrainer
from peft import LoraConfig

dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

trainer = KTOTrainer(
    "Qwen/Qwen2-0.5B-Instruct",
    train_dataset=dataset,
    peft_config=LoraConfig(),
)

trainer.train()

Você também pode continuar o treinamento do seu PeftModel. Para isso, primeiro carregue um PeftModel fora do KTOTrainer e passe-o diretamente para o trainer sem passar o argumento peft_config.

Ao treinar adaptadores, você geralmente usa uma taxa de aprendizado mais alta do que no ajuste fino completo, pois apenas novos parâmetros estão sendo aprendidos.

Tool Calling com KTO

O KTOTrainer oferece suporte completo ao ajuste fino de modelos com capacidades de tool calling. Nesse caso, cada exemplo do conjunto de dados deve incluir:

• As mensagens da conversa (prompt e completion), incluindo qualquer chamada de ferramenta (tool_calls) e respostas de ferramenta (mensagens com o papel tool)
• A lista de ferramentas disponíveis na coluna tools, geralmente fornecida como esquemas JSON

Para mais detalhes sobre a estrutura esperada do conjunto de dados, consulte a seção Formato do Conjunto de Dados — Chamada de Ferramenta.

Treinando Modelos de Visão e Linguagem

KTOTrainer oferece suporte completo ao treinamento de Vision-Language Models (VLMs). Para treinar um VLM, forneça um conjunto de dados com uma coluna image (uma única imagem por amostra) ou uma coluna images (lista de imagens por amostra). Para mais informações sobre a estrutura esperada do conjunto de dados, consulte a seção Dataset Format — Vision Dataset.

from trl import KTOConfig, KTOTrainer
from datasets import load_dataset

trainer = KTOTrainer(
    model="Qwen/Qwen2.5-VL-3B-Instruct",
    args=KTOConfig(max_length=None),
    train_dataset=load_dataset("trl-internal-testing/zen-image", "conversational_unpaired_preference", split="train"),
)
trainer.train()

Para VLMs, a truncagem pode remover tokens de imagem, levando a erros durante o treinamento. Para evitar isso, defina max_length=None no KTOConfig. Isso permite que o modelo processe o comprimento total da sequência sem truncar tokens de imagem.

Copiado

KTOConfig(max_length=None, ...)

Use max_length apenas quando você tiver verificado que a truncagem não removerá os tokens de imagem para todo o conjunto de dados.

Script de exemplo

Fornecemos um script de exemplo para treinar um modelo usando o método KTO. O script está disponível em trl/scripts/kto.py

Para testar o script KTO com o modelo Qwen2 0.5B no conjunto de dados UltraFeedback, execute o seguinte comando:

accelerate launch trl/scripts/kto.py \
    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
    --dataset_name trl-lib/kto-mix-14k \
    --num_train_epochs 1 \
    --output_dir Qwen2-0.5B-KTO

Dicas de uso

Para Modelos Mixture of Experts: Habilitando a perda auxiliar

MOEs são mais eficientes quando a carga é distribuída de forma aproximadamente igual entre os especialistas.
Para garantir que treinemos MOEs de maneira semelhante durante o ajuste de preferências, é benéfico adicionar a perda auxiliar do balanceador de carga à perda final.

Esta opção é ativada definindo output_router_logits=True na configuração do modelo (por exemplo, MixtralConfig).
Para ajustar a contribuição da perda auxiliar para a perda total, use o hiperparâmetro router_aux_loss_coef=... (padrão: 0.001) na configuração do modelo.

Recomendações de tamanho de lote

Use um tamanho de lote por etapa de pelo menos 4 e um tamanho de lote efetivo entre 16 e 128. Mesmo que o seu tamanho de lote efetivo seja grande, se o seu tamanho de lote por etapa for inadequado, a estimativa de KL no KTO será inadequada.

Recomendações de taxa de aprendizado

Cada escolha de beta tem uma taxa de aprendizado máxima que pode tolerar antes que o desempenho do aprendizado se degrade. Para a configuração padrão de beta = 0.1, a taxa de aprendizado normalmente não deve exceder 1e-6 para a maioria dos modelos. À medida que beta diminui, a taxa de aprendizado também deve ser reduzida de forma correspondente. Em geral, recomendamos fortemente manter a taxa de aprendizado entre 5e-7 e 5e-6. Mesmo com conjuntos de dados pequenos, desaconselhamos o uso de uma taxa de aprendizado fora dessa faixa. Em vez disso, opte por mais épocas para obter melhores resultados.

Dados desbalanceados

Os parâmetros desirable_weight e undesirable_weight da KTOConfig referem-se aos pesos aplicados às perdas para exemplos desejáveis/positivos e indesejáveis/negativos. Por padrão, ambos são 1. No entanto, se você tiver mais de um ou do outro, deve aumentar o peso do tipo menos comum para que a proporção entre (desirable_weight $\times$ número de positivos) e (undesirable_weight $\times$ número de negativos) esteja na faixa de 1:1 a 4:3.

KTOTrainer

>>> from trl import KTOTrainer
>>> from datasets import load_dataset

>>> dataset = load_dataset("trl-lib/kto-mix-14k", split="train")

>>> trainer = KTOTrainer(
...     model="Qwen/Qwen2.5-0.5B-Instruct",
...     train_dataset=dataset,
... )
>>> trainer.train()

KTOConfig