WTA-CRS: Un Nuevo Enfoque para Ajustar Modelos de Lenguaje Grandes

El tamaño de los modelos de lenguaje ha crecido rápidamente, lo que hace más difícil ajustar estos modelos debido a sus altas necesidades de memoria. La mayoría de los trabajos anteriores se han centrado en reducir el número de parámetros que se pueden entrenar en el modelo. Sin embargo, aunque los parámetros del modelo sí afectan el uso de memoria, hay un problema mayor con el almacenamiento de mapas de características o Activaciones, que son necesarios para calcular los gradientes.

En el entrenamiento de redes neuronales, generalmente se utiliza el descenso de gradiente estocástico. Este método puede lidiar con gradientes ruidosos siempre que el estimador de gradiente sea imparcial y tenga una varianza razonable. Basándonos en esta idea, presentamos un nuevo tipo de estimador imparcial para la Multiplicación de matrices que reduce el uso de memoria. Nuestro enfoque solo requiere almacenar un subconjunto de activaciones para calcular gradientes. Mostramos que este método conduce a un menor uso de memoria y mantiene niveles de rendimiento similares a los métodos existentes.

#Antecedentes sobre Modelos de Lenguaje

Los modelos de lenguaje con una estructura de transformador han tenido un gran éxito en muchas tareas relacionadas con el procesamiento del lenguaje natural. Estos modelos suelen aprender de grandes cantidades de texto para crear representaciones generales que se pueden ajustar para tareas específicas.

Investigaciones recientes muestran que aumentar el número de parámetros en modelos de lenguaje preentrenados puede mejorar significativamente su rendimiento. Como resultado, hay una creciente necesidad de ajustar estos grandes modelos, que a menudo tienen miles de millones de parámetros, para una variedad de tareas.

Desafortunadamente, hay una gran brecha entre la memoria que estos modelos necesitan y la capacidad del hardware actual, especialmente las GPU. Por ejemplo, incluso una GPU poderosa con 24GB de memoria tiene problemas para ajustar algunos modelos, como el T5-3B, si solo se usa una muestra en un lote. Esto hace que sea imposible ajustar estos grandes modelos en una sola GPU sin métodos adicionales.

Aunque el ajuste paralelo de modelos podría funcionar, a menudo no es práctico debido a los costos y la complejidad de coordinar múltiples GPU. Por lo tanto, es crucial adaptar estos modelos para que funcionen de manera efectiva en configuraciones de una sola GPU.

#Desafíos de Memoria en el Ajuste

Para abordar estos problemas de memoria, se han propuesto varios métodos que ofrecen un ajuste eficiente en parámetros. Por ejemplo, los adaptadores insertan módulos pequeños en bloques de transformadores y solo actualizan estos, manteniendo el resto de los parámetros fijos. Otro enfoque, el ajuste de indicaciones, agrega un pequeño vector a las incrustaciones de entrada que se actualiza durante el ajuste. LoRA, que significa Adaptación de Bajo Rango, implica agregar matrices entrenables en los bloques de transformadores mientras se mantienen congelados los otros parámetros.

Si bien estos métodos eficientes en parámetros ayudan a reducir el uso de memoria de los estados del optimizador, a menudo no reducen significativamente la memoria requerida para almacenar activaciones, que sigue siendo la principal preocupación durante el entrenamiento.

Un método poco explorado para abordar este problema implica usar un enfoque eficiente en memoria para la multiplicación de matrices. Dado que los transformadores dependen principalmente de capas lineales, proponemos un método para submuestrear activaciones y usar estas cantidades reducidas para realizar cálculos.

#El Concepto de Muestreo Columna-Fila

La idea detrás de nuestro método, llamado Muestreo Columna-Fila (CRS), es muestrear partes de tensores en espacios de menor dimensión y luego realizar operaciones regulares usando estas representaciones más pequeñas. Esto es particularmente relevante al trabajar con operaciones lineales entre matrices.

Para lograr esto, primero muestreamos algunos pares columna-fila de acuerdo con una distribución predeterminada. Específicamente, tomamos columnas de una matriz y filas correspondientes de otra en función de estos pares seleccionados, lo que conduce a un producto estimado de las dos matrices.

Este método nos permite almacenar solo las sub-matrices más pequeñas en la memoria de la GPU, reduciendo significativamente la memoria necesaria para las activaciones durante el entrenamiento.

Al entrenar modelos de transformadores, es común utilizar optimizadores estocásticos de primer orden, como Adam. Estos modelos pueden trabajar con gradientes ligeramente ruidosos siempre que sean imparciales y tengan una varianza manejable.

Dada esta comprensión, nos preguntamos por qué invertir en gradientes precisos si ya utilizamos optimización estocástica. Por lo tanto, buscamos una manera de derivar gradientes imparciales usando multiplicación de matrices aproximada que reduzca el uso de memoria.

