Resumen Técnico: La sensibilidad del enrutador bajo el ajuste fino ligero identifica expertos podables en modelos de Mezcla de Expertos (MoE)

Declaración del Problema

Los modelos de Mezcla de Expertos (MoE) desacoplan el recuento de parámetros del cómputo por token, permitiendo modelos de lenguaje de frontera con recuentes masivos de parámetros (por ejemplo, Mixtral-8×7B con 46.7B de parámetros) manteniendo una inferencia eficiente. Sin embargo, el despliegue sigue estando limitado por los costes de memoria, ya que cada experto debe residir en la memoria aunque solo un subconjunto esté activo por token. La poda de expertos —la eliminación de expertos completos— ofrece una solución directa para reducir la memoria y la latencia.

El desafío central es identificar qué expertos eliminar sin una degradación catastrófica del rendimiento. Las heurísticas existentes (por ejemplo, magnitud de los pesos, frecuencia de activación) a menudo fallan porque ignoran los roles funcionales o las magnitudes de contribución. Un trabajo teórico reciente de Chowdhury et al. [3] propuso un criterio demostrablemente efectivo: la poda de los expertos cuyos pesos del enrutador cambian menos durante el ajuste fino preserva la precisión en las pruebas. Sin embargo, esta garantía depende del ajuste fino completo (full fine-tuning), lo cual es computacionalmente prohibitivo a la escala de los grandes modelos MoE (costando más de lo que el ahorro de despliegue permite). El artículo investiga si este signo de "sensibilidad del enrutador" puede ser eliciado mediante el ajuste fino de parámetros eficientes (PEFT) para hacer que el criterio sea práctico a escala.

Metodología

Los autores proponen un marco de poda de "un solo paso" (one-shot) donde un método PEFT actúa no como una técnica de entrenamiento para producir un modelo final, sino como un instrumento de medición para identificar expertos podables.

1. La adaptación como medición: El modelo se somete a una breve ejecución de ajuste fino utilizando un adaptador de parámetros eficientes. El objetivo es únicamente inducir un cambio medible en el comportamiento del enrutador relativo a la tarea.
2. Puntuación de sensibilidad:
   • Basada en pesos (Enrutador sin congelar): Para los métodos que actualizan los pesos del enrutador (por ejemplo, LoRA en el enrutador), la puntuación $s_e$ es la magnitud del cambio de la norma $\ell_2$ del vector de puerta del experto: $s_e = |\|w_e^{(T)}\|_2 - \|w_e^{(0)}\|_2|$. Los expertos con los cambios más pequeños son podados.
   • Basada en logits (Enrutador congelado): Para los métodos que dejan los pesos del enrutador congelados (por ejemplo, adaptadores IA3, Houlsby), la puntuación es el desplazamiento de la raíz cuadrática media de los logits del enrutador para ese experto a través de un lote (batch) de validación: $s_e^{logit} = \frac{1}{\sqrt{N}} \|\Delta L_{:,e}\|_2$.
3. Procedimiento de poda: Los expertos se clasifican y eliminan independientemente por capa. La poda es estructural (eliminando parámetros del checkpoint) y de un solo paso (no sigue reentrenamiento ni destilación). El adaptador se descarta después de la clasificación.
4. Variables experimentales: El estudio varía sistemáticamente:
   • Módulos objetivo: Adjuntar adaptadores solo al enrutador frente a enrutador + atención + pesos de los expertos.
   • Rango del adaptador: Variar el rango de LoRA ($r \in \{2, 4, 8, 16\}$).
   • Geometría de perturbación: Comparar actualizaciones aditivas de bajo rango (LoRA), reescalado multiplicativo (IA3) y residuales aditivos no restringidos (Houlsby).
5. Modelos y Benchmarks: Evaluación primaria en Mixtral-8×7B-Instruct usando MMLU-Pro (un benchmark de razonamiento intensivo). La transferibilidad se prueba en Qwen1.5-MoE (razonamiento matemático) y a través de otros benchmarks (GSM8K, ARC-Challenge, BIG-Bench Hard, BoolQ).

Contribuciones Clave

1. Validación de la sensibilidad del enrutador a escala: Este trabajo proporciona la primera validación del criterio de cambio de norma del enrutador de Chowdhury et al. [3] en modelos de lenguaje MoE, extendiéndolo de la clasificación de visión a modelos ajustados por instrucciones evaluados en benchmarks de razonamiento.
2. Desacoplamiento del ajuste fino completo: El artículo demuestra que la señal de poda no requiere un ajuste fino completo. Un LoRA solo para el enrutador, entrenando aproximadamente el 0.002% de los parámetros (específicamente 0.0022% con rango 8), recupera con éxito la señal. Expandir el adaptador a otros módulos degrada la calidad de la señal.
3. Caracterización de la geometría de la adaptación: El estudio aísla cómo la "geometría" de la perturbación afecta la fidelidad de la señal:
   • Restringido vs. No restringido: Las perturbaciones restringidas (LoRA en el enrutador, IA3) producen mejores clasificaciones que las perturbaciones aditivas no restringidas (Houlsby).
   • Directo vs. Indirecto: Los métodos que nunca escriben en los pesos del enrutador (IA3) funcionan de manera comparable a los que sí lo hacen (LoRA), siempre que la perturbación sea restringida.
