💬 NLP

Learning how to Forget: Fine-tuning for Long-Context Sparse Attention

Este artículo introduce un método de ajuste fino para modelos transformer que permite una inferencia eficiente de contexto largo con atención dispersa en hardware moderado al permitir que los modelos se co-adapten con diversas políticas de caché KV, superando a menudo los enfoques de atención exacta.

Autores originales: Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Publicado 2026-08-21

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Los sistemas de inteligencia artificial modernos que generan texto similar al humano dependen de un mecanismo que actúa como una memoria a corto plazo, permitiéndoles recordar lo que se dijo anteriormente en una conversación o en un documento largo. Esta memoria se almacena en un búfer digital que crece más con cada nueva palabra que el sistema procesa. Para que estos sistemas funcionen bien a lo largo de textos extensos, esta memoria debe ser vasta, pero el hardware informático necesario para albergarla es costoso y limitado. Cuando el búfer de memoria se llena, el sistema debe decidir qué piezas de información antigua debe descartar para hacer espacio a la nueva. Si desecha la información equivocada, el sistema pierde su capacidad para razonar o responder preguntas con precisión. Esto crea un difícil compromiso: mantener la memoria pequeña ahorra dinero y permite que el sistema funcione en equipos estándar, pero conlleva el riesgo de perder el contexto necesario para ser inteligente.

Los investigadores han intentado resolver esto durante mucho tiempo enseñando al sistema a ser selectivo sobre lo que conserva, un proceso conocido como atención dispersa (sparse attention). Sin embargo, un nuevo estudio revela una falla crítica en la forma en que estos sistemas han sido entrenados hasta ahora. La mayoría de los métodos existentes entrenan a la IA utilizando una memoria perfecta e ilimitada y luego intentan forzarla a operar con una limitada más tarde. Los investigadores descubrieron que este enfoque falla porque la IA nunca aprendió cómo funcionar bajo las restricciones específicas que realmente enfrentaría. Al entrenar al modelo para olvidar intencionalmente y adaptarse a un tamaño de memoria fijo desde el principio, el equipo demostró que el sistema podía desempeñarse significativamente mejor que aquellos entrenados con recursos ilimitados, incluso ejecutándose en un solo chip de computadora moderadamente potente.

El equipo, liderado por científicos de Amazon Web Services y la Universidad de Ámsterdam, desarrolló una nueva forma de ajustar estos grandes modelos de lenguaje. En lugar de utilizar supercomputadoras masivas para simular una memoria perfecta, enseñaron a los modelos a co-adaptarse con una política específica de gestión de memoria. Imagine a un bibliotecario que es entrenado para organizar libros en una biblioteca con estantes infinitos, solo para que luego se le diga que trabaje en una habitación diminuta con un solo estante. El bibliotecario entrenado en la gran biblioteca probablemente tendría dificultades para priorizar qué conservar en la habitación pequeña. En contraste, el método propuesto en este artículo entrena al bibliotecario directamente en la habitación pequeña, enseñándole exactamente qué libros conservar y cuáles descartar basándose en las reglas de ese espacio específico. Esto permite que el modelo aprenda el ritmo de sus propias limitaciones, en lugar de intentar desaprender los hábitos de tener demasiado espacio.

Los investigadores probaron este enfoque en un modelo de cuatro mil millones de parámetros, un tamaño sustancial pero manejable. Realizaron sus experimentos en una sola tarjeta gráfica con 40 gigabytes de memoria, una configuración asequible para muchas organizaciones en comparación con los grupos de docenas de tarjetas requeridos por los métodos anteriores. Compararon esta nueva técnica de entrenamiento contra el método estándar, que utiliza una técnica llamada paralelismo de secuencia para dividir la carga de memoria entre múltiples dispositivos costosos. Los resultados mostraron que los modelos entrenados con el nuevo método a menudo superaron a los estándar, particularmente cuando la tarea requería que el sistema generara respuestas específicas y concisas en lugar de texto largo y divagante. En varias pruebas que involucraban preguntas complejas y extracción de datos, el método estándar produjo salidas que eran demasiado largas y llenas de números aleatorios y sin sentido, mientras que el nuevo método aprendió a detenerse en el momento adecuado y proporcionar el valor único correcto.

