A medida que crece la demanda de tareas de razonamiento, se espera cada vez más que los modelos de lenguaje grande (LLM) generen secuencias más largas o cadenas paralelas de razonamiento. Sin embargo, el rendimiento del tiempo de inferencia está severamente limitado por la huella de memoria del caché clave-valor (KV), no solo el número de tokens producidos. En un artículo reciente, los investigadores de Nvidia y la Universidad de Edimburgo introducen Sparsificación de memoria dinámica (DMS)—Se un método de modernización y eficiente de datos que comprime los cachés de KV y desbloquea Hiper-escala de tiempo de inferencia sin degradación de la precisión del modelo.
El cuello de botella: kv caché en la inferencia del transformador
Los modelos basados en transformadores como GPT, LLAMA y QWEN usan cachés de KV para almacenar representaciones de tokens pasadas para la generación autorregresiva. Este caché crece linealmente con la longitud y el ancho de la secuencia (hilos paralelos), que consume grandes cantidades de memoria de GPU y conduce a una inferencia más lenta debido al acceso frecuente a la memoria.
Las técnicas existentes para la optimización de caché de KV dependen de la heurística sin entrenamiento, como el desalojo de token basado en el peso de atención, o requieren una fuerte modernización posterior al entrenamiento como la compresión de la memoria dinámica (DMC). Ambos tienen desventajas significativas: el primero tiende a dañar la precisión, mientras que el segundo es computacionalmente costoso.
Sparsificación de memoria dinámica DMS: compresión sin compromiso
Sparsificación de memoria dinámica DMS aborda estas limitaciones con un enfoque híbrido: esparea el caché de KV como los métodos de poda tradicionales, pero lo hace con una sobrecarga mínima de entrenamiento (~ 1,000 pasos) y desalojo retrasadoque conserva los tokens temporalmente después de que están marcados para la eliminación. Este diseño conserva una información de contexto importante y evita caídas de precisión abrupta.
La idea central es hacer que las decisiones de desalojo diferenciables durante el entrenamiento utilizan un mecanismo de muestreo basado en gumbel-sigmoides. Los tokens predichos para el desalojo futuro siguen siendo utilizables para una duración de la ventana deslizante antes de ser descartadas, lo que permite que el modelo absorba su valor informativo de manera más efectiva.
Modernización eficiente con datos mínimos
A diferencia de DMC, que requiere miles de pasos de entrenamiento y una optimización compleja basada en gradientes, DMS no presenta parámetros adicionales por cabeza de atención. Reutiliza una pequeña parte del mecanismo de atención (una sola neurona) para predecir el desalojo. Esto hace que DMS sea ideal para modernizar modelos existentes sin cambios arquitectónicos.
Los resultados empíricos muestran que con tan pocos como Pasos de entrenamiento de 1kLos DM pueden lograr Compresión de caché de 8 × kVpreservar o incluso mejorar el rendimiento del modelo en las tareas de razonamiento.
Resultados de referencia: rendimiento de escala sin costo de escala
El equipo de investigación probó DMS en puntos de referencia de razonamiento como:
- AIME 2024 (Matemáticas avanzadas)
- Matemáticas 500 (Resolución de problemas matemáticos)
- Diamante gpqa (QA de ciencia dura)
- LivecodeBench (Generación de código)
A través de los tamaños del modelo (Qwen-R1 1.5B, 7B y 32B) mejoró el rendimiento del partido exacto 9.1 puntos en AIME, 7.6 en GPQAy 9.6 en LivecodeBenchtodos bajo la misma memoria y presupuestos de calcular.
En comparación con las líneas de base de alto rendimiento como Quest y TOVA, los DMS constantemente los superaron en ambos KV Cache Leer eficiencia (proxy de tiempo de ejecución) y Uso de la memoria máximalogrando mejores fronteras de Pareto.
Utilidad general de propósito
DMS también se mantiene en tareas que no son de condición. En puntos de referencia de corto contexto como MMLU, GSM8K y HELLASWAG, el rendimiento mantenido por DMS a relaciones de compresión hasta 4 × con degradación mínima (~ 3.5 puntos). En tareas de contexto largo como Needle-in a-Haystack y el seguimiento de variables, los DM incluso superaron los modelos de vainilla, lo que sugiere su potencial para mitigar temas como la información sobre las secuencias largas.
Conclusión
En conclusión, la dispersión de memoria dinámica (DMS) presenta una solución práctica y escalable para mejorar la eficiencia de tiempo de inferencia de los modelos de lenguaje basados en transformadores. Al comprimir de manera inteligente el caché de KV con un reentrenamiento mínimo, DMS permite que los modelos razonen sobre secuencias más largas o en paralelo sin aumentar el tiempo de ejecución o las demandas de memoria. Sus ganancias consistentes en una variedad de razonamiento y tareas de uso general resaltan su versatilidad y efectividad. A medida que los LLM se implementan cada vez más en entornos limitados por recursos, DMS ofrece un camino convincente: la compresión, la precisión y la facilidad de integración de las cargas de trabajo de inferencia del mundo real.
Mira el Papel. Todo el crédito por esta investigación va a los investigadores de este proyecto. Además, siéntete libre de seguirnos Gorjeo Y no olvides unirte a nuestro 99k+ ml de subreddit y suscribirse a Nuestro boletín.
▶ ¿Busca mostrar su producto, seminario web o servicio a más de 1 millón de ingenieros de IA, desarrolladores, científicos de datos, arquitectos, CTO y CIO? Explore una asociación estratégica
Nikhil es consultor interno en MarktechPost. Está buscando un doble grado integrado en materiales en el Instituto Indio de Tecnología, Kharagpur. Nikhil es un entusiasta de AI/ML que siempre está investigando aplicaciones en campos como biomateriales y ciencias biomédicas. Con una sólida experiencia en la ciencia material, está explorando nuevos avances y creando oportunidades para contribuir.
