POP: Prefill-Only Pruning for Efficient Large Model Inference

1. El Problema: Coste Computacional y Limitaciones de la Poda Existente

Los Grandes Modelos de Lenguaje (LLMs) y los Modelos Visión-Lenguaje (VLMs) han demostrado capacidades excepcionales, pero su despliegue está limitado por los altos costes computacionales durante la inferencia.

• Ineficiencia de la Poda Estructurada: Los métodos actuales de poda estructurada (que eliminan capas o canales completos) suelen ofrecer compatibilidad con el hardware, pero sufren de una degradación significativa de la precisión, especialmente en tareas generativas abiertas.
• Enfoque "Agnóstico a la Etapa": El fallo principal de los métodos existentes es tratar el proceso de inferencia como un todo uniforme. Ignoran la asimetría funcional entre las dos etapas clave de la inferencia autoregresiva:
  1. Prefill (Relleno): Procesa el historial de entrada en paralelo para codificarlo en la memoria caché de claves y valores (KV Cache). Su objetivo es solo proporcionar contexto.
  2. Decode (Decodificación): Genera token a token. Debe codificar el token actual en la caché y simultáneamente modelar la distribución de probabilidad del siguiente token.
• Hipótesis: Los autores argumentan que las capas profundas son críticas para la predicción del siguiente token (Decode) pero redundantes para la codificación del contexto (Prefill). Los métodos actuales, al podar las mismas capas en ambas etapas, destruyen la capacidad de generación.

2. Metodología: Prefill-Only Pruning (POP)

El artículo propone POP, una estrategia de inferencia consciente de la etapa que poda selectivamente las capas profundas solo durante la fase de prefill, manteniendo el modelo completo para la fase de decode.

A. Mecanismo de Puerta Virtual (Virtual Gate)

Para identificar qué capas son redundantes en cada etapa, los autores introducen un mecanismo de "puerta virtual":

• Se multiplica la salida de las ramas residuales (Atención y FFN) de cada capa por un parámetro escalar virtual $g_l$.
• Si $g_l = 1$, la capa funciona normalmente; si $g_l = 0$, la capa se poda (actualización residual suprimida).
• Se estima la importancia de la capa calculando el aumento de pérdida (loss) cuando $g_l$ pasa de 1 a 0, utilizando una expansión de Taylor de segundo orden y la matriz de información de Fisher para aproximar la derivada de forma eficiente sin necesidad de múltiples pasadas de inferencia.

B. Análisis de Importancia Sensible a la Etapa

Al aplicar este mecanismo por separado para las etapas de prefill y decode, se descubrió una disparidad notable (visualizada en la Figura 1 del paper):

• Decode: Las capas profundas son extremadamente importantes (alta sensibilidad).
• Prefill: Las capas profundas tienen una importancia muy baja (cercana a cero), indicando redundancia.
• Conclusión: Se puede eliminar la parte inferior de la red (últimas capas) durante el prefill sin afectar significativamente la codificación del contexto.

C. Implementación Técnica de POP

Para ejecutar esta estrategia asimétrica, se implementan dos mecanismos clave:

1. Proyecciones KV Independientes:
   • Si se salta una capa durante el prefill, no se generarían las claves y valores (KV) necesarios para la etapa de decode (que usa el modelo completo).
   • Solución: Se desacopla la proyección KV del bloque computacional pesado. Para las capas podadas en el prefill, se ejecutan solo las proyecciones lineales ($W^K, W^V$) para generar y almacenar el KV Cache, saltándose los cálculos costosos de Atención y FFN. Esto mantiene la integridad del caché con un coste computacional mínimo (<5%).
2. Manejo de la Frontera (Boundary Handling):
   • El último token de entrada ($x_N$) es crítico para predecir el primer token generado. Si se poda en el prefill, la precisión cae drásticamente.
   • Solución: Se redefine la frontera. El prefill podado procesa $x_1$ a $x_{N-1}$. El procesamiento de $x_N$ se trata como el primer paso de decode, utilizando el modelo completo. Esto garantiza que la predicción inicial sea precisa.

3. Contribuciones Clave

1. Descubrimiento de Asimetría: Demostración empírica de que la redundancia de las capas en los LLMs es dependiente de la etapa (las capas profundas son redundantes en el prefill pero esenciales en el decode).
2. Estrategia POP: Propuesta de un método de poda que acelera la etapa intensiva en cómputo (prefill) mientras preserva la capacidad total del modelo para la generación sensible (decode).
3. Mecanismos de Transición: Introducción de proyecciones KV independientes y manejo de fronteras para asegurar la integridad del caché y la precisión del primer token.
4. Evaluación Exhaustiva: Validación en múltiples familias de modelos (Llama-3.1, Qwen3-VL, Gemma-3) y modalidades (texto, visión-lenguaje).

4. Resultados Experimentales

Los experimentos se realizaron en GPUs NVIDIA A100 con modelos de texto y visión-lenguaje.

• Precisión (Accuracy):
  • POP vs. Poda Estructurada: Los métodos existentes (SliceGPT, ShortGPT) sufrieron un colapso catastrófico en tareas de generación abierta (ej. GSM8K, HumanEval), cayendo a menos del 1% de precisión en algunos casos. POP mantuvo una precisión casi idéntica al modelo completo (ej. retuvo el 97% del rendimiento en GSM8K para Llama-3.1).
  • POP vs. Poda No Estructurada: POP logró una precisión comparable a Wanda (un método no estructurado), pero con la ventaja de la compatibilidad de hardware estructurado.
• Velocidad (Speedup):
  • Prefill (TTFT - Time to First Token): POP logró aceleraciones de hasta 1.37x en secuencias largas (2048 tokens) y resoluciones de imagen altas.
  • Dependencia de la longitud: La aceleración es mayor en contextos largos, donde el coste computacional del prefill domina. En secuencias muy cortas, la ganancia es menor debido al manejo de la frontera (el último token requiere el modelo completo).
  • Comparativa: POP superó consistentemente a SliceGPT y ShortGPT en velocidad sin sacrificar la precisión.

5. Significado y Limitaciones

Significado:
POP representa un cambio de paradigma en la optimización de inferencia. En lugar de buscar una arquitectura reducida estática, propone una estrategia dinámica y consciente de la etapa. Esto permite superar la clásica compensación (trade-off) entre eficiencia y precisión en la poda estructurada, ofreciendo una solución "plug-and-play" que no requiere reentrenamiento ni hardware especializado.

Limitaciones:

1. Uso de Memoria (VRAM): A diferencia de la poda tradicional que reduce el tamaño del modelo en memoria, POP requiere cargar los pesos completos para la etapa de decode. Por lo tanto, no reduce el pico de uso de VRAM; es ideal para escenarios limitados por cómputo, no por capacidad de memoria.
2. Integración en Sistemas Distribuidos: La implementación actual asume un pipeline monolítico. Futuras investigaciones podrían integrar POP en sistemas de inferencia desacoplados (como DistServe o Splitwise) que ejecutan prefill y decode en hardware separado, maximizando aún más el rendimiento del clúster.

En resumen, POP demuestra que la poda inteligente basada en la etapa de inferencia puede ofrecer aceleraciones significativas en la latencia de inicio (prefill) manteniendo la calidad de generación intacta, resolviendo uno de los cuellos de botella principales en la implementación de LLMs y VLMs.