Una parte clave de este éxito fue la mejora del "oráculo de heavy-hitter" (heavy-hitter oracle), una estrategia popular para decidir qué información conservar. Esta estrategia funciona rastreando a qué piezas de información presta más atención el modelo a lo largo del tiempo. Los investigadores descubrieron que la versión original de esta estrategia era lenta e ineficiente. Reescribieron el código subyacente para que funcionara mucho más rápido, permitiendo que el sistema calculara estas puntuaciones de importancia sin ralentizar todo el proceso. Esta optimización significó que el sistema podía tomar decisiones inteligentes sobre qué olvidar en tiempo real, sin necesidad de la enorme potencia computacional que usualmente acompaña a tales tareas. El equipo también lanzó una nueva biblioteca de software de código abierto para poner estas técnicas a disposición de otros, con el objetivo de reducir la barrera para cualquiera que desee construir sistemas de IA de contexto largo sin necesitar una fortuna en hardware.

El estudio destaca que la forma en que se entrena un modelo es tan importante como el hardware en el que se ejecuta. Cuando los investigadores obligaron al modelo a entrenar con las mismas restricciones de memoria que enfrentaría durante su uso, este aprendió a navegar esas restricciones de manera efectiva. En una prueba específica que involucraba datos JSON, el método estándar falló por completo, incapaz de encontrar los puntos de datos correctos, mientras que el nuevo método tuvo éxito en identificarlos aproximadamente la mitad de las veces. Esto sugiere que la capacidad de manejar contextos largos no es solo una cuestión de tener más memoria, sino de enseñar al sistema cómo gestionar la memoria que posee. Los hallazgos indican que, para muchas aplicaciones, el camino más efectivo hacia adelante no es construir computadoras más grandes, sino enseñar al software a ser más eficiente con los recursos que ya posee.

Resumen Técnico: Aprendiendo a Olvidar: Ajuste Fino para la Atención Dispersa de Contexto Largo

1. Planteamiento del Problema

Los modelos de lenguaje de gran tamaño (LLM) modernos requieren la capacidad de procesar contextos extremadamente largos para tareas como conversaciones de múltiples turnos, uso de herramientas y razonamiento de cadena de pensamiento (chain-of-thought). Mientras que las implementaciones estándar de transformadores escalan cuadráticamente en cómputo y linealmente en memoria con el ancho del contexto, se han propuesto métodos de atención dispersa (sparse attention) para lograr un escalado de tiempo lineal y de memoria constante. Estos métodos utilizan un caché de Clave-Valor (KV) de tamaño fijo donde la información se expulsa basándose en políticas específicas (por ejemplo, Heavy-Hitter Oracle, H2O) una vez que el caché está lleno.

Sin embargo, existe una brecha significativa entre la inferencia y el entrenamiento en este dominio:

Restricciones de Entrenamiento: El ajuste fino de modelos con políticas de atención dispersa es notoriamente difícil con presupuestos de hardware moderados. La diferenciación automática estándar (autograd) requiere almacenar el grafo de computación completo, incluyendo los buffers del caché KV para cada fragmento (chunk) de la secuencia. Para contextos largos, este requisito de memoria excede la memoria GPU disponible, lo que a menudo requiere configuraciones multi-GPU costosas con paralelismo de secuencia o de contexto (SP/CP).
Brecha de Co-adaptación: Los modelos entrenados con atención exacta (usando SP/CP) suelen rendir significativamente peor cuando se despliegan con atención dispersa en la inferencia. Esto se debe a que el modelo nunca es expuesto a las restricciones de memoria específicas y a la lógica de expulsión que enfrentará durante la inferencia.
Ineficiencia de Implementación: Las principales políticas de atención dispersa como H2O están subutilizadas en las librerías de inferencia de vanguardia (SotA) (por ejemplo, vLLM) porque las implementaciones actuales son lentas, careciendo de kernels optimizados para los requisitos específicos de la atención dispersa (por ejemplo, devolver pesos de atención sumados).

2. Metodología

El artículo propone un nuevo método de ajuste fino que permite a los modelos de transformadores co-adaptarse con políticas arbitrarias de caché KV en hardware moderado (demostrado en una sola Nvidia A100 40GB).

