SigmaScale: LLM Compression with SVD-based Low-Rank Decomposition and Learned Scaling Matrices

Resumen Técnico: SigmaScale

Planteamiento del Problema

Los Modelos de Lenguaje de Gran Escala (LLMs) han logrado un rendimiento notable mediante el escalado de parámetros a decenas o cientos de miles de millones, pero este crecimiento crea barreras significativas respecto a la accesibilidad tecnológica y el impacto ambiental debido al alto consumo de energía en la inferencia. Aunque existen técnicas de compresión de modelos como la cuantización, la poda (pruning) y la destilación de conocimiento, muchas requieren hardware especializado o procesos de reentrenamiento complejos. La descomposición de bajo rango, específicamente la Descomposición en Valores Singulares Truncada (SVD), ofrece una alternativa agnóstica al hardware al aproximar las matrices de pesos $W$ como el producto de matrices de menor rango.

Sin embargo, la SVD estándar enfrenta dos limitaciones primarias en el contexto de los LLMs:

1. Suboptimalidad Analítica: El teorema de Eckart–Young–Mirsky proporciona la solución analítica óptima para minimizar la norma de Frobenius $||W - W'||_F$. Sin embargo, esta minimización estructural suele tener un desempeño deficiente en los pesos de las redes neuronales porque ignora los valores atípicos (outliers) de las activaciones y el comportamiento funcional de la red.
2. Desafíos de Optimización: La SVD de forma ingenua es computacionalmente costosa ($O(n^3)$) y posee derivadas inestables, lo que dificulta su integración en rutinas de optimización basadas en gradientes. Además, los métodos previos que intentan mejorar la SVD mediante matrices de escalado (por ejemplo, para absorber valores atípicos de activación) dependen de derivaciones analíticas (como la descomposición de Cholesky de la covarianza de activación) en lugar de aprender estas transformaciones directamente.

Metodología: SigmaScale

Los autores proponen SigmaScale, un método que aprende matrices de escalado auxiliares $S$ para ayudar a la compresión basada en SVD truncada. En lugar de derivar las matrices de escalado analíticamente, SigmaScale las optimiza mediante descenso de gradiente bajo una pérdida consciente de la activación.

El proceso consta de tres etapas principales:

1. Sondeo de Sensibilidad: Antes de la compresión, el método determina los rangos de truncamiento óptimos para cada capa y módulo (MLP y atención) para lograr una tasa de compresión global objetivo minimizando el aumento de la perplejidad. Esto utiliza un algoritmo de búsqueda binaria similar a ASVD.
2. Matrices de Escalado Aprendidas:
   • Para una matriz de pesos $W \in \mathbb{R}^{m \times n}$, dos vectores $d_r \in \mathbb{R}^m$ y $d_c \in \mathbb{R}^n$ se inicializan a partir de una distribución Gaussiana escalada.
   • Estos vectores definen matrices de escalado diagonal positivas mediante la exponenciación: $S_r = \text{diag}(\exp(d_r))$ y $S_c = \text{diag}(\exp(d_c))$.
   • La matriz de pesos se escala ($S_r W S_c$), se realiza la SVD truncada para obtener factores de rango-$k$, y se aplica el escalado inverso para reconstruir la matriz comprimida $W'$.
   • Los parámetros $d_r$ y $d_c$ se optimizan para minimizar una función de pérdida consciente de la activación:
     $$L_F = \frac{1}{mn} ||WX - W'X||_F^2$$
     donde $X$ representa las activaciones de un conjunto de calibración. Esto asegura que la compresión preserve la funcionalidad del modelo en lugar de solo su estructura.
   • Nota: Para manejar la derivada inestable de la SVD, el bucle de entrenamiento omite pasos de actualización donde las diferencias de valores singulares se acercan a cero para prevenir valores NaN.
3. Ajuste Fino Post-Compresión: Después de la compresión, el modelo se somete a un ajuste fino para realinear los pesos. Los autores comparan el Ajuste Fino Supervisado (SFT) frente a la Destilación de Conocimiento (KD) utilizando un modelo profesor no comprimido para recuperar el rendimiento.

