Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

1. El Problema: Rigidez en la Poda de Tokens por Capa

Los Modelos de Lenguaje Grandes (LLM) enfrentan un desafío crítico en escenarios de contexto largo: el alto consumo de memoria debido a la caché de clave-valor (KV). Para mitigar esto, se han desarrollado métodos de poda de tokens por capa (layer-wise token pruning), que seleccionan un subconjunto de tokens importantes en capas específicas para retener en la caché KV y descartar el resto.

Sin embargo, los métodos existentes (como FastKV, GemFilter, PyramidInfer) sufren de una limitación fundamental:

• Capas de Selección Fijas: Eligen una capa predefinida e inmutable para realizar la selección de tokens, independientemente de la tarea.
• Falta de Flexibilidad: Esta rigidez provoca variaciones significativas en el rendimiento.
  • En tareas fáciles (ej. preguntas de comprensión lectora donde la información relevante se identifica temprano), una selección en capas medias funciona bien.
  • En tareas difíciles (ej. recuperación de pares clave-valor o "Needle in a Haystack" con alta similitud semántica), la información necesaria no se estabiliza hasta capas más profundas. Forzar una selección temprana en estas tareas degrada severamente la precisión.
• Compromiso Ineficiente: Para mantener la precisión en tareas difíciles, los métodos actuales deben posponer la selección a capas más profundas (aumentando el costo de prellenado) o aumentar el presupuesto de memoria, perdiendo así sus beneficios de eficiencia.

2. Metodología: ASL (Adaptive Selection Layer)

Los autores proponen ASL, un método sin entrenamiento (training-free) que determina dinámicamente la capa óptima para la selección de tokens basándose en los patrones de atención observados durante la inferencia.

Mecanismo Central: Variación de Rangos de Tokens

La idea fundamental es monitorear la varianza de los rangos de los tokens ordenados por su puntuación de atención a través de capas consecutivas.

1. Observación de Patrones: En capas tempranas, las puntuaciones de atención suelen estar distribuidas uniformemente. En capas profundas, se localizan en un subconjunto estable de tokens.
2. Cálculo de Variación:
   • ASL calcula las puntuaciones de atención promediadas (pooled) sobre una ventana de tokens recientes.
   • Ordena los tokens por estas puntuaciones y calcula la varianza de sus rangos en un rango de $L_{obs}$ capas consecutivas.
   • Se normaliza esta varianza respecto a la varianza inicial ($L_{min}$) para obtener una varianza relativa.
3. Umbral Adaptativo:
   • Si la varianza relativa cae por debajo de un umbral $\tau$ especificado por el usuario, ASL asume que las puntuaciones de atención se han estabilizado en un subconjunto fijo de tokens.
   • En ese momento, la capa actual se convierte en la capa de selección.
   • Se seleccionan los tokens (top-k) y se propagan a todas las capas subsiguientes.

Integración y Eficiencia

• Fase de Prefilling: ASL opera durante el prellenado con una sobrecarga mínima, almacenando solo las puntuaciones de atención agrupadas de las últimas capas.
• Fase de Decodificación: Una vez seleccionados los tokens, la inferencia continúa solo con ellos.
• Compatibilidad: ASL se puede combinar con métodos existentes para cumplir con presupuestos de memoria específicos:
  • Con SnapKV para optimizar la compresión antes de la capa de selección.
  • Con GemFilter para una estrategia de dos pasos (two-pass) que mejora la precisión.

3. Contribuciones Clave

1. Identificación de la Rigidez: Se demuestra empíricamente que la selección de capas fijas es subóptima y que la dificultad de la tarea dicta la profundidad necesaria para una selección precisa.
2. Propuesta de ASL: Un nuevo método adaptativo que no requiere fine-tuning, utilizando la varianza de los rangos de atención como indicador de estabilidad semántica para decidir cuándo podar.
3. Integración Modular: ASL funciona como una capa de control que se integra con métodos de reducción de caché existentes (SnapKV, GemFilter) para optimizar tanto el prellenado como la decodificación.
4. Evaluación Exhaustiva: Validación en tres benchmarks de contexto largo (InfiniteBench, RULER, NIAH) con modelos de hasta 256k de longitud de contexto.

4. Resultados Experimentales

Los experimentos se realizaron en modelos como Llama-3.1-8B-UL y Qwen2.5-7B.

• Precisión (Accuracy):
  • ASL supera consistentemente a los métodos de estado del arte (FastKV, GemFilter, PyramidInfer) en tareas difíciles como la recuperación de pares clave-valor (KV retrieval) y NIAH (Needle in a Haystack).
  • En InfiniteBench, ASL logra puntuaciones cercanas a la caché completa (Full KV) en tareas complejas donde FastKV falla drásticamente.
  • En RULER (128k context), ASL mantiene una ventaja significativa sobre FastKV, especialmente en contextos largos.
• Eficiencia (TTFT y TPOT):
  • TTFT (Tiempo hasta el primer token): ASL es ligeramente más lento que FastKV en tareas fáciles porque a menudo selecciona tokens en capas más profundas (necesarias para tareas difíciles). Sin embargo, ofrece un mejor equilibrio precisión-velocidad.
  • TPOT (Tiempo por token de salida): Una vez seleccionados los tokens, la velocidad de decodificación es comparable a otros métodos de poda, ya que todos procesan el mismo número reducido de tokens en las capas profundas.
• Uso de Memoria:
  • La sobrecarga de memoria de ASL es insignificante (aprox. 0.016 veces el tamaño de la caché completa), manteniendo un perfil de memoria similar a otros métodos de poda.

5. Significado e Impacto

El trabajo ASL representa un avance significativo en la inferencia eficiente de LLMs de contexto largo al introducir conciencia de la tarea en el proceso de poda de tokens.

• Adaptabilidad Dinámica: En lugar de un enfoque "talla única", ASL adapta la estrategia de compresión a la complejidad semántica de cada entrada específica.
• Solución al Dilema Precisión-Velocidad: Permite a los usuarios ajustar el umbral $\tau$ para equilibrar la velocidad de inferencia (selección temprana) con la precisión (selección tardía), logrando un mejor compromiso que los métodos estáticos.
• Viabilidad Práctica: Al ser un método sin entrenamiento y compatible con infraestructuras existentes, ASL es una solución práctica para desbloquear el potencial de los LLMs en aplicaciones de contexto ultra-largo sin requerir reentrenamiento costoso de modelos.

En resumen, ASL demuestra que la estabilidad de los patrones de atención es un indicador fiable para determinar cuándo es seguro reducir la información en la caché KV, superando las limitaciones de los enfoques basados en capas fijas.