Retrofitting Linear Attention into Diffusion Language Models

Imagina el mundo de la inteligencia artificial como una biblioteca enorme y bulliciosa donde un robot bibliotecario superinteligente intenta escribir una historia. Durante años, la forma estándar en que este bibliotecario trabajaba era "autorregresiva": escribía una palabra, se detenía, pensaba en todo lo que acababa de escribir, escribía la siguiente palabra, se detenía y volvía a pensar. Era como construir una torre ladrillo a ladrillo, esperando a que el pegamento se secara antes de añadir el siguiente. Aunque era fiable, este método era lento, especialmente para historias largas.

Recientemente, llegó un nuevo estilo llamado "Difusión". En lugar de construir ladrillo a ladrillo, el bibliotecario de Difusión comienza con una página en blanco llena de signos de interrogación e intenta adivinar toda la historia de una vez, refinando las conjeturas una y otra vez hasta que las palabras cobran sentido. Esto es como tener un equipo de artistas pintando un mural simultáneamente; pueden trabajar mucho más rápido porque no están esperando a que se seque el ladrillo anterior. Sin embargo, incluso este método rápido tiene un inconveniente. A medida que el mural se hace más grande, los artistas tienen que mirar constantemente cada una de las palabras que ya han decidido para asegurarse de que las nuevas encajen. Es como intentar recordar un libro de 10.000 páginas mientras escribes la página 101, lo que eventualmente ralentiza a todos y llena la memoria de la biblioteca.

Este es el rompecabezas que un equipo de investigadores abordó. Se preguntaron: "¿Podemos hacer que este método de Difusión rápido sea aún más rápido cambiando la forma en que el bibliotecario recuerda el pasado?". Su respuesta fue un truco ingenioso llamado "Atención Híbrida por Bloques" (Block-Hybrid Attention). Se dieron cuenta de que, si bien el bibliotecario necesita recordar el párrafo actual perfectamente (para que las frases fluyan), no necesita volver a leer cada una de las palabras de los capítulos anteriores en alta definición. En su lugar, pueden resumir los capítulos antiguos en una "hoja de resumen" diminuta y de tamaño fijo. Esto permite al bibliotecario escribir las nuevas partes de la historia a la velocidad del rayo sin verse agobiado por el peso de toda la historia.

La Gran Idea: Un Sistema de Memoria de Dos Velocidades

El artículo presenta un método para modernizar (o actualizar) un modelo de lenguaje de Difusión existente y potente (específicamente un modelo de 16 mil millones de parámetros llamado LLaDA 2.1) para que pueda utilizar esta estrategia de la "hoja de resumen" sin necesidad de ser reentrenado desde cero.

Piensa en el cerebro del modelo como si tuviera 20 diferentes capas de pensamiento. En el modelo original, cada una de las 20 capas actúa como un bibliotecario súper atento que lee cada palabra anterior para comprender la actual. Esto es preciso pero pesado. La innovación de los investigadores, la Atención Híbrida por Bloques, divide el trabajo en dos modos:

1. El Modo "Ahora" (Atención Exacta): Cuando el modelo está trabajando en el bloque actual de palabras que está generando activamente, mantiene su modo de "súper atención". Utiliza la memoria estándar y de gran capacidad para observar cada palabra en el párrafo actual y asegurar que las frases tengan un sentido perfecto.
2. El Modo "Pasado" (Atención Lineal): Cuando el modelo necesita recordar lo que sucedió en los bloques anteriores (los capítulos terminados), cambia a la "Atención Lineal". En lugar de releer todo el libro, consulta un estado de resumen de tamaño fijo. Es como consultar el índice de un libro o un resumen de una página en lugar de releer todo el texto. Este resumen mantiene el mismo tamaño sin importar cuán largo sea el libro.

Cómo lo Hicieron: La Actualización de Dos Etapas

Los investigadores no solo cambiaron las piezas y esperaron lo mejor; utilizaron un proceso cuidadoso de dos etapas para asegurar que el modelo no perdiera su inteligencia durante la actualización.

