LEAP: Learnable End-to-End Adaptive Pruning of Large Language Models

Resumen Técnico: LEAP: Poda Adaptativa Aprendizable de Extremo a Extremo de Modelos de Lenguaje Grandes

Enunciado del Problema
El despliegue de Modelos de Lenguaje Grandes (LLM) está cada vez más limitado por restricciones de memoria y computación. Aunque la dispersión de pesos ofrece una solución, el panorama se divide en regímenes estructurados, semiestructurados (por ejemplo, 2:4) y no estructurados. Los avances recientes en núcleos de GPU (por ejemplo, FlashLLM, SpInfer) y hardware de flujo de datos han desplazado el cuello de botella de ejecutar la dispersión no estructurada a inducirla con una pérdida mínima de precisión.

Los métodos actuales más avanzados para la poda no estructurada, como Wanda, SparseGPT, Thanos y ADMM, dependen de sustitutos por capas derivados del principio del Cirujano Cerebral Óptimo (OBS). Estos métodos minimizan los errores de reconstrucción locales en lugar de la pérdida del modelo de extremo a extremo. En consecuencia, acumulan errores locales que se compounding en redes profundas, lo que lleva a una degradación significativa de la precisión, particularmente bajo dispersión agresiva (50–60%).

Por el contrario, los métodos de máscara aprendible de extremo a extremo como MaskLLM y PATCH optimizan con éxito las máscaras directamente contra la pérdida de modelado de lenguaje, pero están restringidos a patrones semiestructurados. Su parametrización depende de una distribución categórica sobre patrones válidos dentro de un grupo. Aunque es manejable para grupos pequeños (por ejemplo, $\binom{4}{2}=6$ patrones para dispersión 2:4), este enfoque se vuelve combinatoriamente intratable para la dispersión no estructurada. Para una fila de ancho $d=4096$ con una dispersión del 50%, el número de máscaras válidas es $\binom{4096}{2048} \approx 10^{1229}$, lo que hace imposible el almacenamiento y la indexación de los logits de patrones.

Metodología: LEAP
Los autores presentan LEAP (Poda Adaptativa Aprendizable de Extremo a Extremo), un marco diseñado para hacer manejable el aprendizaje de máscaras no estructuradas de extremo a extremo, reemplazando la parametrización intratable de categorías-sobre-patrones con un enfoque por peso.

• Relajación Bernoulli por Peso: En lugar de asignar un logit a cada patrón válido, LEAP asigna un logit aprendible a cada peso individual. Para una matriz de pesos $W \in \mathbb{R}^{m \times n}$, se introduce una matriz de logits paralela $P \in \mathbb{R}^{m \times n}$. La máscara estocástica $M$ se genera mediante la relajación sigmoide-Gumbel:
  $$M = \sigma\left(\frac{\alpha P + g}{\tau}\right)$$
  donde $g$ es ruido Gumbel, $\sigma$ es la función sigmoide, $\alpha$ es un factor de escala y $\tau$ es un parámetro de temperatura. Esta formulación escala linealmente con el número de pesos ($O(mn)$) en lugar del número de patrones válidos.
• Estrategia de Optimización:
  • Pesos Congelados: Los pesos preentrenados $W$ se mantienen congelados. Solo se entrenan los logits de la máscara $P$. Esto preserva la calibración del modelo preentrenado y simplifica el despliegue.
  • Inicialización: $P$ se inicializa utilizando una máscara Wanda de un solo disparo (estableciendo las entradas seleccionadas en $+s$ y las demás en $-s$) para proporcionar un punto de partida razonable, evitando problemas de inicio en frío.
  • Programaciones: Una programación de escala ($\alpha$) aumenta para empujar la sigmoide hacia valores binarios, mientras que una programación de temperatura ($\tau$) decae para afilar la máscara.
  • Función de Pérdida: El objetivo combina la pérdida de modelado de lenguaje en un flujo de calibración con dos términos de regularización:
    1. Regularizador de Dispersión Global: Impone la densidad objetivo globalmente en todas las capas, permitiendo que las capas individuales ajusten su densidad basándose en la importancia de extremo a extremo.
    2. Término Consciente de la Magnitud: Sesga la optimización hacia la retención de pesos de mayor magnitud para evitar mínimos degenerados donde se descartan pesos críticos.
