Resumen Técnico: Poda Guiada por Atribución y de Cobertura Maximizada para la Compresión Estructural de MoE

1. Planteamiento del Problema

Los modelos de Mezcla de Expertos (MoE) ofrecen una alta capacidad de parámetros con una computación manejable al activar solo un subconjset de expertos por token. Sin embargo, el despliegue de estos modelos sigue siendo costoso debido a las sustanciales huellas de memoria y la sobrecarga de inferencia. Si bien la poda estructural es una solución prometedora, los métodos existentes enfrentan limitaciones significativas:

• Granularidad Gruesa: Los enfoques previos operan al nivel de experto, ya sea eliminando expertos completos o clasificándolos por importancia. Esto no logra capturar la redundancia interna dentro de los expertos de "alta importancia", donde la información suele concentrarse en solo unos pocos canales.
• Estimación de Importancia Inexacta: Las métricas existentes para la importancia de los expertos (por ejemplo, pesos del router, uso de tokens, estadísticas brutas de activación/gradiente) a menudo están desalineadas con la contribución real de un experto a la pérdida final. Las estadísticas del router reflejan la frecuencia de selección en lugar de la utilidad, y las estadísticas brutas exhiben sesgos dependientes de la capa que hacen que la comparación entre capas sea poco fiable.
• Asignación de Presupuesto Ineficiente: La asignación de presupuestos de poda basados en clasificaciones de expertos gruesas conduce a un desperdicio de capacidad en estructuras de baja contribución y no logra optimizar la relación entre compresión y precisión en MoEs de gran escala con cientos de expertos.

2. Metodología

Los autores proponen un marco de poda estructural diseñado para MoE, que consta de tres componentes principales: Aproximación de Pérdida Guiada por Atribución (ALA), Asignación de Presupuesto de Cobertura Maximizada (CBA) y Redistribución Consciente de la Alineación (AAR).

2.1 Aproximación de Pérdida Guiada por Atribución (ALA)

Para abordar el alto costo de la ablación por experto (que escala linealmente con el número de expertos), los autores derivan un proxy eficiente para la importancia del experto utilizando una expansión de Taylor de primer orden.

• Derivación: El cambio en la pérdida inducido por la eliminación de un experto $e$ se aproxima como $\Delta L^{(e)} \approx -(\frac{\partial L}{\partial z_{\ell,e}})^\top z_{\ell,e}$, donde $z_{\ell,e}$ es la salida del experto y $\frac{\partial L}{\partial z_{\ell,e}}$ es el gradiente con respecto a esa salida.
• Eficiencia: Esto permite la estimación de la importancia para todos los expertos en una capa mediante una única pasada hacia atrás (backward pass), reduciendo el tiempo de calibración entre 14 y 26$\times$ en comparación con la ablación exhaustiva.

2.2 Asignación de Presupuesto de Cobertura Maximizada (CBA)

Motivados por la observación de que la información en los expertos de MoE está altamente concentrada en una pequeña fracción de los canales, los autores reformulan el objetivo de poda. En lugar de asignar una razón de poda fija por experto, el objetivo es maximizar la cobertura de la puntuación de canales acumulada bajo un presupuesto global.

• Mecanismo: Los canales dentro de cada experto se clasifican por sus puntuaciones (derivadas de ALA). El algoritmo busca una razón de cobertura objetivo $\rho$ para cada experto de tal manera que el número total de canales retenidos cumetre con el presupuesto de poda global mientras se maximiza la suma de las puntuaciones de los canales retenidos.
• Beneficio: Este enfoque prioriza la retención de canales de alta contribución y poda agresivamente las colas de baja puntuación, evitando el desperdicio de capacidad en estructuras redundantes que incurren los métodos uniformes o basados en ratios.

2.3 Redistribución Consciente de la Alineación (AAR)

Para asegurar la compatibilidad con la cuantización de bajos bits y los kernels eficientes de hardware (por ejemplo, GEMM), las dimensiones de los canales deben ser múltiplos de un tamaño de bloque $a$ (por ejemplo, 64 o 128).

• Proceso: Después de que CBA determina el recuento óptimo de canales, AAR realiza una alineación hacia abajo al múltiplo más cercano de $a$. Los canales "remanentes" liberados por este redondeo se redistribuyen a los expertos con los mayores remanentes fraccionarios (usando la regla del mayor remanente de Hamilton).
• Objetivo: Esto asegura que el modelo podado satisfaga las restricciones del kernel sin una desviación significativa de la asignación óptima basada en la cobertura, evitando la degradación del rendimiento debido al padding o al acceso ineficiente a la memoria.