Contribuciones Clave

1. Aprendizaje de Parámetros de Escalado: Hasta donde llega el conocimiento de los autores, este es el primer trabajo en explorar el aprendizaje de los parámetros de las matrices de escalado por filas y columnas $S$ específicamente para mejorar el rendimiento de la compresión SVD, alejándose de las derivaciones analíticas.
2. Análisis de Recuperación de Rendimiento: El artículo proporciona un análisis comparativo de KD frente a SFT para la recuperación de rendimiento en el contexto de la compresión basada en SVD, utilizando diversos conjuntos de datos post-compresión.
3. Dataset Personalizado: Los autores introducen una variante personalizada del dataset Alpaca, generada utilizando salidas de instrucciones de Llama 3.1-8B, para introducir varianza en los experimentos de realineación de pesos.

Resultados y Análisis

Los experimentos se realizaron en modelos Llama 3.1 8B Instruct y Qwen3-8B a través de varias tasas de compresión (retención de 0.90x, 0.75x, 0.50x).

• Competitividad: SigmaScale es competitivo con los métodos basados en SVD de vanguardia (SVD-LLM y ASVD+).
  • En la retención de 0.90x, SigmaScale mejora sustancialmente la perplejidad sobre SVD-LLM para ambos modelos, recuperando la mayor parte del rendimiento en benchmarks de zero-shot.
  • En la retención de 0.75x, SigmaScale generalmente mejora varios benchmarks de zero-shot, aunque las ganancias en perplejidad son marginales.
  • En la retención de 0.50x, el rendimiento se degrada drásticamente, particularmente para Llama 3.1. Los autores señalan que el método es más efectivo cuando el subespacio retenido es lo suficientemente grande como para que el escalado aprendido compense las direcciones singulares descartadas; no es una solución completa para la compresión de rango extremadamente bajo.
• KD vs. SFT: Contrario a la literatura previa que sugiere que KD es superior para la recuperación, los resultados de los autores no muestran mejoras significativas de KD sobre SFT para SigmaScale en su configuración específica.
• Reducción de Rango Intrínseco: Los autores analizan la "entropía de rango efectivo" de los valores singulares ($\Sigma$) durante la optimización. Observan una fuerte correlación (aprox. 0.81–0.92) entre la reducción en la pérdida de compresión y la reducción en la entropía de rango efectivo. Esto sugiere que el escalado aprendido reduce el rango intrínseco efectivo de las matrices de pesos, alineándolas mejor con las restricciones de compresión.

Significancia y Limitaciones

El artículo posiciona a SigmaScale como una opción válida para aplicaciones que requieren reducir los costos de inferencia de los LLM, particularmente en regímenes de compresión leve a moderada. Al adaptarse a la estructura de los pesos individuales del modelo mediante transformaciones aprendidas y conscientes de la activación, ofrece una ruta más flexible hacia la compresión de bajo rango que los métodos analíticos.

Los autores reconocen explícitamente varias limitaciones:

• Costo Computacional: El método aún requiere calcular la SVD en cada paso de actualización durante el entrenamiento de la matriz de escalado, manteniendo la complejidad $O(n^3)$, y no explora alternativas de SVD aproximadas más rápidas.
• Compresión Extrema: El método se degrada significamente bajo una compresión agresiva (por ejemplo, 0.50x), lo que indica que es un mecanismo para mejorar la calidad del truncamiento en el régimen de rango retenido más que una solución para la compresión extrema.
• Alcance de la Evaluación: La evaluación se limita a la perplejidad y benchmarks específicos de zero-shot; no se exploraron los efectos en la generación de textos largos o tareas de programación.
• Datos de Calibración: No se ha verificado formalmente la robustez ante diferentes distribuciones de calibración, y se señala que el uso de Wikitext es una elección subóptima realizada principalmente para mantener la consistencia con las líneas base previas.