MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs

Resumen Técnico: MaskPro – Aprendizaje Probabilístico de Espacio Lineal para Esparsidad Estricta (N:M) en LLMs

Enunciado del Problema
El escalado rápido de los Modelos de Lenguaje Grandes (LLM) ha creado un cuello de botella significativo en la eficiencia de la inferencia, haciendo necesarias técnicas de compresión efectivas. La esparsidad semi-estructurada, específicamente el patrón (N:M) donde se retienen $N$ elementos de cada $M$ pesos consecutivos, ofrece una solución amigable con el hardware para la aceleración y la reducción de memoria. Sin embargo, identificar la máscara (N:M) óptima es un formidable problema de optimización combinatoria. Los enfoques existentes se dividen en dos categorías con limitaciones distintas:

1. Métodos basados en reglas: Estos emplean búsquedas codiciosas, capa por capa, utilizando métricas creadas a mano (por ejemplo, minimizar el error de reconstrucción por capa). Aunque son computacionalmente baratos, sufren de sesgos inherentes y brechas considerables con el rendimiento en tareas de extremo a extremo.
2. Métodos basados en aprendizaje (por ejemplo, MaskLLM): Estos optimizan máscaras directamente mediante descenso de gradiente sobre una pérdida de generación. Aunque son efectivos, incurren en costos de entrenamiento prohibitivos. Específicamente, MaskLLM requiere almacenar logits para todas las máscaras posibles dentro de un grupo de $M$ pesos, lo que conduce a una complejidad de memoria de $O(\binom{M}{N} \frac{d}{M})$. Para $M$ grande, esto escala exponencialmente, haciéndolo inviable para modelos grandes. Además, el vasto espacio combinatorio induce una alta varianza en los gradientes de la política, causando inestabilidad en el entrenamiento.

Metodología: MaskPro
Los autores proponen MaskPro, un nuevo marco probabilístico de espacio lineal diseñado para aprender esparsidad (N:M) de manera eficiente y estable.

1. Relajación Probabilística de Espacio Lineal:
   En lugar de definir una distribución categórica sobre todas las $\binom{M}{N}$ máscaras posibles para un grupo de $M$ pesos, MaskPro reformula el problema. Establece una distribución categórica sobre los $M$ elementos individuales dentro de cada grupo. La máscara (N:M) se genera luego realizando un muestreo N-ario sin reemplazo desde esta distribución.

   • Base Teórica: Los autores demuestran que cualquier máscara (N:M) puede representarse como una suma probabilística de $N$ vectores de base distintos.
   • Eficiencia de Memoria: Esta reformulación reduce el número de parámetros aprendibles (logits) de $O(\binom{M}{N} \frac{d}{M})$ a un espacio lineal $O(d)$, donde $d$ es el número total de pesos. Esto elimina la explosión de memoria asociada con los métodos basados en aprendizaje anteriores.

2. Estimador de Gradiente de Política Refinado (PGE):
   Para optimizar los logits, MaskPro utiliza gradientes de política. Los autores identifican que los gradientes de política estándar que utilizan valores de pérdida crudos son inestables en LLM debido a la alta varianza causada por las diferencias en los mini-lotes (donde la variación de la pérdida por muestreo de datos a menudo supera con creces la variación de la pérdida por muestreo de máscaras).

   • Residuos de Pérdida: En lugar de usar la pérdida cruda $f(m \odot w, \xi)$, MaskPro actualiza los logits basándose en el residuo de pérdida relativo a una máscara de referencia inicial $m_0$: $f(m \odot w, \xi) - f(m_0 \odot w, \xi)$. Esto aísla el impacto del cambio de máscara.
   • Rastreador de Promedio Móvil: Para estabilizar aún más el entrenamiento y prevenir explosiones numéricas, se introduce un rastreador de promedio móvil $\delta$ para suavizar la referencia. La regla de actualización se convierte en:
     $$\pi_{t+1} = \pi_t - \eta (f(m_t \odot w, \xi) - f(m_0 \odot w, \xi) - \delta) \nabla \log p(m_t|\pi_t)$$
   • Garantías Teóricas: El artículo proporciona demostraciones que muestran que estos estimadores refinados permanecen como estimadores sin sesgo del gradiente de política verdadero. Además, bajo la condición de que la pérdida de la máscara muestreada sea menor que el doble de la pérdida de referencia (una condición fácilmente cumplida en la práctica), la varianza del estimador refinado es estrictamente menor que la del gradiente de política estándar.

Contribuciones Clave

1. Marco de Espacio Lineal: MaskPro reduce el requisito de memoria para aprender esparsidad (N:M) de exponencial $O(\binom{M}{N} \frac{d}{M})$ a lineal $O(d)$ al modelar el proceso de selección como un muestreo N-ario sin reemplazo desde una distribución categórica sobre $M$ elementos.
2. Gradiente de Política Mejorado: Los autores introducen un mecanismo de actualización de varianza reducida utilizando residuos de pérdida y un rastreador de promedio móvil, abordando la inestabilidad del entrenamiento causada por la alta varianza de los gradientes de política en vastos espacios combinatorios.
3. Análisis Teórico: Se proporcionan demostraciones exhaustivas sobre la ausencia de sesgo de los estimadores y las propiedades de reducción de varianza de las reglas de actualización propuestas.

Resultados Experimentales
Se realizaron experimentos extensos en múltiples LLM (Gemma-7B, Vicuna-7B, LLaMA-2-7B, DeepSeek-7B) en diversas tareas posteriores (por ejemplo, HellaSwag, RACE, PIQA) y patrones de esparsidad ((2:4), (4:8), (8:16)).

• Rendimiento: MaskPro supera consistentemente a los métodos basados en reglas más avanzados (SparseGPT, Wanda, Pruner-Zero) y logra un rendimiento comparable al MaskLLM intensivo en memoria, a menudo con una mejora promedio superior al 2% en precisión.
• Escalabilidad de Memoria: Mientras que MaskLLM encuentra errores de memoria insuficiente (OOM) en 8x GPUs A100 para esparsidad (4:8) en modelos de 7B, MaskPro mantiene un uso de memoria estable (~36-37 GB) independientemente del patrón de esparsidad.
• Eficiencia de Datos: MaskPro demuestra una notable robustez ante el volumen de datos. Logra un rendimiento estable incluso cuando se entrena con solo 1 muestra de datos, whereas MaskLLM requiere cientos de miles de muestras para converger.
• Eficiencia de Entrenamiento: El método requiere significativamente menos muestras de entrenamiento y menos tiempo de entrenamiento en comparación con MaskLLM, mientras mantiene un uso de memoria comparable a los métodos basados en reglas.

Significado y Afirmaciones
El artículo afirma que MaskPro proporciona una vía viable para el despliegue local práctico y escalable de LLM al resolver los cuellos de botella de memoria y estabilidad del aprendizaje de esparsidad semi-estructurada basada en aprendizaje. Los autores enfatizan que su enfoque cierra la brecha entre el alto rendimiento de los métodos basados en aprendizaje y la eficiencia de los métodos basados en reglas. Al reducir la huella de memoria a espacio lineal y estabilizar el proceso de entrenamiento mediante gradientes de política refinados, MaskPro permite el aprendizaje de máscaras (N:M) óptimas en modelos a gran escala sin la sobrecarga computacional prohibitiva asociada previamente con tales tareas. El método se presenta como un marco generalizable para el aprendizaje probabilístico de máscaras que es robusto ante la escasez de datos y la escala del modelo.