Mejorando la Eficiencia de la Memoria en el Aprendizaje por Refuerzo con Retroalimentación Humana

El Aprendizaje por Refuerzo con Retroalimentación Humana (RLHF) ha cambiado la forma en que entrenamos modelos de lenguaje para que reflejen mejor lo que la gente quiere. Pero una parte clave de este proceso, llamada Optimización de Políticas Proximales (PPO), usa mucha memoria. Puede requerir tres veces más memoria que los métodos tradicionales. Esto lo hace difícil de usar para muchas personas. Para solucionar este problema, analizamos de cerca cuánta memoria utilizan estos métodos, qué tan bien funcionan y cuánto tiempo tardan en entrenarse.

Introdujimos un nuevo enfoque llamado Hydra-RLHF. Este enfoque combina diferentes modelos y apaga ciertas partes durante el entrenamiento para ahorrar memoria. Nuestros tests mostraron dos cosas principales: Primero, usar una técnica llamada LoRA durante PPO reduce el uso de memoria por debajo de los métodos tradicionales mientras hace que el modelo se alinee mejor con las preferencias humanas en base a cuatro tests diferentes. Segundo, nuestro enfoque Hydra-PPO reduce el tiempo que toma cada muestra en hasta un 65% sin perder rendimiento. Esto hace que sea más fácil para más personas usar RLHF en su trabajo.

Desde que modelos como ChatGPT, GPT-4 y Llama-2 se hicieron populares, han sorprendido a los usuarios con lo útiles que pueden ser en varias tareas. Un aspecto crucial de su éxito proviene del uso de RLHF para alinear estos modelos con las expectativas humanas. Entrenar modelos de lenguaje grandes les da mucho conocimiento, pero a menudo luchan por aplicar ese conocimiento de manera correcta. Esta discrepancia puede llevar a errores y potenciales daños. Para manejar esto, la alineación ajusta los modelos para que se comporten de maneras esperadas. Ahora es una parte vital para asegurar que estos modelos sean seguros y útiles.

Sin embargo, aunque RLHF mejora esta alineación, también presenta desafíos. Puede ser muy complejo y necesita mucha memoria para ejecutar múltiples modelos al mismo tiempo durante PPO. Dado que RLHF todavía es un área nueva de investigación, hay una gran necesidad de evaluar sus diferentes formas en términos de velocidad y efectividad.

Para satisfacer esa necesidad, nos enfocamos en los pasos de entrenamiento y estructuras del RLHF-PPO estándar. Encontramos grandes oportunidades para reducir costos de memoria y computación al compartir modelos entre los modelos de Referencia, Recompensa, Actor y Crítico.

Nuestras comparaciones mostraron cuánta memoria y tiempo usaron diferentes métodos cuando se probaron en un modelo específico. También presentamos una visión general detallada de cuántos modelos son necesarios en diferentes métodos de PPO, demostrando que nuestro método Hydra-PPO usa menos modelos en memoria, haciéndolo más eficiente.

#Etapas del Proceso RLHF

El método RLHF consiste en tres etapas principales:

1. Ajuste Fino Supervisado (SFT): Esta etapa implica entrenar un modelo de lenguaje en un conjunto de datos para aprender patrones de lenguaje. Hay dos versiones: una donde se entrenan todos los parámetros (Ajuste Fino Completo) y otra donde se usa una técnica específica (LoRA) para reducir la cantidad de parámetros.

2. Modelo de Recompensa (RM): Aquí, modificamos la salida del modelo de lenguaje, enfocándonos en predecir lo que los humanos prefieren basado en un conjunto de pares de prompt y respuesta. Después del entrenamiento, aseguramos que la recompensa dada por este modelo sea estable para ayudar en el paso de PPO.

3. PPO: En esta última etapa, entrenamos tanto a un actor (la parte creativa del modelo) como a un crítico (que evalúa la salida) usando el modelo de recompensa definido anteriormente. Durante este entrenamiento, se están utilizando al menos cuatro modelos, incluyendo un modelo de Referencia congelado para asegurar estabilidad.

#Introduciendo Hydra-RLHF

Proponemos Hydra-RLHF, que modifica el RLHF tradicional para ahorrar memoria durante la fase de PPO mientras mantiene el rendimiento.

1. Hydra-SFT: Este nuevo método de entrenamiento usa un conjunto de datos similar al entrenamiento estándar del modelo de recompensa, optimizando dos tareas simultáneamente. Este método requiere nuevos datos que incluyan comparaciones emparejadas para entrenar efectivamente.

2. LoRA Dinámico: Este enfoque ayuda a ahorrar memoria apagando los pesos de LoRA cuando no son necesarios. Dado que hay dos modelos idénticos (actor y crítico), podemos recuperar uno del otro, reduciendo significativamente el uso de memoria mientras mantenemos el rendimiento intacto.

3. Hydra-PPO: Al usar pesos de LoRA separados para el actor y el crítico, reducimos aún más la necesidad de múltiples modelos en memoria durante PPO.

#Resultados y Comparaciones

Probamos diferentes métodos entre sí para determinar su rendimiento. Encontramos que nuestros nuevos métodos generalmente superan a los métodos tradicionales en promedio. Hydra-PPO mostró mejor alineación que LoRA-PPO, probablemente debido a la mejora en el modelo de recompensa.

En términos de tiempo, Hydra-PPO se volvió más rápido a medida que aumentaba la cantidad de texto. Al aumentar el tamaño del lote de entrenamiento, logramos una disminución sustancial en el tiempo tomado por muestra durante PPO.

También evaluamos otros conjuntos de datos, como StackExchange y Aprendiendo a Resumir, y encontramos patrones interesantes a través de los resultados. Por ejemplo, mientras los modelos estándar a menudo rinden bien, los métodos PPO mostraron mejor recuerdo pero a veces se quedaron atrás en precisión.

#Desafíos con Joined-Hydra-PPO

También probamos Joined-Hydra-PPO, que usa un conjunto de pesos de LoRA tanto para el actor como para el crítico. Este método ahorró algo de memoria, pero su rendimiento no fue tan bueno como el de Hydra-PPO. Creemos que esto proviene de la inestabilidad que surge al combinar los dos modelos en uno.

#Direcciones Futuras

Nuestra investigación señala nuevas vías para mejorar RLHF. Hay una necesidad de equilibrar mejor los conjuntos de datos utilizados para el entrenamiento de SFT y RM. Un desarrollo adicional podría mejorar el rendimiento de métodos como J-Hydra-PPO, así como hacer que otras técnicas de ajuste fino eficientes en parámetros sean más efectivas en entornos RLHF.

#Conclusión

A través de nuestro estudio, hemos demostrado que es posible mejorar la eficiencia de RLHF al ahorrar memoria durante la fase de PPO. Nuestro método Hydra-RLHF combina modelos y ajusta estrategias de entrenamiento para permitir el uso de tamaños de lote más grandes, lo que lleva a procesos de entrenamiento más rápidos y accesibles. Esperamos que nuestros hallazgos alienten una adopción más amplia de RLHF e inspiren mejoras futuras en esta emocionante área de la tecnología.