- StreamingLLM es un framework para desplegar LLM en aplicaciones de streaming con entradas de longitud infinita sin sacrificar eficiencia ni rendimiento
- Aborda el problema de que, en conversaciones de múltiples rondas que requieren interacciones largas, el cacheo de los estados Key/Value de tokens anteriores consume mucha memoria, y de que los LLM comunes no generalizan a textos más largos que la longitud de secuencia usada en el entrenamiento
- La window attention que cachea solo los KV recientes falla cuando la longitud del texto supera el tamaño de la caché; StreamingLLM recupera gran parte del rendimiento de la window attention con un attention sink que conserva los KV de los tokens iniciales
- Permite que LLM entrenados con una attention window de longitud finita generalicen a longitudes de secuencia infinitas sin fine-tuning, y realiza modelado de lenguaje estable y eficiente con más de 4 millones de tokens en Llama-2, MPT, Falcon y Pythia
- En configuraciones de streaming, muestra hasta 22.2 veces de speedup frente al baseline de recomputación con sliding window
- No amplía la ventana de contexto; conserva solo los tokens recientes y el attention sink, y descarta los tokens intermedios
- Si Llama-2 fue preentrenado con una ventana de contexto de 4096 tokens, el tamaño máximo de caché de Llama-2 en StreamingLLM también es 4096
- Si se ingresa un libro largo, el modelo solo reconoce los tokens más recientes, por lo que solo puede resumir la parte de la conclusión
- Los casos de uso adecuados son aplicaciones de streaming que requieren operación continua y deben evitar depender de datos pasados o usar mucha memoria; se mencionan como ejemplos las conversaciones de múltiples rondas y los asistentes diarios basados en LLM
- Es ortogonal a los métodos recientes de extensión de contexto y puede integrarse con ellos; en el contexto de StreamingLLM, context extension significa la posibilidad de almacenar más tokens recientes con un tamaño de caché mayor
- Un ejemplo de ejecución es
examples/run_streaming_llama.py --enable_streaming; para la configuración del entorno se usan Python 3.8,torch,transformers==4.33.0,accelerate,datasets,evaluate,wandb,scikit-learn,scipyysentencepiece - El código principal ya está publicado, incluyendo Llama-2, MPT, Falcon y Pythia; también están publicados el código de evaluación de perplexity y la demo de Streaming Llama Chatbot, mientras que el dataset StreamEval y el código de evaluación aún no están publicados
1 comentarios
Opiniones en Hacker News
Parece que algunos confunden este trabajo con algo como atención densa completa.
Lo que se plantea aquí no es permitir ver contenido muy lejano, sino más bien una mejora de eficiencia: reutilizar la caché para mantener la perplejidad, en lugar de volver a calcular la ventana deslizante cada vez y pagar el costo L² unas T veces.
Las pruebas también se midieron concatenando algo como
Q A Q A Q A Q A..., no haciendo que encontrara una respuesta mucho después, como enQ Q Q Q A A A A....Medir la perplejidad significa producir “texto legible”, es decir, oraciones localmente plausibles; no es evidencia de que “extraiga” algo desde un gran vacío triangular al que la atención no llega.
Si le das un libro y le pides escribir la primera palabra de cada párrafo, o resumir cada capítulo en una oración, creo que fallaría.
Al probarlo directamente, no me pareció que hiciera extensión de la longitud de contexto, y sí se ejecutaba bastante rápido.
Usó alrededor de 35 GB de memoria en una A100, y el uso se mantuvo fijo durante toda la ejecución.
Tomé un libro de Project Gutenberg, lo dividí por párrafos, los fui ingresando uno por uno y le pedí que respondiera “okay” a cada párrafo; luego, al hacerle una pregunta al final, alucinó por completo la respuesta.
Como referencia, durante los unos 10 minutos que estuve probándolo, tampoco logré hacer que el modelo base,
lmsys/vicuna-13b-v1.3, respondiera bien en inglés.https://gist.github.com/bluecoconut/9cae9e91fe3b1616ed650a96...
Aun así, es un trabajo interesante, y el punto clave parece ser el hallazgo de la Figure 2.
Las primeras dos capas muestran un patrón local que presta más atención a los tokens recientes, pero después de pasar por las capas inferiores, el modelo presta mucha atención a los tokens iniciales en todas las capas y cabezas.
Los autores llaman a esto “attention sinks” y consideran que, aunque no sean semánticamente importantes, como por Softmax la suma de las puntuaciones de atención debe ser 1, la atención sobrante tiene que ir a algún lado.
La explicación es que, en un modelo de lenguaje autorregresivo, los tokens iniciales son visibles desde casi todos los tokens posteriores, por lo que es fácil que el modelo aprenda a usarlos como esos sumideros.
StreamingLLM se parece más a un “hack” para corregir este comportamiento extraño que aparece al recortar la ventana de atención de un LLM, y como es un caso que muestra las grietas del uso de Softmax, me hace pensar que, si queremos LLM flexibles en cuanto a longitud de contexto, quizá otra función sería mejor.
A primera vista suena tan bueno que da para dudar de si es real, pero la calidad del trabajo parece razonable y la técnica es sorprendentemente simple.
La idea es aplicar atención, en cada capa, solo al primer token y a la ventana deslizante de contexto, ignorando los tokens intermedios.
Esto parece significar que cada capa va empujando la información relevante poco a poco hacia la parte posterior de la secuencia, de modo que la ventana de atención deslizante al final de la capa superior pueda verla.
Sin embargo, si el rango cubierto por todas las ventanas deslizantes no alcanza para conectar toda la secuencia, es posible que no pueda transmitir hacia adelante toda la información importante.
Por ejemplo, si todas las ventanas tienen la misma longitud y
profundidad del modelo × longitud de la ventana < longitud de la secuencia, aparece una limitación.Esto parece posible gracias a la observación de que Softmax debe hacer que la suma sea 1.
Viéndolo rápido, cuando el modelo no necesita prestar atención a tokens anteriores, tiende a usar el primer token como marcador de posición.
La primera vez que vi este problema fue en un comentario de Evan Miller en HN, donde decía que obligar a una cabeza de atención a distribuir toda su atención entre los tokens anteriores está mal, y que habría que permitir “no prestar atención” agregando 1 al denominador de Softmax.
Es bueno que hayan aprovechado esta observación sin reentrenar, y también me da curiosidad cómo habría cambiado el modelo si hubieran seguido la propuesta de Evan.
[2] https://news.ycombinator.com/item?id=36851494
Entrenaron el modelo con un token sumidero dedicado cuyos valores eran todos 0, pero aun así otros tokens iniciales terminaban usándose como sumideros, así que la conclusión parece ser que conviene tener un token sumidero dedicado.
Hasta donde sé, nunca había logrado mejorar realmente el rendimiento.
Al manipular la ventana de atención después del entrenamiento, Softmax + 1 quizá encaje mejor, pero no sé si alguien lo ha probado a gran escala.
Agregar memoria de caché de atención es una solución muy interesante para este problema.
Hace unos días también salió un paper con una observación relacionada en Vision Transformer.
Los modelos Transformer parecen elegir tokens para almacenar información global, y parece que necesitan una especie de “token para pensar”.
Al proporcionar tokens específicos para este propósito, el rendimiento mejora un poco, y las visualizaciones explicativas también resultan bastante interesantes.
[0] https://arxiv.org/pdf/2309.16588.pdf
En el ajuste fino, se podrían congelar los parámetros del modelo original y ajustar solo los parámetros que entran y salen de las nuevas unidades de caché de “tuning”.
Así se podrían intercambiar distintos conjuntos de unidades de tuning o usarlos juntos.
Sería una forma de mezclar una especie de superprompt, como una unidad para evitar groserías + una unidad de terminología específica + una unidad para escribir de forma concisa.
Si la cantidad de parámetros nuevos es lo bastante pequeña, aunque use más memoria, también podría ser posible un tuning rápido y efectivo mediante optimización de orden superior.
También se podría pensar en aumentar juntos la longitud de secuencia y la cantidad de unidades durante el entrenamiento.
Por ejemplo, usar solo unas pocas unidades para secuencias cortas y, al aumentar la longitud de las secuencias de entrenamiento, agregar unidades y seguir entrenando.
En lugar de usar un calendario arbitrario, quizá se podría controlar la expansión de la caché mediante análisis de rendimiento o de gradientes.
Los autores publicaron una FAQ, que puede aclarar parte de la confusión: https://github.com/mit-han-lab/streaming-llm/blob/main/READM...
Ante “¿se puede meter un texto largo, como un libro, en StreamingLLM para resumirlo?”, responden que se puede ingresar un texto largo, pero como el modelo solo reconoce los tokens más recientes, si se le mete un libro solo podría resumir los últimos párrafos y quizá no sea muy útil.
Es decir, no extiende la ventana de contexto del LLM ni refuerza la memoria de largo plazo; la fortaleza de StreamingLLM está en generar texto fluido a partir de tokens recientes sin refrescar la caché.
Puede que me equivoque, pero no parece ser una tecnología que permita al LLM referirse a contenido más allá de la longitud con la que fue entrenado, como mucha gente cree.
Más bien parece tratar el problema de mantener el rendimiento del modelo en textos largos; más precisamente, el rendimiento sobre contenido que todavía está dentro de la ventana de contexto.
La explicación es que el modelo aprende a poner una especie de carga en la atención de los tokens iniciales del texto, y cuando eso desaparece fuera de la ventana, se rompe; pero no tengo claro por qué ocurre eso.
Si no se trata de una entrada de instrucciones, pensaría que el texto intermedio es tan bueno como el texto inicial.
Me pregunto cómo manejan estas técnicas de ventana deslizante los casos en que aparece una instrucción inesperada solo al final.
Por ejemplo, supongamos que metes un libro en el modelo y en la última oración hay una instrucción que dice “devuelve la cantidad de letras m en la entrada anterior”; una persona suspiraría, lo volvería a leer y contaría, pero un LLM no tiene la capacidad de volver atrás y releer la entrada.
En este ejemplo, incluso ignorando las limitaciones propias de los LLM para contar letras, para resolverlo de verdad parecería que el LLM tendría que poder hacer loops y saltos arbitrarios.
Claro que eso traería problemas completamente nuevos, y quizá requeriría una arquitectura totalmente nueva.
La idea sería convertir la pregunta en palabras clave relevantes, buscar, y volver a revisar el contenido para encontrar la información pertinente.
Si existe el preprocesamiento necesario, un LLM que pueda “ir a investigar lo suficiente y luego responder” sería muy potente.
Durante los últimos aproximadamente 10.000 años hemos mejorado las tecnologías de gestión del conocimiento para superar la capacidad y el tiempo del cerebro individual, así que los modelos de lenguaje también deberían aprovechar métodos reales de investigación y predigestión, no solo una simple búsqueda de Bing.
La memoria de corto plazo no necesita recordar qué hacía cada fragmento de código; basta con etiquetarlo al leerlo y apoyarse en un índice compartido de etiquetas escalable.
Aunque cuanto más lo pienso, más se parece al preentrenamiento general de un LLM, y el índice de conocimiento se siente como si fuera un enorme bloque de pesos de LLM.
Sería más bien una capa por encima del LLM que un cambio en el funcionamiento del LLM en sí.
Puede que lo haya entendido mal, pero en ese caso, ¿no resuelve el problema de consulta el estado oculto?
Antes de responder, de todos modos tiene que absorber toda la entrada, así que no parece importar mucho si la instrucción está al principio o al final, aparte de la atención.
Claude 100K le pide al usuario que ponga las instrucciones al final.
O bien se podría usar un modelo rápido para verificar si hay una instrucción al final y moverla al principio.
No estoy seguro de que los modelos actuales puedan hacer esto ni siquiera con entradas cortas.
Dicho un poco en broma, los LLM se están esforzando muchísimo por reinventar las RNN, y si les das herramientas, creo que eventualmente lo harán.
Visto de otra manera, los modelos Transformer intentan predecir qué partes de una red RNN “vale la pena conservar” cuando hay restricciones de recursos.
Los Transformer actuales usan una heurística simple, y este resultado mejora esa heurística.
Como en muchos problemas NP-completos, puede haber aproximaciones útiles aunque no sean perfectas, y los Transformer muestran que eso también es posible en redes neuronales.
Durante un tiempo estuvo en una posición intermedia en los rankings open source, así que es un enfoque bastante legítimo; simplemente no está de moda.
[1]: https://huggingface.co/blog/rwkv
La principal ventaja de los Transformer frente a las RNN es la paralelización del entrenamiento.
Las RNN tienen desvanecimiento del gradiente durante el entrenamiento, y también es difícil lograr alta utilización general, porque se necesitan lotes grandes, así que son complicadas.
La existencia de modelos como RWKV muestra que podría haber un futuro en el que se entrene como Transformer y se infiera como RNN.
Relacionado con esto, el profesor Han del MIT está dictando un curso público de TinyML.
https://news.ycombinator.com/item?id=37620507
https://efficientml.ai