El entrenamiento permite reutilizar cachés de inferencia en modelos de lenguaje en bucle
BeNice! Closing Looping's Inference-Memory Gap via KV-Sharing · arXiv:2610.31633
En palabras sencillas
Algunos modelos de lenguaje repiten las mismas capas para hacer más cálculos, pero normalmente guardan una caché de atención distinta en cada repetición. Este artículo presenta un método de entrenamiento que ayuda a las primeras repeticiones a usar la caché de la última, lo que reduce la necesidad de conservar varias cachés. El método primero entrena el modelo con distintas cantidades de repeticiones y luego enseña a la primera a trabajar con la caché de la última.
En las pruebas, esto a menudo mantuvo la calidad, aunque entrenar con distintas profundidades no bastó para corregir todos los modelos.
La investigación
Los modelos de lenguaje en bucle pueden añadir cálculos al repetir las mismas capas sin añadir parámetros, pero su caché de inferencia crece con cada repetición. ¿Pueden mantener esa forma de ampliar el cálculo usando menos memoria de caché y sin perjudicar la calidad lingüística del modelo?
Los autores combinan el preentrenamiento con profundidad estocástica y BeNice!, que enseña a la primera recurrencia de un modelo a usar la caché de la recurrencia final sin perder calidad lingüística.
Primero, los autores entrenaron modelos con profundidad de recurrencia estocástica; después, añadieron BeNice! para enseñar a la primera recurrencia a usar la caché de la recurrencia final, conservando a la vez la pérdida estándar de modelado del lenguaje.
Compararon evaluaciones con caché compartida y no compartida, también frente al ajuste fino estándar.
En un modelo público de 770M, la perplejidad de validación con caché compartida fue de 11.86 tras BeNice! y de 21.61 tras el ajuste fino estándar, mientras que las puntuaciones con caché no compartida fueron casi iguales.
Esto respalda la afirmación de que BeNice! puede mantener la calidad al compartir cachés en este modelo evaluado.
Los resultados de velocidad de inferencia y de tiempo por solicitud provienen de reproducciones simuladas de solicitudes derivadas de producción con modelos proxy, no de un servicio en funcionamiento.
La prueba de convergencia también depende de que la dinámica recurrente sea contractiva y de que las proyecciones de claves y valores sean Lipschitz —es decir, que sus salidas no cambien de forma desproporcionada ante cambios en las entradas—, condiciones que el artículo no garantiza para todas las arquitecturas en bucle.
Qué muestran las pruebas
- El problema de memoria
Los modelos de lenguaje en bucle repiten las mismas capas, lo que permite aumentar la profundidad de recurrencia sin añadir parámetros al modelo. Sin embargo, durante la inferencia convencional se mantiene una caché de claves y valores separada para cada recurrencia, así que la memoria de caché crece con el número de repeticiones.
El artículo pregunta cómo pueden estos modelos reutilizar la caché de la recurrencia final en las anteriores sin perder calidad.
- Un primer paso de entrenamiento
El método de dos etapas de los autores combina el preentrenamiento con profundidad estocástica y el posentrenamiento BeNice!.
En su comparación de modelos Parcae 770M del mismo tamaño, el entrenamiento con profundidad estocástica produjo una diferencia de +0.01 puntos en la perplejidad de validación entre caché compartida y no compartida, mientras que el entrenamiento con profundidad fija produjo +3.10 puntos. Los modelos compartían la misma profundidad media de recurrencia, los mismos datos y la misma profundidad de retropropagación; el resultado relaciona la profundidad estocástica con la posibilidad de compartir cachés en esta configuración.
- Por qué hacía falta otro paso
El preentrenamiento con profundidad estocástica no resolvió el problema en todos los modelos evaluados. Con ocho recurrencias sobre datos de validación de ClimbMix, las representaciones ocultas finales del modelo Parcae 1.3B diferían en un 177.5% relativo medio entre las evaluaciones con caché compartida y no compartida.
En este ejemplo, esa diferencia indica que ambas evaluaciones se acercaron a puntos distintos y ayuda a explicar por qué hacía falta una corrección adicional.
- Corregir la discrepancia de caché
BeNice! entrena la primera recurrencia para que use la caché de la recurrencia final, a la vez que mantiene la pérdida habitual de modelado del lenguaje.
En cuatro clases de tamaño de Parcae recién entrenadas, las puntuaciones con caché compartida mejoraron hasta 8.35 puntos en Core, un conjunto de 22 tareas de lenguaje, y hasta 4.56 puntos en Core-Extended, un conjunto de 46 tareas.
Son mejoras máximas entre clases de tamaño. En una prueba aparte con un checkpoint público de 770M, BeNice! también acercó la perplejidad de validación con caché compartida al resultado con caché no compartida.
- Un resultado acotado a la carga evaluada
En una reproducción simulada de solicitudes derivadas de producción a dos solicitudes por segundo, compartir la caché redujo de 139.2 a 23.2 segundos la espera hasta el primer token de salida en el percentil 90 y aumentó el rendimiento un 23.6% frente al uso de cachés separadas.
Esto respalda un beneficio de rendimiento en esa condición simulada de alta carga. No establece un resultado en un servicio en funcionamiento.
Detalles
- Autores del paper original
- Hayden Prairie, Jacob M. Ryan, Taylor Berg-Kirkpatrick, Daniel Y. Fu
- Enviado a arXiv
- Licencia
- CC BY 4.0