Por Adán Zewe Cuando una conversación entre humanos e IA implica muchas rondas de diálogo continuo, los poderosos modelos de aprendizaje automático de lenguajes grandes que impulsan los chatbots como ChatGPT a veces comienzan a colapsar, lo que hace que el rendimiento de los bots se deteriore rápidamente. Un equipo de investigadores del MIT y de otros lugares ha identificado una causa sorprendente de este problema y ha desarrollado una solución sencilla que permite a un chatbot mantener una conversación ininterrumpida sin bloquearse ni ralentizarse. Su método implica un ajuste en la caché de valores-clave (que es como una memoria de conversación) en el centro de muchos modelos de lenguaje grandes. En algunos métodos, cuando este caché necesita contener más información de la que tiene capacidad, los primeros datos se eliminan. Esto puede hacer que el modelo falle. Al garantizar que estos primeros puntos de datos permanezcan en la memoria, el método de los investigadores permite que un chatbot siga charlando sin importar cuánto dure la conversación. El método, llamado StreamingLLM, permite que un modelo siga siendo eficiente incluso cuando una conversación se extiende por más de 4 millones de palabras. En comparación con otro método que evita fallas al recalcular constantemente parte de las conversaciones pasadas, StreamingLLM funcionó más de 22 veces más rápido. Esto podría permitir que un chatbot mantenga largas conversaciones durante la jornada laboral sin necesidad de reiniciarlo continuamente, lo que permitiría asistentes de IA eficientes para tareas como redacción, edición o generación de código. “Ahora, con este método, podemos implementar de manera persistente estos grandes modelos de lenguaje. Al crear un chatbot con el que siempre podamos chatear y que siempre pueda respondernos en función de nuestras conversaciones recientes, podríamos utilizar estos chatbots en algunas aplicaciones nuevas”, dice Guangxuan Xiao, un estudiante graduado en ingeniería eléctrica e informática (EECS). y autor principal de un artículo sobre StreamingLLM. Los coautores de Xiao incluyen a su asesor, Song Han, profesor asociado en EECS, miembro del MIT-IBM Watson AI Lab y científico distinguido de NVIDIA; así como Yuandong Tian, ​​científico investigador de Meta AI; Beidi Chen, profesora asistente de la Universidad Carnegie Mellon; y el autor principal Mike Lewis, científico investigador de Meta AI. El trabajo será presentado en la Conferencia Internacional sobre Representaciones del Aprendizaje. Un fenómeno desconcertante Los modelos de lenguaje grandes codifican datos, como palabras en una consulta de usuario, en representaciones llamadas tokens. Muchos modelos emplean lo que se conoce como mecanismo de atención que utiliza estos tokens para generar texto nuevo. Por lo general, un chatbot de IA escribe texto nuevo basándose en el texto que acaba de ver, por lo que almacena tokens recientes en la memoria, llamado KV Cache, para usarlos más adelante. El mecanismo de atención crea una cuadrícula que incluye todos los tokens en el caché, un "mapa de atención" que traza la fuerte relación de cada token o palabra entre sí. Comprender estas relaciones es una característica que permite que los modelos de lenguaje grandes generen texto similar al humano. Pero cuando la memoria caché se vuelve muy grande, el mapa de atención puede volverse aún más masivo, lo que ralentiza el cálculo. Además, si la codificación de contenido requiere más tokens de los que puede contener la caché, el rendimiento del modelo disminuye. Por ejemplo, un modelo popular puede almacenar 4.096 tokens, pero hay alrededor de 10.000 tokens en un artículo académico. Para solucionar estos problemas, los investigadores emplean un "caché deslizante" que elimina los tokens más antiguos para agregar tokens nuevos. Sin embargo, el rendimiento del modelo a menudo cae en picado tan pronto como se elimina el primer token, lo que reduce rápidamente la calidad de las palabras recién generadas. En este nuevo artículo, los investigadores se dieron cuenta de que si mantienen el primer token en la caché deslizante, el modelo mantendrá su rendimiento incluso cuando se exceda el tamaño de la caché. Pero esto no tenía ningún sentido. La primera palabra de una novela probablemente no tenga nada que ver con la última palabra, entonces, ¿por qué la primera palabra sería tan importante para que el modelo genere la palabra más nueva? En su nuevo artículo, los investigadores también descubrieron la causa de este fenómeno. La atención se hunde Algunos modelos utilizan una operación Softmax en su mecanismo de atención, que asigna una puntuación a cada token que representa cuánto se relaciona entre sí. La operación Softmax requiere que todos los puntajes de atención sumen 1. Dado que la mayoría de los tokens no están fuertemente relacionados, sus puntajes de atención son muy bajos. El modelo descarta cualquier puntuación de atención restante en el primer token. Los investigadores llaman a este primer token un "sumidero de atención". “Necesitamos un receptor de atención, y el modelo decide usar el primer token como receptor de atención porque es globalmente visible; todos los demás tokens pueden verlo. Descubrimos que siempre debemos mantener el receptor de atención en la memoria caché para mantener la dinámica del modelo”, afirma Han. Al crear StreamingLLM, los investigadores descubrieron que tener cuatro tokens de sumidero de atención al comienzo del caché deslizante conduce a un rendimiento óptimo. También descubrieron que la codificación posicional de cada token debe permanecer igual, incluso cuando se agregan nuevos tokens y otros se eliminan. Si se elimina el token 5, el token 6 debe permanecer codificado como 6, aunque ahora sea el quinto token en la caché. Al combinar estas dos ideas, permitieron a StreamingLLM mantener una conversación continua y al mismo tiempo superar un método popular que utiliza recálculo. Por ejemplo, cuando el caché tiene 256 tokens, el método de recálculo tarda 63 milisegundos en decodificar un nuevo token, mientras que StreamingLLM tarda 31 milisegundos. Sin embargo, si el tamaño de la caché aumenta a 4096 tokens, el recálculo requiere 1411 milisegundos para un nuevo token, mientras que StreamingLLM necesita solo 65 milisegundos. "El enfoque innovador de StreamingLLM, centrado en el mecanismo de sumidero de atención, garantiza un uso y un rendimiento estables de la memoria, incluso cuando se procesan textos de hasta 4 millones de tokens de longitud", afirma Yang You, joven profesor presidencial de informática en la Universidad Nacional de Singapur, que no participó en este trabajo. “Esta capacidad no sólo es impresionante; es transformador y permite aplicar StreamingLLM en una amplia gama de aplicaciones de IA. El rendimiento y la versatilidad de StreamingLLM la marcan como una tecnología muy prometedora, preparada para revolucionar la forma en que abordamos las aplicaciones de generación impulsadas por IA”. Tianqi Chen, profesor asistente en los departamentos de aprendizaje automático y ciencias de la computación de la Universidad Carnegie Mellon, que tampoco participó en esta investigación, estuvo de acuerdo y dijo: “La transmisión de LLM permite la extensión fluida de la duración de la conversación de modelos de lenguaje grandes. Lo hemos estado utilizando para permitir la implementación de modelos Mistral en iPhones con gran éxito”. Los investigadores también exploraron el uso de receptores de atención durante el entrenamiento del modelo anteponiendo varios tokens de marcador de posición en todas las muestras de entrenamiento. Descubrieron que el entrenamiento con receptores de atención permitía que un modelo mantuviera el rendimiento con un solo receptor de atención en su caché, en lugar de los cuatro que normalmente se requieren para estabilizar el rendimiento de un modelo previamente entrenado. Pero si bien StreamingLLM permite que un modelo lleve a cabo una conversación continua, el modelo no puede recordar palabras que no estén almacenadas en el caché. En el futuro, los investigadores planean abordar esta limitación investigando métodos para recuperar tokens que han sido desalojados o permitir que el modelo memorice conversaciones anteriores. StreamingLLM se ha incorporado a la gran biblioteca de optimización de modelos de lenguaje de NVIDIA, TensorRT-LLM . ****Este trabajo está financiado, en parte, por el MIT-IBM Watson AI Lab, el MIT Science Hub y la Fundación Nacional de Ciencias de EE. UU.