• Etapa 1: El Entrenamiento de Sombra. Tomaron el modelo original y brillante (el "Maestro") y lo congelaron. Luego, reemplazaron 6 de las 20 capas de atención con su nueva versión "Híbrida por Bloques" (el "Estudiante"). El Estudiante fue entrenado para imitar exactamente la salida del Maestro, utilizando una versión "corrupta" del texto como entrada. Era como tener a un estudiante observando de cerca a un maestro chef, intentando copiar sus movimientos exactos y probar los mismos sabores, pero solo para los platos específicos que al estudiante le correspondía cambiar. Esta etapa tomó alrededor de 24 horas.
• Etapa 2: El Ajuste Fino (Fine-Tuning). Una vez que el Estudiante aprendió a imitar al Maestro, los investigadores dejaron que todo el modelo trabajara junto de nuevo. Utilizaron una técnica llamada LoRA (Adaptación de Bajo Rango) para realizar ajustes diminutos y precisos en las conexiones del modelo. Esto fue como darle al modelo actualizado un rápido "pulido" para corregir cualquier pequeño fallo que ocurriera debido a que las nuevas partes ahora estaban interactuando con las antiguas. Esta etapa tomó 6 horas.

Los Resultados: Más Rápido, Más Ligero y Todavía Inteligente

El equipo probó su modelo actualizado, al que llamaron LLaDA-HYBRID, para ver si seguía siendo bueno escribiendo y resolviendo problemas, y si era realmente más rápido.

¿Todavía sabe pensar?
Sorprendentemente, sí. A pesar de que reemplazaron 6 de las 20 capas con este método de la "hoja de resumen", el rendimiento del modelo se mantuvo muy cercano al original.

• En una prueba de programación llamada HumanEval, el modelo original obtuvo un 75.6% y el modelo híbrido obtuvo un 72.0%.
• En una prueba de matemáticas llamada CMATH, el original obtuvo un 88.3% y el híbrido un 86.7%.
• Curiosamente, en otra prueba de programación llamada MBPP+, el modelo híbrido mejoró, pasando de 57.7% a 63.0%.
  El artículo sugiere que, aunque el modelo es ligeramente menos perfecto en las tareas de razonamiento más difíciles (como un difícil cuestionario científico llamado GPQA-Diamond, donde bajó de 38.9% a 30.8%), preservó en gran medida su capacidad para escribir código y resolver problemas matemáticos.

¿Es realmente más rápido?
Aquí es donde ocurre la magia. Debido a que el modelo ya no tiene que releer todo el historial de la conversación por cada nueva palabra, puede manejar más solicitudes a la vez.

• Al probar cuántas palabras podía generar el modelo por segundo, el modelo híbrido fue 1.7 veces más rápido que el original en su rendimiento máximo (manejando 128 solicitudes concurrentes).
• El uso de memoria del modelo original crecía a medida que la historia se hacía más larga, llegando eventualmente a un muro donde no podía manejar más usuarios. El modelo híbrido, gracias a su "hoja de resumen" de tamaño fijo, pudo manejar más solicitudes concurrentes antes de quedarse sin memoria.

La Conclusión

Este artículo sugiere que no siempre necesitamos construir un nuevo IA más rápido desde cero. En su lugar, podemos tomar una IA potente ya existente y darle un sistema de "memoria híbrida". Al mantener su enfoque agudo en el momento presente mientras resume el pasado en una forma compacta, podemos hacer que estos modelos sean significativamente más rápidos y eficientes. Los investigadores descubrieron que esta actualización podía realizarse en unas 60 horas de entrenamiento con datos públicos, ofreciendo un camino prometedor para hacer que los servicios de IA sean más rápidos y económicos sin sacrificar demasiado su inteligencia. Sin embargo, los autores señalan que esto es solo el comienzo; solo actualizaron 6 capas, y el trabajo futuro podría encontrar formas de actualizar incluso más capas o manejar historias aún más largas.