KTO Trainer

Resumen

TRL admite el entrenador de Optimización Kahneman-Tversky (KTO) para entrenar modelos de lenguaje, tal como se describe en el artículo KTO: Model Alignment as Prospect Theoretic Optimization de Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, Douwe Kiela.

El resumen del artículo es el siguiente:

La teoría de la perspectiva de Kahneman & Tversky nos dice que los humanos perciben las variables aleatorias de manera sesgada pero bien definida; por ejemplo, los humanos son notoriamente aversos a las pérdidas. Mostramos que los objetivos para alinear LLMs con retroalimentación humana incorporan implícitamente muchos de estos sesgos: el éxito de estos objetivos (p. ej., DPO) sobre la minimización de la entropía cruzada puede atribuirse en parte a que son funciones de pérdida conscientes de lo humano (HALOs). Sin embargo, las funciones de utilidad que estos métodos atribuyen a los humanos aún difieren de las de la literatura sobre la teoría de la perspectiva. Utilizando un modelo de utilidad humana de Kahneman-Tversky, proponemos un HALO que maximiza directamente la utilidad de las generaciones, en lugar de maximizar la verosimilitud logarítmica de las preferencias, como hacen los métodos actuales. Llamamos a este enfoque Optimización de Kahneman-Tversky (KTO), y iguala o supera el rendimiento de los métodos basados en preferencias a escalas de 1B a 30B. Crucialmente, KTO no necesita preferencias: solo una señal binaria de si una salida es deseable o indeseable para una entrada dada. Esto lo hace mucho más fácil de usar en el mundo real, donde los datos de preferencia son escasos y costosos.

El código oficial se puede encontrar en ContextualAI/HALOs.

Este método de post-entrenamiento fue aportado por Kashif Rasul, Younes Belkada, Lewis Tunstall, Pablo Vicente y, posteriormente, refactorizado por Albert Villanova del Moral.

Inicio rápido

Este ejemplo demuestra cómo entrenar un modelo utilizando el método KTO. Utilizamos el modelo Qwen 0.5B como modelo base. Utilizamos los datos de preferencia de KTO Mix 14k. Puede ver los datos del conjunto de datos aquí:

A continuación se muestra el script para entrenar el 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()

Ejecute el script usando el siguiente comando:

accelerate launch train_kto.py

Distribuido en 8 x H100 GPUs, el entrenamiento toma aproximadamente 30 minutos. Puedes verificar el progreso del entrenamiento revisando el gráfico de recompensa. Una tendencia creciente en el margen de recompensa indica que el modelo está mejorando y generando mejores respuestas con el tiempo.

Para ver cómo rinde el modelo entrenado, puedes usar el 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 y formato de datos esperados

KTO requiere un conjunto de datos de preferencias sin pares. Alternativamente, puede proporcionar un conjunto de datos de preferencias con pares (también conocido simplemente como conjunto de datos de preferencias). En este caso, el entrenador lo convertirá automáticamente a un formato sin pares separando las respuestas elegidas y rechazadas, asignando label = True a las finalizaciones elegidas y label = False a las rechazadas.

El KTOTrainer es compatible con los formatos de datos estándar y conversacional. Cuando se proporciona un conjunto de datos conversacional, el entrenador aplica automáticamente la plantilla de chat al conjunto de datos.

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

En teoría, el conjunto de datos debe contener al menos una finalización elegida y una rechazada. Sin embargo, algunos usuarios han ejecutado KTO con éxito utilizando únicamente datos elegidos o únicamente datos rechazados. Si se utilizan solo datos rechazados, es aconsejable adoptar una tasa de aprendizaje conservadora.

Mirando más a fondo el método KTO

La optimización de Kahneman-Tversky (KTO) es un método de entrenamiento diseñado para alinear un modelo de lenguaje utilizando únicamente una señal binaria de si una salida es deseable o indeseable para una entrada dada, en lugar de pares de finalizaciones preferidas/no preferidas. Basándose en la teoría de la perspectiva de Kahneman y Tversky, define una pérdida consciente del humano (HALO) que maximiza directamente la utilidad de las generaciones, ponderando de forma asimétrica los ejemplos deseables e indeseables para reflejar la aversión a las pérdidas de los humanos.

Esta sección detalla cómo funciona KTO en la práctica, cubriendo los pasos clave: preprocesamiento y cálculo de la pérdida.

Preprocesamiento y tokenización

Durante el entrenamiento, se espera que cada ejemplo contenga un prompt, una completion y una etiqueta booleana label que indique si la completion es deseable (True) o indeseable (False). Para más detalles sobre los formatos esperados, consulte Dataset formats. El KTOTrainer tokeniza cada entrada utilizando el tokenizer del modelo.

