Resumen Técnico: CLHA – Agregación de Hessiano entre Capas para la Cuantización de Transformers con Pesos Compartidos entre Capas (Cross-Layer Weight-Shared)

1. Declaración del Problema

Las arquitecturas recientes de Transformers de Mezcla de Expertos (MoE) eficientes, como las de Megrez2, emplean Compartición de Pesos entre Capas (CLWS). En este diseño, capas adyacentes del transformador (por ejemplo, $\ell$ y $\ell+1$) reutilizan los mismos tensores de pesos de los expertos ruteados para reducir a la mitad la huella de parámetros, manteniendo el mismo número de parámetros activados por token.

Sin embargo, los métodos actuales de Cuantización Post-Entrenamiento (PTQ), incluyendo enfoques de tipo GPTQ, tratan estos pesos compartidos como si pertenecieran a una sola capa. Estos calibran la cuantización utilizando estadísticas de activación de solo una "capa primaria". Esto crea un desajuste fundamental:

• El flujo residual de la representación evoluciona entre capas, lo que significa que las distribuciones de tokens que activan el experto compartido en la capa $\ell_a$ difieren de las de la capa $\ell_b$.
• Los métodos existentes ignoran las estadísticas de activación de la capa secundaria, lo que conduce a una minimización subóptima del error de reconstrucción, particularmente en regímenes de bajos bits (por ejemplo, 2 bits).

2. Metodología

2.1 Fundamento Teórico: Agregación de Hessiano entre Capas

El artículo formaliza el objetivo para cuantizar un peso $W$ compartido a través de las capas $\ell_1, \dots, \ell_n$. El objetivo correcto es minimizar el error de reconstrucción total en todas las capas utilizando el peso:
$$L(W_q) = \sum_{i} \alpha_i \text{tr}[(W_q - W)H_{\ell_i}(W_q - W)^\top]$$
donde $H_{\ell_i}$ es el Hessiano (covarianza de activación) de la capa $\ell_i$, y $\alpha_i$ representa la frecuencia de ruteo (proporción de tokens) para esa capa.

La Proposición 1 establece que minimizar este error combinado es matemáticamente equivalente a aplicar GPTQ con un Hessiano combinado:
$$H_{comb} = \sum_{i} \alpha_i H_{\ell_i}$$
Esto contrasta con la Calibración Independiente por Capa (PLIC), que minimiza el error utilizando únicamente el Hessiano de una sola capa ($H_{\ell_a}$).

2.2 Algoritmo CLHA

Los autores proponen CLHA (Agregación de Hessiano entre Capas), un algoritmo de calibración de tres pasadas:

1. Colección: Para cada secuencia de calibración, se acumulan los Hessianos por capa y el conteo de tokens para cada experto activo.
2. Agregación: Para cada par de expertos compartidos $(\ell_a, \ell_b)$, se calcula un Hessiano combinado. El artículo evalúa dos estrategias de ponderación:
   • Ponderación por Frecuencia de Ruteo ($\alpha = n_a / (n_a + n_b)$): Pondera los Hessianos según el conteo real de tokens recibido por cada capa.
   • Ponderación Igualitaria ($\alpha = 0.5$): Actúa como un regularizador, útil cuando los datos de calibración son limitados o las estimaciones del Hessiano de la capa secundaria son ruidosas.
3. Cuantización: Se aplica la actualización secuencial de columnas estándar de GPTQ utilizando el $H_{comb}$ agregado.

2.3 CLHA-MP (Precisión Mixta)

Reconociendo que los expertos compartidos siempre activos reciben la totalidad de los $N$ tokens (a diferencia de los expertos ruteados que reciben una fracción $K/E$), el artículo señala que su sensibilidad de cuantización es $E/K$ veces mayor. CLHA-MP asigna $b+1$ bits a estos expertos compartidos mientras mantiene los expertos ruteados a $b$ bits, incurriendo en una sobrecarga de almacenamiento mínima (10–20% de los parámetros de los expertos).

2.4 Corrección de Implementación: Ambigüedad de Hooks

