Pruning via Causal Attribution Preserves Reasoning Performance in Large Language Models

Resumen Técnico: Poda por Atribución Causal (CAP)

Declaración del Problema
Los Modelos de Lenguaje Extensos (LLM) demuestran sólidas capacidades de razonamiento de múltiples pasos, pero incurren en costos de inferencia sustanciales, lo que dificulta su despliegue en el mundo real. Si bien la poda de modelos ofrece una solución al eliminar parámetros redundantes sin necesidad de reentrenamiento, los métodos existentes suelen degradar el rendimiento en tareas de razonamiento complejas. Los enfoques actuales dependen de señales correlacionales, como la magnitud de los pesos (por ejemplo, Wanda) o estadísticas de gradiente locales. Estos métodos asumen que los pesos pequeños no son importantes; sin embargo, los parámetros de baja magnitud pueden ser críticos para los circuitos de razonamiento de múltiples pasos. Además, los criterios puramente correlacionales no logran capturar las interacciones entre capas esenciales para la inferencia de cadena de pensamiento (CoT), lo que frecuentemente conduce a un "colapso de capa" donde módulos enteros pierden su función.

Metodología: Poda por Atribución Causal (CAP)
Los autores introducen la Poda por Atribución Causal (CAP), un método de poda de un solo paso y sin entrenamiento diseñado para preservar el rendimiento del razonamiento. A diferencia de los métodos que clasifican los pesos basándose en propiedades estáticas, CAP mide el impacto causal de componentes específicos del modelo en el rendimiento de la tarea. El marco opera en tres etapas:

1. Calibración y Puntuación Causal (Etapa 1):

   • Se construye un pequeño conjunto de datos de calibración ($D_{cal}$) mediante el muestreo de evaluaciones centradas en el razonamiento: GSM8K (matemáticas), StrategyQA (sentido común de múltiples saltos) y ARC-Challenge (ciencia).
   • Para cada cabeza de atención $h_\ell$ en el modelo, CAP realiza un experimento de enmascaramiento intervencional. La cabeza se enmascara temporalmente durante la pasada hacia adelante (forward pass), y se mide el aumento resultante en la pérdida de entropía cruzada a nivel de token ($\Delta(h_\ell)$) con respecto al modelo intacto.
   • Para garantizar la robustez contra la varianza derivada de tamaños de muestra pequeños, el conjunto de datos se divide en $K$ subconjuntos, y se calcula la mediana del aumento de la pérdida.
   • Este proceso identifica las cabezas que son funcionalmente críticas para el razonamiento, distinguiéndolas de aquellas que son meramente grandes numéricamente pero redundantes.

2. Conversión de Importancia a Nivel de Peso (Etapa 2):

   • Las puntuaciones causales a nivel de cabeza se convierten en factores de importancia normalizados $I(h_\ell)$, mapeados a un rango de $[1.0, 10.0]$. Las cabezas de alto impacto reciben factores más altos.
   • Estos factores se propagan a los pesos individuales dentro de sus matrices de proyección correspondientes ($W_q, W_k, W_v, W_o$). La puntuación de importancia para un peso específico $W_{ij}$ se calcula como el producto de su magnitud y el factor de importancia de su región de cabeza asociada: $S_{ij} = |W_{ij}| \cdot I(h_{region})$.
   • Para las capas MLP, que carecen de una estructura natural particionada por cabezas, el método asigna una importancia promedio por capa derivada de las cabezas de atención colocalizadas. Esto se señala como una aproximación.

3. Poda Global (Etapa 3):

   • Todos los pesos en todo el modelo se clasifican por sus puntuaciones de magnitud ponderadas por importancia ($S_{ij}$).
   • Se determina un umbral global mediante búsqueda binaria para alcanzar la dispersión objetivo (por ejemplo, 10%, 20%, 50%). Los pesos por debajo del umbral se eliminan (se vuelven cero), creando un patrón de dispersión no estructurado compatible con la infraestructura de inferencia estándar.