Cálculo de la pérdida

El término de divergencia KL utilizado por la pérdida KTO se estima emparejando cada prompt con una finalización no coincidente extraída de otra parte del lote. Por esta razón, los tipos de pérdida que estiman el término KL (todos excepto apo_zero_unpaired) requieren un tamaño de lote de entrenamiento por dispositivo mayor que 1 y una estrategia de muestreo secuencial, para que los pares no coincidentes sean estables a lo largo de un lote.

La pérdida utilizada en KTO (Ec. 7 del artículo) ) se define de la siguiente manera: $$\mathcal{L}_{\mathrm{KTO}}(\theta) = \mathbb{E}_{(x,y)}\!\left[ w(y)\Big(1 - v(x, y)\Big) \right]$$

donde el valor $v(x, y)$ es $$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}$$

Aquí $x$ es el prompt, $y$ es la finalización, $\pi_{\theta}$ es el modelo de política que se está entrenando, $\pi_{\mathrm{ref}}$ es el modelo de referencia, $\sigma$ es la función sigmoide, $\beta > 0$ controla la desviación respecto al modelo de referencia, y $\mathrm{KL}$ es el término de divergencia KL estimado. El peso $w(y)$ es desirable_weight para las finalizaciones deseables y undesirable_weight para las indeseables, utilizado para contrarrestar un desequilibrio entre el número de ejemplos deseables e indeseables.

Tipos de pérdida

Métricas registradas

Durante el entrenamiento y la evaluación, registramos las siguientes métricas:

• global_step: El número total de pasos del optimizador realizados hasta el momento.
• epoch: El número de época actual, basado en la iteración del conjunto de datos.
• num_tokens: El número total de tokens procesados hasta el momento.
• loss: La pérdida KTO promedio durante el intervalo de registro actual.
• entropy: La entropía media de la distribución de tokens predicha por el modelo sobre los tokens no enmascarados.
• aux_loss: La pérdida auxiliar de equilibrio de carga de un modelo Mixture-of-Experts, antes de que se escale por router_aux_loss_coef y se sume a la pérdida. Se registra solo cuando el modelo es un modelo MoE y router_aux_loss_coef es distinto de cero.
• kl: La divergencia KL promedio estimada entre el modelo de política y el modelo de referencia, utilizada como punto de referencia en la pérdida KTO.
• learning_rate: la tasa de aprendizaje actual, que puede cambiar dinámicamente si se utiliza un programador.
• grad_norm: La norma L2 de los gradientes, calculada antes del recorte de gradientes.
• logits/chosen: Los valores de logit promedio asignados por el modelo a los tokens de la finalización elegida (deseable).
• logits/rejected: Los valores medios de logit asignados por el modelo a los tokens de la finalización rechazada (no deseable).
• logps/chosen: La probabilidad logarítmica media asignada por el modelo a los tokens de la finalización elegida (deseable).
• logps/rejected: La probabilidad logarítmica media asignada por el modelo a los tokens de la finalización rechazada (no deseable).
• rewards/chosen: La recompensa implícita promedio calculada para la finalización elegida (deseable), calculada como $\beta \log \frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}$.
• rewards/rejected: La recompensa implícita promedio calculada para la finalización rechazada (indeseable), calculada como $\beta \log \frac{\pi_{\theta}(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}$.
• rewards/margins: El margen medio de recompensa implícita entre las finalizaciones elegidas y rechazadas.

Personalización

Compatibilidad y restricciones

Algunas combinaciones de argumentos están restringidas intencionalmente en la implementación actual de KTOTrainer:

• Los adaptadores PEFT que apuntan a lm_head no son compatibles (usa modules_to_save=["lm_head"] en su lugar).
• sync_ref_model=True no se admite al entrenar con modelos PEFT que no conservan un ref_model independiente.
• sync_ref_model=True no se puede combinar con precompute_ref_log_probs=True.
• precompute_ref_log_probs=True no se admite con IterableDataset (entrenamiento o evaluación) ni con conjuntos de datos de visión.
• Los tipos de pérdida que estiman el término de divergencia KL (todos excepto "apo_zero_unpaired") requieren train_sampling_strategy="sequential" y un tamaño de lote de entrenamiento por dispositivo mayor que 1.

Inicialización del modelo

Puede pasar directamente los kwargs del método from_pretrained() a KTOConfig. Por ejemplo, si desea cargar un modelo en una precisión diferente, de manera análoga a

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

puede hacerlo pasando el argumento model_init_kwargs={"dtype": torch.bfloat16} a KTOConfig.

from trl import KTOConfig

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

Tenga en cuenta que se admiten todos los argumentos de palabra clave de from_pretrained().

Entrenar adaptadores con PEFT

Soportamos una integración estrecha con la biblioteca 🤗 PEFT, lo que permite a cualquier usuario entrenar adaptadores de manera conveniente y compartirlos en el Hub, en lugar de entrenar el modelo completo.

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()

También puedes continuar entrenando tu PeftModel. Para ello, primero carga un PeftModel fuera de KTOTrainer y pásalo directamente al trainer sin pasar el argumento peft_config.

Al entrenar adaptadores, normalmente se utiliza una tasa de aprendizaje más alta que en el ajuste fino completo, ya que solo se están aprendiendo parámetros nuevos.

Llamadas a herramientas con KTO

El KTOTrainer admite por completo el ajuste fino de modelos con capacidades de tool calling. En este caso, cada ejemplo del conjunto de datos debe incluir:

• Los mensajes de la conversación (el prompt y la respuesta), incluidos los llamados a herramientas (tool_calls) y las respuestas de las herramientas (mensajes con el rol tool)
• La lista de herramientas disponibles en la columna tools, generalmente proporcionada como esquemas JSON

Para obtener detalles sobre la estructura esperada del conjunto de datos, consulte la sección Formato de conjunto de datos — Llamadas a herramientas.

Entrenamiento de modelos de visión y lenguaje

KTOTrainer admite por completo el entrenamiento de modelos de visión-lenguaje (VLM). Para entrenar un VLM, proporcione un conjunto de datos con una columna image (una imagen por muestra) o una columna images (lista de imágenes por muestra). Para obtener más información sobre la estructura esperada del conjunto de datos, consulte la sección Formato del conjunto de datos — Conjunto de datos de visión.

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()

En el caso de los VLM, la truncación puede eliminar los tokens de imagen, lo que provoca errores durante el entrenamiento. Para evitarlo, establece max_length=None en KTOConfig. Esto permite que el modelo procese la longitud completa de la secuencia sin truncar los tokens de imagen.

Copiado

KTOConfig(max_length=None, ...)

Solo usa max_length cuando hayas verificado que la truncación no eliminará los tokens de imagen para todo el conjunto de datos.

Script de ejemplo

Proporcionamos un script de ejemplo para entrenar un modelo utilizando el método KTO. El script está disponible en trl/scripts/kto.py

Para probar el script de KTO con el modelo Qwen2 0.5B en el conjunto de datos UltraFeedback, ejecute el siguiente 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

Consejos de uso

Para los modelos Mixture of Experts: habilitar la pérdida auxiliar

Los MOEs son más eficientes cuando la carga se distribuye de manera aproximadamente equitativa entre los expertos.
Para garantizar que entrenemos los MOEs de forma similar durante el ajuste de preferencias, es beneficioso añadir la pérdida auxiliar del equilibrador de carga a la pérdida final.

Esta opción se habilita configurando output_router_logits=True en la configuración del modelo (p. ej. MixtralConfig).
Para escalar cuánto contribuye la pérdida auxiliar a la pérdida total, use el hiperparámetro router_aux_loss_coef=... (valor predeterminado: 0.001) en la configuración del modelo.

Recomendaciones de tamaño de lote

Utilice un tamaño de lote por paso de al menos 4 y un tamaño de lote efectivo entre 16 y 128. Incluso si su tamaño de lote efectivo es grande, si su tamaño de lote por paso es deficiente, la estimación de KL en KTO también será deficiente.

Recomendaciones de tasa de aprendizaje

Cada elección de beta tiene una tasa de aprendizaje máxima que puede tolerar antes de que el rendimiento del aprendizaje se degrade. Para la configuración predeterminada de beta = 0.1, la tasa de aprendizaje normalmente no debería superar 1e-6 para la mayoría de los modelos. A medida que beta disminuye, la tasa de aprendizaje también debe reducirse en consecuencia. En general, recomendamos encarecidamente mantener la tasa de aprendizaje entre 5e-7 y 5e-6. Incluso con conjuntos de datos pequeños, desaconsejamos utilizar una tasa de aprendizaje fuera de este rango. En su lugar, elija más épocas para obtener mejores resultados.

Datos desequilibrados

Los parámetros desirable_weight y undesirable_weight de KTOConfig se refieren a los pesos asignados a las pérdidas de los ejemplos deseables/positivos e indeseables/negativos. Por defecto, ambos son 1. Sin embargo, si tienes más de uno u otro tipo, debes aumentar el peso del tipo menos común para que la proporción entre (desirable_weight $\times$ número de positivos) y (undesirable_weight $\times$ número de negativos) esté en el rango 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