El artículo identifica un error de ingeniería crítico: los hooks de forward estándar de PyTorch adjuntos a un módulo ExpertFFN compartido se activan para cualquier capa que llame a dicho módulo. Esto causa que las activaciones de diferentes capas se mezclen en un único acumulador de Hessiano, corrompiendo la atribución por capa.

• Solución: Los autores proponen un mecanismo de callback sensible a la capa pasado a través de la firma MoELayer.forward(). Esto proporciona explícitamente los índices de capa al callback, asegurando que se acumulen Hessianos separados para $H_a$ y $H_b$, incluso cuando el módulo de peso subyacente es compartido.

3. Contribuciones Clave

1. Proposición Teórica: Formalizó que el Hessiano óptimo para pesos compartidos es la suma ponderada por la frecuencia de ruteo de los Hessianos por capa.
2. Algoritmo CLHA: Una implementación práctica de la calibración de Hessiano combinado dentro del marco de trabajo de GPTQ.
3. CLHA-MP: Una extensión de precisión mixta que asigna un bit adicional a los expertos compartidos altamente sensibles.
4. Corrección de Implementación: Identificó y resolvió el problema de la "ambigüedad de hooks" que corrompe silenciosamente la calibración en arquitecturas de módulos compartidos.
5. Validación Empírica: Demostró que las ganancias de CLHA se correlacionan con la divergencia de ruteo entre capas.

el 4. Resultados Experimentales

Los experimentos se realizaron en un MoE de 34.5M de parámetros entrenado desde cero con WikiText-2.

• Rendimiento a 2 bits: Este es el régimen crítico donde la mala calibración se potencia.
  • CLHA-rt (ponderado por ruteo) redujo la Perplejidad (PPL) en un 4.3% en comparación con PLIC.
  • CLHA-eq (ponderación igualitaria) redujo la PPL en un 5.4%.
  • CLHA-MP logró una mejora del 18.6% sobre PLIC.
  • En 4 bits y 3 bits, las diferencias estuvieron dentro de los márgenes de ruido (<1%).
• Calibración vs. Error de Prueba: Curiosamente, PLIC logró un menor Error de Hessiano de Calibración (CHE) en el conjunto de entrenamiento porque optimiza directamente para esa distribución. Sin embargo, CLHA generalizó mejor en el conjunto de prueba, indicando que el Hessiano combinado captura mejor la distribución real del peso compartido que la visión de una sola capa.
• Experimento de Agudeza de Ruteo (Routing Sharpness): Al agudizar artificialmente los logits del ruteador (aumentando la divergencia entre capas), los autores observaron:
  • El rendimiento de PLIC se degradó un 5% a medida que aumentaba la divergencia.
  • CLHA se mantuvo estable (<1% de degradación).
  • La brecha de rendimiento entre CLHA y PLIC creció de un +4.3% (ruteo natural) a un +8.0% (alta divergencia), confirmando la predicción teórica de que las ganancias de CLHA escalan con la divergencia entre capas.

5. Significado y Reivindicaciones

El artículo afirma proporcionar el primer análisis de PTQ específicamente para la compartición de pesos entre capas en modelos MoE. Su importancia radica en:

• Corrección del Objetivo: Demostrar que tratar los pesos compartidos como entidades de una sola capa es matemáticamente subóptimo y que la agregación de Hessiano es la corrección necesaria.
• Escalabilidad de las Ganancias: Mostrar que a medida que los modelos se vuelven más grandes y los patrones de ruteo más especializados (aumentando la divergencia entre capas), los beneficios de CLHA probablemente crecerán, convirtiéndolo en esencial para futuros modelos CLWS-MoE de producción (ej. Megrez2, DeepSeek, Qwen).
• Necesidad Práctica: Destacar que la aplicación ingenua de las herramientas de PTQ existentes a arquitecturas CLWS está fundamentalmente rota debido a la ambigüedad de los hooks, y proporcionar la corrección necesaria para cualquier método de calibración futuro en este dominio.

Los autores mantienen la modestia respecto al rendimiento absoluto, señalando que su modelo es pequeño (34.5M) y fue entrenado durante pasos limitados, pero enfatizan que las mejoras relativas sobre la calibración estándar son robustas y están fundamentadas teóricamente.