Resumen Técnico: Asimetría Geométrica en la Especialización de MoE

Enunciado del Problema

Aunque las arquitecturas de Mezcla de Expertos (MoE) se han convertido en un mecanismo estándar para escalar modelos Transformer permitiendo grandes cantidades de parámetros con costos computacionales fijos, la estructura geométrica interna de la especialización de expertos sigue siendo poco comprendida. Los análisis existentes se centran principalmente en estadísticas de enrutamiento, equilibrio de carga o rendimiento a nivel de sistema. Existe una brecha crítica en la comprensión de cómo difieren los expertos en el espacio de funciones (las transformaciones que realizan) versus el espacio de representaciones (cómo se organizan los tokens enrutados en los estados ocultos).

La intuición actual sugiere que las capas de MoE dividen la distribución de entrada en clusters disjuntos donde cada experto actúa como un aproximador local. Sin embargo, la evidencia empírica sugiere que los expertos a menudo operan en regímenes de tokens superpuestos, planteando la pregunta: ¿Qué estructura geométrica induce el enrutamiento dentro de un Transformer MoE preentrenado? Específicamente, ¿implementan los expertos funciones distintas y cómo se relacionan sus representaciones entre sí?

Metodología

Los autores introducen un marco unificado Jacobian-PCA-Grassmann para analizar capas de MoE tanto en espacios de funciones como de representaciones. Este marco se aplica a modelos preentrenados (Mistral, Qwen) y a modelos Transformer controlados de 3 capas entrenados desde cero.

1. Geometría Funcional: Jacobianos Locales de Experto

Para caracterizar la función de expertos individuales independientemente de la dinámica de enrutamiento, los autores analizan el Jacobiano local del experto ($J_e(x) = \partial f_e(x) / \partial x$).

• Métrica: Calculan el Jacobiano promedio ($\bar{J}_e$) para cada experto sobre un conjunto de datos y lo vectorizan.
• Alineación: La alineación funcional par a par se mide mediante el producto interno normalizado de Frobenius (similitud coseno) entre estos Jacobianos vectorizados.
• Interpretación: Una baja similitud coseno indica decorrelación funcional, lo que significa que los expertos implementan transformaciones localmente distintas en lugar de variantes redundantes de la misma función.

2. Geometría de Representación: PCA Enrutado y Distancias Grassmannianas

Para caracterizar la geometría de las representaciones enrutadas, los autores recopilan estados ocultos ponderados por los pesos de enrutamiento ($h_e(x) = g_e(x)h(x)$).

• PCA: Se realiza un Análisis de Componentes Principales (PCA) estándar sobre los estados ocultos enrutados para identificar direcciones principales y espectros de varianza.
• Distancia de Subespacio: La separación geométrica entre expertos se cuantifica utilizando la Distancia Geodésica Grassmanniana entre los subespacios PCA de los $n$ primeros componentes de diferentes expertos.
• Interpretación: Esta métrica captura la separación multidireccional entre subespacios completos. Una distancia cercana a cero implica roles funcionales redundantes, mientras que una distancia que se acerca al máximo teórico ($\frac{\pi}{2}\sqrt{n}$) implica ortogonalidad.

3. Experimentos de Enrutamiento Controlado

Para aislar el efecto de los mecanismos de enrutamiento de la dinámica de entrenamiento, los autores comparan dos regímenes en un modelo controlado de 3 capas con parámetros de expertos fijos:

• Enrutamiento Top-k: Selección dispersa y discreta.
• Enrutamiento Completamente Suave: Mezcla densa y continua.

Resultados Clave

1. Asimetría Estructural: Decorrelación Funcional vs. Superposición de Representaciones

En Transformers MoE preentrenados (Mistral-8x7B y Qwen1.5-MoE-A2.7B), los autores observan una asimetría estructural consistente:

• Fuerte Decorrelación Funcional: La alineación de Jacobianos entre expertos es consistentemente cercana a cero (media $\approx$ 0.000–0.068). Esto indica que los expertos no son redundantes; implementan transformaciones locales distintas.
• Superposición Parcial de Representaciones: A pesar de la decorrelación funcional, las representaciones enrutadas ocupan subespacios distintos pero parcialmente superpuestos. Las distancias Grassmannianas son sustanciales (indicando separación) pero permanecen significativamente por debajo del máximo teórico para la ortogonalidad completa.
• Conclusión: La decorrelación funcional y la superposición de representaciones coexisten. Los expertos actúan como operadores localmente decorrelacionados sobre regímenes superpuestos en una variedad compartida, en lugar de dividir el espacio de entrada en clusters disjuntos.

2. Dispersión del Enrutamiento como Control Geométrico

Los experimentos controlados revelan que la nitidez del enrutamiento moldea directamente esta geometría:

• Enrutamiento Top-k: Induce una separación funcional más nítida (baja alineación de Jacobianos) y distancias Grassmannianas más grandes (mayor divergencia de subespacios).
• Enrutamiento Completamente Suave: Produce estructuras de expertos entrelazadas con alta alineación de Jacobianos (indicando colapso hacia mapeos redundantes) y distancias Grassmannianas colapsadas (alta superposición de subespacios).
• Implicación: La dispersión del enrutamiento es un factor clave en la imposición de la especialización.

3. Comparación con Capas Densas

Los expertos de MoE exhiben un decaimiento más pronunciado de los valores propios en sus espectros de varianza en comparación con las capas de alimentación frontal densas. Los primeros componentes principales capturan un porcentaje significativamente mayor de varianza en los expertos de MoE, lo que sugiere que el enrutamiento concentra los tokens en subespacios de menor dimensión y más enfocados.

4. Consistencia entre Modelos

Los patrones geométricos observados (baja alineación de Jacobianos + alta pero submáxima separación Grassmanniana) son consistentes entre diferentes modelos (Mistral, Qwen), capas (tempranas, medias, tardías) y conjuntos de datos, lo que sugiere que esta es una propiedad robusta de las arquitecturas MoE en lugar de un artefacto de ejecuciones de entrenamiento específicas.

Significado y Afirmaciones

El artículo afirma proporcionar un marco de diagnóstico general para estudiar el cálculo condicional en arquitecturas Transformer modernas. Sus contribuciones principales son:

1. Herramienta de Diagnóstico Unificada: Ofrece un método sistemático para analizar la especialización de expertos más allá de las estadísticas simples de enrutamiento o patrones de activación, caracterizando conjuntamente el espacio de funciones y el espacio de representaciones.
2. Interpretación Geométrica de MoE: Los resultados apoyan una visión de las capas de MoE como la implementación de una partición suave de una variedad de representación compartida. Los expertos son funcionalmente distintos pero operan sobre subvariedades superpuestas.
3. Perspectiva Mecanística: El estudio identifica la nitidez del enrutamiento (dispersión) como un mecanismo de control directo para el grado de especialización, vinculando las decisiones de diseño arquitectónico con estructuras geométricas medibles.
4. Implicaciones para el Diseño de Modelos: Los autores sugieren que las estrategias para la fusión o poda de expertos basadas únicamente en la similitud de representaciones (por ejemplo, superposición de activaciones) pueden ser defectuosas. Dado que los expertos pueden tener representaciones superpuestas mientras realizan funciones distintas, las métricas funcionales (como la alineación de Jacobianos) pueden ser necesarias para identificar correctamente la verdadera redundancia.

El artículo mantiene un alcance modesto, reconociendo que su análisis es local (Jacobianos de primer orden) y lineal (PCA), y no captura el comportamiento no lineal global ni la evolución de la geometría durante el entrenamiento. Sin embargo, establece una perspectiva geométrica fundamental sobre cómo los modelos MoE organizan el cálculo.