• Eficiencia de Entrenamiento: El método entrena durante aproximadamente 2.000 iteraciones en un pequeño flujo de calibración de texto general (SlimPajama).

Contribuciones Clave

1. Identificación de la Obstrucción Combinatoria: El artículo identifica formalmente que las parametrizaciones de categorías-sobre-patrones (utilizadas por MaskLLM/PATCH) no pueden escalar a la dispersión no estructurada debido al crecimiento exponencial de los patrones válidos.
2. Reformulación Manejable: Los autores proponen una reformulación Bernoulli por peso vía sigmoide-Gumbel que permite el aprendizaje diferenciable de extremo a extremo para máscaras no estructuradas a escala de LLM.
3. Marco LEAP: Un marco de poda ligero que opera sobre pesos congelados, requiriendo solo un pequeño conjunto de datos de calibración y aproximadamente 2.000 pasos de entrenamiento.
4. Validación Empírica: Evaluación exhaustiva en cinco familias de LLM (Qwen-2.5, Gemma-3, LLaMA-3.2, LLaMA-3.1) que van desde 0,5B hasta 8B de parámetros con dispersión del 50% y 60%.

Resultados
LEAP fue evaluado frente a líneas base por capas (Wanda, SparseGPT, Thanos, ADMM) y la línea base aprendible semiestructurada MaskLLM.

• Ganancias de Precisión: En diez configuraciones (modelo, dispersión), LEAP mejoró la precisión cero-shot promedio de seis tareas en +2,59 puntos sobre ADMM, la mejor línea base por capas.
• Mejoras Específicas: La mayor ganancia fue de +5,40 puntos en LLaMA-3.2 1B con una dispersión del 60% (50,39 frente a 44,99).
• Perplejidad: LEAP logró consistentemente una perplejidad (PPL) más baja que las líneas base. Por ejemplo, en LLaMA-3.1 8B con una dispersión del 50%, LEAP alcanzó 7,66 PPL en comparación con los 9,12 de ADMM.
• Asignación de Dispersión: Los experimentos mostraron que bajo la optimización de extremo a extremo con un presupuesto global, LEAP converge a una asignación de densidad por bloque casi uniforme, en contraste con las variaciones no uniformes observadas en los métodos semiestructurados.
• Costo Computacional: Aunque es más costoso que los métodos de un solo disparo (tomando horas en H100 frente a minutos), LEAP es significativamente más barato que MaskLLM (por ejemplo, 276 horas de GPU para LLaMA-3.1 8B frente a miles de horas de A100 para MaskLLM en escalas comparables).

Significado y Afirmaciones
El artículo afirma que LEAP resuelve la obstrucción fundamental que impide el aprendizaje de máscaras de extremo a extremo en la dispersión no estructurada. Al desacoplar el aprendizaje de la máscara de las actualizaciones de pesos, ofrece una vía práctica para comprimir modelos fundamentales sin sacrificar la calibración de los pesos preentrenados.

Los autores enfatizan que LEAP produce máscaras en el régimen no estructurado del 50%–60%, que son directamente compatibles con los núcleos acelerados existentes (por ejemplo, FlashLLM, SpInfer, MACKO). El significado no radica en crear nuevos núcleos de inferencia, sino en permitir el despliegue de estos núcleos en modelos podados de alta calidad que anteriormente eran inalcanzables mediante métodos por capas. El trabajo sugiere que la dispersión no estructurada, cuando se induce mediante aprendizaje de extremo a extremo, puede ofrecer una compensación favorable entre compresión y precisión para el despliegue en hardware de consumo.

Limitaciones y Alcance

• Pesos Congelados: El alcance actual mantiene los pesos congelados. Aunque la optimización conjunta de pesos y máscaras se señala como una extensión natural, no se implementa en este trabajo.
• Asignación Uniforme: La observación de una densidad por bloque uniforme se presenta como un hallazgo específico del régimen para los modelos estudiados y los niveles de dispersión, no como una regla general contra la asignación no uniforme.
• Seguridad: Los autores advierten que los modelos comprimidos pueden exhibir diferentes modos de fallo en entradas de cola larga y recomiendan una reevaluación para despliegues críticos para la seguridad.