4. Sensibilidad al rango: La fidelidad de la señal de poda aumenta monótonamente con el rango de LoRA, aunque con rendimientos decrecientes. El rango 8 captura la mayor parte de la señal disponible.
5. Condicionamiento de la tarea: La señal está condicionada por el corpus de adaptación. Adaptar sobre un corpus genérico preserva capacidades amplias, mientras que adaptar sobre un dominio específico (por ejemplo, matemáticas) protege a los expertos de dicho dominio.
6. Baselines empíricos: El artículo proporciona las primeras evaluaciones de MMLU-Pro y BIG-Bench Hard para Mixtral podado por enrutador, estableciendo una curva de decaimiento cuasi-lineal para la poda guiada por el enrutador, contrastando con el colapso catastrófico visto en la poda basada en magnitud o aleatoria.

Resultados

• Calidad de la señal: En Mixtral-8×7B-Instruct, restringir LoRA al enrutador (entrenando el 0.0022% de los parámetros con rango 8) logra un 27.54% de precisión en MMLU-Pro tras eliminar el 50% de los expertos. Esto supera significativamente la adaptación de todos los módulos (24.42%) y excede ampliamente la poda basada en magnitud (15.88%) o la aleatoria (16.31%).
• Características de decaimiento: La poda guiada por el enrutador exhibe un decaimiento cuasi-lineal en la precisión a medida que se eliminan expertos, lo que hace que el coste de la compresión sea predecible. En contraste, la poda por magnitud y la aleatoria colapsan abruptamente tras eliminar uno o dos expertos.
• Transferibilidad: En Qwen1.5-MoE (completamente ajustado para matemáticas), la sensibilidad del enrutador retiene un 49.7% de precisión media en 11 benchmarks de matemáticas con una poda del 50%, mientras que la poda aleatoria colapsa a cifras de un solo dígito (4.4%).
• Eficiencia de despliegue: Podar 4 expertos (reducción del 50%) reduce el uso de memoria en un 49% y la latencia por token en un 37% en GPUs A100. El ahorro de latencia y memoria escala casi linealmente con el número de expertos eliminados.
• Comparación con SOTA: Con una compresión del 25%, la sensibilidad del enrutador es competitiva con métodos que utilizan estadísticas de activación completas o búsqueda diferenciable (por ejemplo, REAP, DiEP), a pesar de entrenar órdenes de magnitud menos parámetros.

Significado y Reivindicaciones

El artículo argumenta que la sensibilidad del enrutador medida bajo un ajuste fino ligero hace que la poda de expertos, motivada de forma demostrable, sea práctica a escala.

• Viabilidad económica: El método transforma la señal de poda de un subproducto del costoso ajuste fino completo en un subproducto de un PEFT económico. Si un profesional ya está realizando PEFT para el despliegue, la clasificación de poda llega con un coste marginal de cero.
• Visión instrumental de PEFT: El trabajo redefine los métodos PEFT como instrumentos para medir la dinámica interna del modelo en lugar de solo técnicas de entrenamiento. Establece que la "mejor" configuración para el entrenamiento (adaptación amplia) es a menudo la "peor" para medir señales de poda (señal diluida), y viceversa.
• Predictibilidad: A diferencia de la poda aleatoria o basada en magnitud que conduce a colapsos de rendimiento impredecibles, la sensibilidad del enrutador ofrece un compromiso cuasi-lineal y predecible entre compresión y precisión, permitiendo a los profesionales "valorar" un objetivo de compresión de antemano.
• Modestia: Los autores reconocen que la señal es competitiva en lugar de dominante frente a los criterios libres de entrenamiento en todos los regímenes. Señalan que el criterio es más efectivo cuando la tarea de adaptación coincide con el dominio de despliegue y que los resultados actuales son de "un solo paso" (no calibrados), dejando espacio para una futura recuperación mediante destilación o reentrenamiento, particularmente en arquitecturas de grano fino.

El artículo concluye que, al concentrar la adaptación en el enrutador y utilizar geometrías de perturbación restringidas, los profesionales pueden identificar y eliminar de forma fiable los expertos redundantes, reduiendo significativamente la huella de memoria y latencia de los modelos MoE sin requerir un ajuste fino completo.