Contribuciones Clave

• Formalización de la Atribución Causal para la Poda: Los autores formalizan un método para medir el aumento esperado en la pérdida específica de la tarea cuando se enmascaran cabezas de atención individuales, yendo más allá de los proxies correlacionales para cuantificar la criticidad funcional.
• Poda a Nivel de Peso Consciente del Presupuesto: Introducen un procedimiento que convierte las puntuaciones causales a nivel de cabeza en valores de importancia a nivel de peso. Esto permite una poda de magnitud fina bajo una restricción de dispersión global mientras se preserva la estructura de los circuitos de razonamiento críticos.
• Validación Empírica: El método se evalúa en Llama-3-8B-Instruct y Mistral-7B-Instruct a través de los benchmarks GSM8K, StrategyQA y ARC-Challenge.

Resultados

• Dispersión Moderada (10–20%): CAP generalmente supera al baseline de Wanda en la preservación del rendimiento del razonamiento.
  • En Llama-3-8B-Instruct con un 20% de dispersión, CAP logra una mejora relativa del 61% sobre Wanda en ARC-Challenge (70.8% vs. 43.9% de precisión).
  • En GSM8K con un 10% de dispersión, CAP muestra una mejora relativa del 20% sobre Wanda (71.5% vs. 59.5%).
  • Los autores señalan que, aunque CAP a veces queda por detrás de Wanda en tareas específicas (por ejemplo, Llama-3 GSM8K al 20%), preserva consistentemente el rendimiento en tareas de inferencia de múltiples pasos donde los criterios correlacionales de Wanda fallan.
• Alta Dispersión (50%): El rendimiento se degrada significativamente para ambos métodos. CAP sufre un fallo catastrófico en Llama-3 (la perplejidad salta a 428.2), atribuido a su tratamiento de grano grueso de las capas MLP. En este nivel, el umbral global elimina pesos MLP esenciales que fueron asignados con baja importancia basándose en la aproximación de las cabezas de atención.
• Perplejidad vs. Razonamiento: En la dispersión moderada, CAP a menudo logra una mejor precisión de razonamiento a pesar de tener una perplejidad ligeramente mayor en comparación con Wanda, lo que sugiere un desacoplamiento entre la calidad del modelado de lenguaje general y la preservación del razonamiento específico de la tarea.
• Limitaciones de la Arquitectura: Experimentos preliminares en modelos de Mezcla de Expertos (MoE) sugieren que CAP es menos efectivo, ya que el enrutamiento dinámico hace que la importancia de las cabezas no sea estacionaria.

Significancia y Reivindicaciones
El artículo sostiene que la atribución causal a nivel de cabeza de atención es un criterio superior para preservar el rendimiento del razonamiento en comparación con los métodos de poda correlacionales a niveles de dispersión equivalentes. Al medir explícitamente la contribución causal de las cabezas a las tareas de razonamiento, CAP puede identificar y proteger parámetros "numéricamente pequeños pero funcionalmente críticos" que los métodos basados en magnitud descartan.

Los autores enfatizan que su trabajo es de alcance modesto:

• Es un método sin entrenamiento, lo que lo hace computacionalmente eficiente en comparación con enfoques que requieren actualizaciones de pesos (como SparseGPT).
• Está alineado con la tarea, calibrado específicamente en datos de razonamiento en lugar de modelado de lenguaje genérico.
• Reconoce limitaciones, particularmente la falta de atribución causal de grano fino para las neuronas MLP y la dificultad de aplicar el método a arquitecturas de enrutamiento dinámico (MoE).

La contribución central es la demostración de que la poda intervencional y alineada con la tarea puede preservar mejor capacidades específicas (razonamiento) que los métodos optimizados para la perplejidad general o las estadísticas estáticas de los pesos, siempre que el nivel de dispersión sea moderado.