Este método viene con un potencial compromiso en Precisión, ya que reducir la memoria puede llevar a más varianza en las salidas. El desafío, entonces, se convierte en integrar nuestra multiplicación aproximada en los transformadores sin introducir demasiada varianza en los gradientes.

#Introduciendo Muestreo Columna-Fila Ganador

Nuestro método propuesto es un nuevo estimador para la multiplicación de matrices llamado Muestreo Columna-Fila Ganador (WTA-CRS). Este método reduce la varianza de los estimadores al centrarse en áreas de la distribución que tienen alta probabilidad.

La belleza de WTA-CRS radica en su capacidad para reemplazar las operaciones lineales en transformadores sin problemas. Esto hace posible ofrecer un gradiente de peso imparcial mientras se utiliza menos memoria.

Proporcionamos tanto pruebas teóricas como experimentales de que nuestro nuevo estimador funciona mejor que los métodos existentes en términos de eficiencia de memoria y precisión.

Al intercambiar las operaciones lineales originales por WTA-CRS, podemos reducir las necesidades de memoria máxima hasta 2.7 veces con poco o ningún descenso en precisión.

#Implementación en Transformadores

Para profundizar, es esencial analizar cómo los transformadores utilizan memoria. Durante el entrenamiento, la retropropagación involucra dos fases: un pase hacia adelante y un pase hacia atrás. Las operaciones lineales en los transformadores son cruciales para este proceso.

Cuando consideramos la fase de retroceso, las activaciones utilizadas deben almacenarse en memoria para calcular gradientes. Estudios anteriores muestran que estas activaciones ocupan una parte significativa de la memoria total cuando se trata de entrenar grandes modelos de transformadores.

Nuestro objetivo es reemplazar partes del almacenamiento de activaciones con nuestro nuevo método de muestreo. Si podemos sustituir efectivamente el uso de memoria requerido para las activaciones de esta manera, podemos reducir drásticamente la memoria necesaria, haciendo que el entrenamiento a gran escala sea factible en hardware estándar.

#Pruebas y Resultados

Para validar nuestro nuevo método, realizamos una serie de pruebas. Nos centramos en dos preguntas clave: ¿qué tan efectivo es nuestro método en términos de precisión y uso de memoria, y cómo se desempeña bajo diferentes hiperparámetros?

1. Precisión vs. Uso de Memoria: Nuestros hallazgos muestran que nuestro método logra un buen equilibrio entre ahorro de memoria y precisión. Específicamente, cuando se usa junto con otras técnicas de ahorro de memoria, nuestro método puede reducir significativamente el uso máximo de memoria mientras mantiene niveles de precisión similares a los de entrenar completamente el modelo.

2. Sensibilidad a Hiperparámetros: También exploramos cómo el presupuesto de pares columna-fila afecta el rendimiento. En general, encontramos que los modelos más grandes son más flexibles en términos de ahorros de memoria debido a sus redundancias inherentes.

En pruebas adicionales, comparamos WTA-CRS contra métodos estándar. Nuestros resultados demuestran consistentemente que WTA-CRS supera tanto a métodos deterministas como a CRS tradicionales en términos de precisión y eficiencia de memoria.

#Abordando la Velocidad de Ajuste

Notamos que nuestro nuevo método no agrega parámetros adicionales al modelo, lo que significa que no debería afectar la velocidad de inferencia. Sin embargo, también exploramos cómo afecta la velocidad de ajuste en relación con el uso de memoria.

Con nuestro método reduciendo la memoria requerida para activaciones, esto permite tamaños de lote más grandes durante el entrenamiento. Los tamaños de lote más grandes, a su vez, mejoran la utilización de la GPU y pueden llevar a velocidades de entrenamiento más rápidas.

Si bien la implementación puede ralentizar ligeramente el cálculo, los beneficios generales de tamaños de lote más grandes pueden compensar esto, llevando a mejoras en la velocidad en la práctica.

#Conclusión

En resumen, presentamos WTA-CRS, un nuevo estimador imparcial para la multiplicación de matrices que puede reducir la varianza y las necesidades de memoria mientras preserva el rendimiento del modelo. Nuestros experimentos revelan su valor para facilitar el uso de grandes modelos de lenguaje en hardware estándar, mejorando enormemente su accesibilidad para diversas tareas de procesamiento del lenguaje natural.

A medida que el ajuste de modelos de lenguaje se vuelve cada vez más crítico, especialmente aquellos con miles de millones de parámetros, estrategias efectivas como WTA-CRS muestran un gran potencial para hacer esto factible dentro de las limitaciones del hardware actual. A medida que trabajamos para optimizar aún más nuestra implementación, creemos que esta investigación allana el camino para un ajuste más eficiente de modelos de lenguaje en el futuro.