🤖 machine learning

MoE-nD: Per-Layer Mixture-of-Experts Routing for Multi-Axis KV Cache Compression

El paper presenta MoE-nD, un marco de mezcla de expertos que optimiza la compresión heterogénea de la memoria caché KV en modelos de lenguaje grandes asignando dinámicamente a cada capa una combinación específica de eliminación de tokens y cuantización, logrando así una reducción de memoria de 14x sin sacrificar la precisión en comparación con los métodos de compresión uniformes existentes.

Autores originales: Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

Publicado 2026-04-21

📖 4 min de lectura☕ Lectura para el café

CC BY 4.0

Autores originales: Libo Sun, Peixiong He, Po-Wei Harn, Xiao Qin

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

¡Claro que sí! Imagina que tienes un cerebro digital (un modelo de Inteligencia Artificial) que necesita leer un libro entero para responder una pregunta. Pero, hay un problema: su "mesa de trabajo" (la memoria) es muy pequeña. Si intenta leer todo el libro de una vez, la mesa se desborda y el cerebro se bloquea.

Para solucionar esto, los científicos han estado intentando dos cosas:

Tirar páginas: Decidir qué partes del libro no son importantes y borrarlas de la mesa (esto se llama evicción).
Escribir con letra más pequeña: Guardar la información con menos detalles para que ocupe menos espacio (esto se llama cuantización).

El problema de los métodos anteriores es que trataban a todas las páginas del libro igual. Decían: "Borremos el 50% de las páginas de todos los capítulos" o "Escribamos todo con letra pequeña en todos los capítulos". Pero, ¿y si el capítulo 1 es crucial y no se puede borrar nada, mientras que el capítulo 10 es solo una lista de ingredientes que se puede escribir con letra muy pequeña? Los métodos antiguos no hacían esa distinción.

La Solución: MoE-nD (El "Jefe de Obra" Inteligente)

Los autores de este paper proponen MoE-nD, que es como tener un jefe de obra muy inteligente que revisa cada capítulo del libro individualmente antes de empezar a trabajar.

Aquí te explico cómo funciona con una analogía sencilla:

1. El Problema de la "Receta Única"

Imagina que eres un chef y tienes que preparar 28 platos diferentes (los 28 "capítulos" o capas del cerebro de la IA).

El método antiguo: Te dice: "Para todos los platos, usa solo la mitad de los ingredientes y ponlos en platos pequeños".
El resultado: Algunos platos salen deliciosos, pero otros (como un pastel delicado) quedan arruinados porque necesitaban todos sus ingredientes y un plato grande.

2. La Innovación: "MoE-nD" (Mezcla de Expertos)

MoE-nD actúa como un chef experto que tiene una hoja de ruta personalizada para cada plato.

Para el Plato 1 (que es muy sensible), el chef dice: "¡No toques nada! Usa todos los ingredientes y un plato gigante".
Para el Plato 15 (que es robusto), el chef dice: "¡Perfecto! Puedes tirar la mitad de los ingredientes y usar un plato minúsculo".
Para el Plato 20, el chef decide: "No tires ingredientes, pero sí escribe la receta con letra más pequeña (menos bits)".

El sistema MoE-nD decide, para cada "capítulo" del cerebro de la IA, exactamente cuánto borrar y cuánto comprimir, basándose en qué tan importante es ese capítulo específico.

¿Por qué es un gran avance?

En el papel, demostraron algo increíble:

Lograron comprimir la memoria de la IA 14 veces más (de 1.9 GB a solo 136 MB).
Lo asombroso: A pesar de haber tirado tanta información, la IA respondió tan bien como si hubiera tenido toda la memoria original.
Los métodos antiguos, al tratar a todos los capítulos por igual, perdían mucha precisión (como si el chef arruinara el pastel por ahorrar espacio).

¿Cuándo funciona y cuándo no?

El paper también es honesto sobre sus límites:

Funciona genial cuando el libro es muy largo (como leer una novela de 16,000 páginas). Ahí, el "jefe de obra" tiene mucho trabajo que hacer y puede ahorrar mucho espacio sin perder calidad.
No hace falta cuando el libro es muy corto (como un tweet o una pregunta simple). Si el texto es corto, el "jefe de obra" dice: "No necesito borrar nada, todo cabe en la mesa". En esos casos, el sistema se vuelve automático y no gana nada extra, pero tampoco pierde nada.

En resumen

MoE-nD es como pasar de tener un martillo (que golpea todo igual) a tener un kit de herramientas de precisión (destornilladores, alicates, etc.) para cada tarea. Permite a la Inteligencia Artificial leer textos larguísimos en computadoras normales (incluso en laptops) sin volverse loca ni perder la cabeza, simplemente sabiendo exactamente qué guardar y qué tirar en cada momento.

¡Es un paso gigante para que la IA sea más rápida, barata y accesible para todos!

1. El Problema

La inferencia de Grandes Modelos de Lenguaje (LLM) con contextos largos se ve limitada principalmente por el consumo de memoria del caché KV (Key-Value), que puede alcanzar decenas de gigabytes por secuencia, causando fallos de memoria (OOM) en hardware comercial.

Las técnicas de compresión existentes abordan este problema actuando sobre un solo eje del tensor KV de cuatro dimensiones:

Evicción (Secuencia): Eliminar tokens.
Cuantización (Precisión): Reducir bits por elemento.
Proyección de bajo rango: Reducir dimensiones.
Compartición entre capas: Reutilizar KV en capas adyacentes.