2.1 Checkpointing de Activación Anidado con Offloading a CPU

Para abordar el cuello de botella de memoria al entrenar en secuencias largas ( $N \gg N_C$ , donde $N_C$ es el tamaño del caché), los autores emplean una estrategia de checkpointing de activación anidada:

Bucle Exterior (Capas): El paso hacia atrás (backward pass) itera sobre las capas del modelo.
Bucle Interior (Celdas): La secuencia se divide en "celdas" (grupos de fragmentos). Para cada celda, se ejecuta un paso hacia adelante (forward pass) para almacenar los buffers de caché KV necesarios en la memoria de la CPU.
Mecanismo de Replay: En lugar de diferenciar a través de la política del caché KV (que suele ser no diferenciable o costosa), el paso hacia adelante registra todas las decisiones de expulsión. Durante el paso hacia atrás, estas decisiones se replican utilizando "cachés de replay" que imitan la lógica de expulsión del paso hacia adelante sin requerir el cálculo de gradientes a través de la política en sí misma.

2.2 Codificación Delta para Buffers de Caché KV

La optimización de memoria más crítica involucra cómo se almacenan los buffers de caché KV dentro del grafo de autograd.

Recurrencia Lineal: Los autores observan que el buffer de caché KV para un nuevo fragmento ($keys'$) es una transformación lineal del buffer anterior ($keys$), diferenciándose solo por los valores de los nuevos tokens y los índices específicos de los slots sobrescritos por la política de expulsión.
Almacenamiento Delta: En lugar de almacenar los tensores completos de $keys $y$ values$ para cada fragmento en el grafo de computación, el método almacena solo el delta (la diferencia) y los índices de los slots sobrescritos.
Hooks de Tensores Guardados de Autograd: Utilizando los autograd_saved_tensors_hooks de PyTorch, los autores implementan un mecanismo de "empaquetado" (packing). Durante el paso hacia adelante, los tensores completos se reemplazan por referencias a tuplas de (índice, delta). Durante el paso hacia atrás, los tensores completos se reconstruyen ("desempaquetados") a partir de los deltas y los buffers subsecuentes. Esto reduce la huella de memoria del grafo de autograd por un factor de $k$ (el número de fragmentos por celda), reduciéndola a niveles comparables con la inferencia.

2.3 Mejoras de la Política H2O

El artículo introduce mejoras específicas a la política Heavy-Hitter Oracle (H2O):

Independencia de Batch: La implementación original de H2O agregaba los puntajes a través de las dimensiones de batch. El nuevo método permite que las decisiones de expulsión se tomen de forma independiente para cada posición de batch.
Normalización: Se introduce un puntaje H2O normalizado para prevenir el sesgo hacia los tokens que han estado en el caché por más tiempo.
Optimización de Kernel: Para permitir un entrenamiento e inferencia eficientes, los autores proporcionan código Triton que devuelve los pesos de atención sumados junto con un kernel de Scaled Dot Product Attention (SDPA) de FlashInfer. Esto permite que H2O funcione eficientemente sin recurrir a implementaciones de SDPA ingenuas y lentas.

3. Contribuciones Clave

Nuevo Método de Ajuste Fino: Un método para realizar el ajuste fino de modelos de transformadores con atención dispersa y políticas arbitrarias de caché KV en recursos comparables a la inferencia de atención dispersa (por ejemplo, una sola A100). Combina checkpointing de activación anidado, offloading a CPU y codificación delta mediante hooks de autograd.
Implementación y Optimización de H2O: Mejoras metodológicas a la política H2O (independencia de batch, normalización) y una implementación de alto rendimiento utilizando kernels SDPA dedicados que devuelven pesos sumados, dando un paso significativo hacia latencias competitivas con las librerías de inferencia SotA.
Evaluación Exhaustiva: Experimentos en el conjunto de pruebas Helmet con ventanas de contexto de 64k y 128k, demostrando que aunque el rendimiento varía según el dataset, el método propuesto permite que los modelos se co-adapten con las políticas de caché KV, superando a menudo a los modelos entrenados con paralelismo de secuencia cuando se evalúan bajo inferencia de atención dispersa.
Librería KeysAndValues: El lanzamiento de una librería de código abierto (awslabs/keys_values) que proporciona código de alto rendimiento y fácil uso para la inferencia de contexto largo y el ajuste fino con atención dispersa.

4. Resultados Experimentales

Los autores evaluaron un modelo Qwen3-4B-Instruct en el benchmark Helmet a través de longitudes de contexto de 64k y 128k.

Comparación: Los modelos fueron entrenados usando:
- SP: Paralelismo de Secuencia con atención exacta (línea base SotA).
- US: El método propuesto con la política de caché KV específica en su lugar.
- NO: El checkpoint base sin ajuste fino.
Hallazgos:
- Resultados Mixtos en los Datasets de la Tabla 1: Para los datasets de la Tabla 1, los resultados fueron mixtos e inconcluyentes. La línea base SP (atención exacta) logró los mejores resultados para nq y hotpot qa. Para trivia qa, el checkpoint base sin ajuste fino (no) funcionó mejor, indicando que el ajuste fino no ayudó en este caso específico. El método propuesto (us) superó a SP en pop qa.
- Fuerte Desempeño en los Datasets de la Tabla 2: Para los datasets de la Tabla 2, us superó fuertemente tanto a sp como a no.
- Modos de Fallo de SP: Los modelos entrenados con SP exhibieron un modo de fallo donde las salidas eran excesivamente largas y contenían disparates aleatorios cuando se sometían a la expulsión de atención dispersa, probablemente porque nunca fueron entrenados para manejar el mecanismo de "olvido". En contraste, us aprendió a detenerse correctamente y usualmente produce un solo número o el objetivo correcto.
- Varianza de la Política: Los resultados variaron significativamente por dataset. Mientras que las variantes de H2O funcionaron bien en varias tareas, la política lastrec (lr) superó a H2O en datasets como nq y hotpot qa. Además, para nq y hotpot qa, la línea base SP (atención exacta) a menudo logró los mejores resultados, mientras que para pop qa y otros datasets de la Tabla 2, el método US con políticas de atención dispersa fue superior.
- Eficiencia: La implementación de H2O con el nuevo soporte de kernel representa una mejora importante en la eficiencia, dando un gran paso hacia latencias competitivas con las librerías SotA, aunque el artículo reconoce que la latencia de la atención dispersa actual sigue siendo significativamente mayor que los métodos que usan paralelismo de secuencia o de contexto.

5. Significancia y Reivindicaciones

El artículo afirma que la atención dispersa ha sido subutilizada en la práctica debido principalmente a problemas de latencia y a la dificultad de ajustar finamente los modelos para adaptarse a políticas de expulsión específicas.

Factibilidad: Los autores demuestran que el ajuste fino de alta calidad para la atención dispersa de contexto largo es factible en hardware moderado, eliminando la necesidad de masivos clústeres de entrenamiento distribuido únicamente para adaptar un modelo a una política de caché.
Necesidad de Co-adaptación: Los resultados sugieren fuertemente que los modelos deben ser entrenados bajo las mismas restricciones (políticas de expulsión) que enfrentarán durante la inferencia para evitar la degradación del rendimiento.
Dirección Futura: El artículo argumenta que al simplificar la estructura del caché KV y clarificar los requisitos de los kernels SDPA (por ejemplo, devolver pesos sumados), la comunidad puede integrar mejor la atención dispersa en las librerías de inferencia SotA.
Limitaciones: Los autores reconocen que, si bien su método mejora la latencia, la naturaleza secuencial de las decisiones de atención dispersa aún puede resultar en una latencia mayor en comparación con RingAttention o el Paralelismo de Secuencia para presupuestos de hardware muy grandes. Planean abordar esto mediante la fusión de kernels y el offloading asíncrono a CPU en trabajos futuros.

En resumen, el artículo proporciona una vía práctica para "aprender a olvidar", permitiendo que los modelos se especialicen para la inferencia de contexto largo en hardware accesible. Si bien el rendimiento relativo al paralelismo de secuencia varía según el dataset, el método a menudo supera al entrenamiento de atención exacta en escenarios que requieren la adaptación a políticas específicas de atención dispersa.