3. Contribuciones Clave

1. Observación de Redundancia Interna: El artículo demuestra que la información de MoE está altamente concentrada en una pequeña fracción de los canales, lo que hace que las métricas de importancia a nivel de experto sean demasiado gruesas para una compresión estructural efectiva.
2. Nuevo Objetivo de Poda: Los autores introducen la cobertura de puntuación de canales como un objetivo de poda, reformulando la asignación de capacidad para maximizar la cobertura bajo un presupuesto global. Es el primer trabajo que optimiza explícitamente la cobertura de puntuación en lugar de simples ratios de poda.
3. Marco Eficiente: Proponen un marco escalable que combina:
   • ALA: Una aproximación basada en atribución para la estimación de la importancia de los expertos que reduce las horas de GPU en 20$\times$ en comparación con la ablación.
   • CBA: Un algoritmo basado en búsqueda binaria para asignar presupuestos que maximicen la cobertura de puntuación.
   • AAR: Una estrategia de redistribución para satisfacer las restricciones de los kernels de bajos bits.
4. Validación Empírica: Experimentos extensos en arquitecturas DeepSeek y Qwen MoE muestran que el método logra más de 5$\times$ de compresión con una pérdida de precisión insignificante en tareas de conocimiento general y mantiene un fuerte rendimiento en tareas de razonamiento.

4. Resultados Experimentales

El marco fue evaluado en DeepSeek-MoE-16B, DeepSeek-V2-Lite, Qwen1.5-MoE-A2.7B y Qwen3-30B-A3B.

• Conocimiento General: Bajo una razón de poda del 25% combinada con cuantización de 4 bits ($P25\% Q4b$), el método logró una reducción de almacenamiento de más de 5$\times$. En Qwen3-30B-A3B, la huella de memoria se redujo en 5.27$\times$ (de 61.06GB a 11.64GB) con una caída de precisión promedio de máximo 1% en benchmarks generales (por ejemplo, MMLU, ARC).
• Tareas de Razonamiento: En benchmarks de razonamiento (GSM8K, MATH500), los modelos comprimidos consistentemente se acercaron o superaron a sus contrapartes originales. Por ejemplo, Qwen3-30B-A3B bajo una poda del 50% retuvo una puntuación de 94.5 en MATH500 (comparado con el 92.8 de la base).
• Eficiencia: El método superó a las líneas base de vanguardia (incluyendo Wanda, MoNE, EAC-MoE y PuzzleMoE) a través de diversos benchmarks. También demostró un rendimiento robusto bajo diferentes estrategias de enrutamiento (top-1 a top-6) y corpora de calibración.
• Throughput: La redistribución consciente de la alineación mejoró significativamente el rendimiento de inferencia. Para Qwen3-30B-A3B, las dimensiones alineadas (128) produjeron 14.21 tokens/s en comparación con 10.23 tokens/s para dimensiones no alineadas bajo un 50% de dispersión.

5. Significado y Reivindicaciones

El artículo afirma proporcionar un camino práctico hacia el despliegue eficiente de MoE al abordar el desajuste entre la importancia a nivel de experto y la redundancia de canales interna.

• Escalabilidad: Al reemplazar la ablación exhaustiva con la aproximación basada en atribución, el método escala a MoEs con cientos de expertos, un régimen donde los métodos previos basados en pérdida eran computacionalmente prohibitivos.
• Optimización de Grano Fino: El cambio de la poda a nivel de experto a la maximización de la cobertura a nivel de canal permite una asignación de capacidad más precisa, preservando la información crítica mientras se elimina el ruido.
• Compatibilidad de Hardware: La integración de la redistribución consciente de la alineación asegura que los modelos comprimidos no solo sean más pequeños, sino también compatibles con la cuantización de bajos bits y los kernels de alto rendimiento, abordando un cuello de botella crítico en el despliegue del mundo real.

Los autores posicionan este trabajo como un paso fundamental para mover la compresión de MoE desde decisiones gruesas a nivel de experto hacia una optimización de grano fino a nivel de canal, permitiendo el despliegue de modelos MoE a gran escala con requisitos de memoria y cómputo significativamente reducidos.