Resumo Técnico: Assimetria Geométrica na Especialização de MoE

Declaração do Problema

Embora as arquiteturas de Mistura de Especialistas (MoE) tenham se tornado um mecanismo padrão para escalar modelos Transformer, permitindo grandes contagens de parâmetros com custos computacionais fixos, a estrutura geométrica interna da especialização dos especialistas permanece pouco compreendida. As análises existentes focam majoritariamente em estatísticas de roteamento, balanceamento de carga ou desempenho em nível de sistema. Existe uma lacuna crítica na compreensão de como os especialistas diferem no espaço de função (as transformações que realizam) versus no espaço de representação (como os tokens roteados são organizados nos estados ocultos).

A intuição atual sugere que as camadas MoE particionam a distribuição de entrada em clusters disjuntos, onde cada especialista atua como um aproximador local. No entanto, evidências empíricas sugerem que os especialistas frequentemente operam em regimes de tokens sobrepostos, levantando a questão: Qual estrutura geométrica o roteamento induz dentro de um Transformer MoE pré-treinado? Especificamente, os especialistas implementam funções distintas e como suas representações se relacionam entre si?

Metodologia

Os autores introduzem um framework unificado Jacobiano-PCA-Grassmann para analisar camadas MoE tanto no espaço de função quanto no espaço de representação. Este framework é aplicado a modelos pré-treinados (Mistral, Qwen) e a modelos Transformer controlados de 3 camadas treinados do zero.

1. Geometria Funcional: Jacobianos Locais ao Especialista

Para caracterizar a função de especialistas individuais independentemente da dinâmica de roteamento, os autores analisam o jacobiano local ao especialista ($J_e(x) = \partial f_e(x) / \partial x$).

• Métrica: Eles computam o Jacobiano médio ($\bar{J}_e$) para cada especialista sobre um conjunto de dados e o vetorizam.
• Alinhamento: O alinhamento funcional par a par é medido via o produto interno de Frobenius normalizado (similaridade de cosseno) entre esses Jacobianos vetorizados.
• Interpretação: Baixa similaridade de cosseno indica decorrelação funcional, significando que os especialistas implementam transformações localmente distintas em vez de variantes redundantes da mesma função.

2. Geometria de Representação: PCA Roteado e Distâncias Grassmannianas

Para caracterizar a geometria das representações roteadas, os autores coletam estados ocultos ponderados pelos pesos de roteamento ($h_e(x) = g_e(x)h(x)$).

• PCA: Uma Análise de Componentes Principais (PCA) padrão é realizada nos estados ocultos roteados para identificar direções principais e espectros de variância.
• Distância de Subespaço: A separação geométrica entre especialistas é quantificada usando a Distância Geodésica Grassmanniana entre os subespaços PCA dos top-$n$ de diferentes especialistas.
• Interpretação: Esta métrica captura a separação multidirecional entre subespaços inteiros. Uma distância próxima de zero implica papéis funcionais redundantes, enquanto uma distância aproximando-se do máximo teórico ($\frac{\pi}{2}\sqrt{n}$) implica ortogonalidade.

3. Experimentos de Roteamento Controlado

Para isolar o efeito dos mecanismos de roteamento da dinâmica de treinamento, os autores comparam dois regimes em um modelo controlado de 3 camadas com parâmetros de especialista fixos:

• Roteamento Top-k: Seleção esparsa e discreta.
• Roteamento Fully-Soft: Mistura densa e contínua.

Principais Resultados

1. Assimetria Estrutural: Decorrelação Funcional vs. Sobreposição de Representação

Em Transformers MoE pré-treinados (Mistral-8x7B e Qwen1.5-MoE-A2.7B), os autores observam uma assimetria estrutural consistente:

• Forte Decorrelação Funcional: O alinhamento de Jacobianos entre especialistas é consistentemente próximo de zero (média $\approx$ 0,000–0,068). Isso indica que os especialistas não são redundantes; eles implementam transformações locais distintas.
• Sobreposição Parcial de Representação: Apesar da decorrelação funcional, as representações roteadas ocupam subespaços distintos, mas parcialmente sobrepostos. As distâncias Grassmannianas são substanciais (indicando separação), mas permanecem significativamente abaixo do máximo teórico para ortogonalidade completa.
• Conclusão: A decorrelação funcional e a sobreposição de representação coexistem. Os especialistas atuam como operadores localmente decorrelacionados sobre regimes sobrepostos em uma variedade compartilhada, em vez de particionar o espaço de entrada em clusters disjuntos.

2. Esparsidade de Roteamento como Controle Geométrico

Experimentos controlados revelam que a nitidez do roteamento molda diretamente essa geometria:

• Roteamento Top-k: Induz uma separação funcional mais nítida (baixo alinhamento de Jacobiano) e maiores distâncias Grassmannianas (maior divergência de subespaço).
• Roteamento Fully-Soft: Produz estruturas de especialistas emaranhadas com alto alinhamento de Jacobiano (indicando colapso em direção a mapeamentos redundantes) e distâncias Grassmannianas colapsadas (alta sobreposição de subespaço).
• Implicação: A esparsidade do roteamento é um fator chave na imposição de especialização.

3. Comparação com Camadas Densas

Especialistas MoE exibem um decaimento mais acentuado de autovalores em seus espectros de variância em comparação com camadas de alimentação direta densas. Os primeiros componentes principais capturam uma porcentagem significativamente maior de variância em especialistas MoE, sugerindo que o roteamento concentra tokens em subespaços de dimensão mais baixa e mais focados.

4. Consistência entre Modelos

Os padrões geométricos observados (baixo alinhamento de Jacobiano + alta, mas submáxima, separação Grassmanniana) são consistentes entre diferentes modelos (Mistral, Qwen), camadas (iniciais, médias, tardias) e conjuntos de dados, sugerindo que esta é uma propriedade robusta das arquiteturas MoE, e não um artefato de execuções de treinamento específicas.

Significado e Alegações

O artigo alega fornecer um framework de diagnóstico geral para estudar computação condicional em arquiteturas Transformer modernas. Suas contribuições primárias são:

1. Ferramenta de Diagnóstico Unificada: Oferece um método sistemático para analisar a especialização de especialistas além de estatísticas simples de roteamento ou padrões de ativação, caracterizando conjuntamente o espaço de função e o espaço de representação.
2. Interpretação Geométrica de MoE: Os resultados apoiam uma visão de camadas MoE como implementando uma particionamento suave de uma variedade de representação compartilhada. Os especialistas são funcionalmente distintos, mas operam em subvariedades sobrepostas.
3. Insight Mecanístico: O estudo identifica a nitidez do roteamento (esparsidade) como um mecanismo de controle direto para o grau de especialização, ligando escolhas de design arquitetural a estruturas geométricas mensuráveis.
4. Implicações para o Design de Modelos: Os autores sugerem que estratégias para mesclagem ou poda de especialistas baseadas apenas na similaridade de representação (por exemplo, sobreposição de ativação) podem ser falhas. Como os especialistas podem ter representações sobrepostas enquanto realizam funções distintas, métricas funcionais (como alinhamento de Jacobiano) podem ser necessárias para identificar corretamente a verdadeira redundância.

O artigo mantém um escopo modesto, reconhecendo que sua análise é local (Jacobianos de primeira ordem) e linear (PCA), e não captura comportamento não linear global ou a evolução da geometria durante o treinamento. No entanto, estabelece uma perspectiva geométrica fundamental sobre como modelos MoE organizam a computação.