La limitación crítica: Estos métodos aplican la misma estrategia de compresión (la misma proporción de evicción o el mismo ancho de bits) a todas las capas del transformador de manera uniforme. El artículo demuestra que esta homogeneidad es ineficiente: diferentes capas responden de manera drásticamente distinta a la evicción frente a la cuantización. Mantener una política uniforme deja precisión sobre la mesa.

2. Metodología: MoE-nD

Los autores proponen MoE-nD, un marco basado en Mezcla de Expertos (MoE) que trata cada eje de compresión como un "experto" y utiliza un enrutador para asignar la combinación óptima de parámetros a cada capa individualmente bajo un presupuesto global de memoria.

Componentes Clave:

Espacio de Diseño: Cada capa $\ell$ $ℓ$ tiene tres parámetros ajustables:
- Ratio de retención (Keep Ratio): Fracción de tokens retenidos tras la evicción.
- Bits para K (Key): Ancho de bits para las claves (ej. 16, 8, 4).
- Bits para V (Value): Ancho de bits para los valores.
Calibración Offline (Tabla de Sensibilidad):
- Se mide la sensibilidad de cada capa a diferentes operaciones de compresión.
- Se utiliza una métrica proxy ligera: el error L2 relativo entre la salida de atención de precisión completa y la comprimida en una sola capa.
- Hallazgo Empírico: La sensibilidad varía enormemente entre capas. Por ejemplo, en evicciones agresivas (75-90%), la diferencia de error entre la capa más sensible y la menos sensible puede ser de 500x a 700x. En cambio, la cuantización de V es casi uniforme.
Solucionador Greedy (Voraz):
- Dado que la búsqueda exhaustiva es intratable ( $\approx 10^{42}$ combinaciones para un modelo de 28 capas), se utiliza un solucionador greedy.
- Comienza con la configuración más barata en todas las capas e itera mejorando la capa que ofrece la mayor reducción de sensibilidad por unidad de memoria adicional ( $\Delta S / \Delta m$ ) hasta agotar el presupuesto global.
Parche de Atención Heterogénea:
- Para manejar la inferencia en tiempo real donde cada capa tiene una longitud de caché y precisión diferentes, se modifica el kernel de atención.
- Se utilizan arrays de posición por capa para rastrear los tokens retenidos y se invierte la rotación RoPE (Rotary Positional Embedding) basándose en las posiciones originales, no en las globales.

3. Contribuciones Clave

Enrutamiento Multi-eje por Capa: Es el primer trabajo que realiza un enrutamiento conjunto de la proporción de evicción y los bits de cuantización (K y V) a nivel de capa bajo un presupuesto global.
Descubrimiento de Heterogeneidad: Demuestra que la compensación (trade-off) entre evicción y cuantización no es global; una capa puede preferir evicción mientras otra prefiere cuantización para el mismo costo de memoria.
Aislamiento de Contribuciones: Mediante una cadena de ablación, identifican que la ganancia principal proviene del enrutamiento de evicción por capa, no del enrutamiento de cuantización (que aporta poco o nada en comparación).
Definición de Alcance: Caracterizan cuándo el método es útil (contextos largos o presupuestos ajustados) y cuándo no (contextos cortos donde el solucionador elige no evictar nada).

4. Resultados Experimentales

Los experimentos se realizaron en el modelo DeepSeek-R1-Distill-Qwen-7B en varios benchmarks:

LongBench-v1 (Contexto Largo - 16k tokens):
- La variante heterogénea (2d_hetero) logra una compresión de 14x (reduciendo la memoria de 1.9 GB a 136 MB) sin pérdida detectable de precisión en comparación con la línea base sin compresión.
- En el mismo punto de memoria (136 MB), los métodos uniformes (2d_uniform) y los basados solo en evicción (1d) pierden entre un 33% y un 64% de la precisión promedio.
- En tareas individuales, la mejora es aún más drástica, con algunas líneas base perdiendo hasta un 99% de precisión.
Benchmarks de Razonamiento (AIME-24/25):
- MoE-nD supera consistentemente a la línea base de dos ejes no heterogénea en todos los presupuestos.
- En presupuestos muy ajustados (b=64), donde otros métodos caen a 0% de precisión, MoE-nD mantiene un rendimiento significativo (hasta 30%).
Resultados Nulos (MATH-500 y TREC):
- En tareas de contexto corto, el método no muestra ventaja. El solucionador greedy detecta que el presupuesto es suficiente para mantener todos los tokens (keep=1.0) en la mayoría de las capas, degenerando el método a una configuración uniforme. Esto valida que el método es específico para escenarios donde la compresión es estrictamente necesaria.

5. Significado e Impacto

Eficiencia Sin Sacrificio: Demuestra que es posible comprimir drásticamente el caché KV en modelos de razonamiento de largo contexto sin sacrificar la precisión, algo que los métodos uniformes no logran.
Paradigma de Compresión: Establece un nuevo estándar para la compresión de LLMs, moviéndose de políticas "talla única" (one-size-fits-all) a políticas adaptativas por capa que reconocen la heterogeneidad interna de la red neuronal.
Viabilidad Práctica: Aunque actualmente tiene limitaciones de implementación (incompatibilidad temporal con FlashAttention y necesidad de recalibración por familia de modelos), ofrece una plantilla clara para futuras optimizaciones de hardware y software.

En resumen, MoE-nD resuelve el cuello de botella de memoria en inferencia de LLMs de largo contexto mediante una asignación inteligente y dinámica de recursos de compresión, logrando un equilibrio óptimo entre memoria y precisión que los métodos actuales no pueden alcanzar.