Hierarchical Copula-Gumbel-Top-\texorpdfstring{$K$}{K} Routing: Two-Sided Dependence Control for Frozen Mixture-of-Experts at Fixed Per-Token Routing Laws

Resumo Técnico: Roteamento Hierárquico Copula-Gumbel-Top-K

Definição do Problema
Em modelos de mistura de especialistas (MoE) esparsos, o mecanismo de roteamento estocástico padrão (Gumbel-Top-K) define uma "lei de roteamento" para cada token: uma distribuição de probabilidade específica sobre listas ordenadas de especialistas e pesos de mistura, condicionada aos logits do token. Enquanto quase todo o trabalho existente modifica essa lei por token (por exemplo, via ajuste fino de logits ou perdas auxiliares), a distribuição conjunta das escolhas de roteamento entre diferentes tokens permanece amplamente inexplorada. O artigo questiona: dado que a lei de roteamento de cada token individual é mantida exatamente fixa (como em um modelo congelado), quais distribuições conjuntas sobre as escolhas de roteamento de diferentes tokens são alcançáveis? Especificamente, é possível controlar a dependência entre tokens para melhorar a coerência dos especialistas ou o balanceamento de carga sem alterar o comportamento marginal de qualquer token individual?

Metodologia: Hierarchical Copula-Gumbel-Top-K (H-CGA)
O autor propõe o H-CGA, uma camada de dependência do lado do roteamento que opera no vetor de ruído Gumbel antes da seleção Top-K. O método baseia-se no teorema de Sklar para separar as distribuições marginais das estruturas de dependência. A construção procede em quatro estágios:

1. Agrupamento Pré-Roteamento: Os tokens são particionados em grupos disjuntos (por exemplo, janelas fixas de tokens adjacentes). Um controlador lê características pré-roteamento congeladas (por exemplo, estados ocultos, entropia do gate) para determinar as forças de acoplamento.
2. Acoplamento Positivo Intra-Grupo: Para tokens dentro de um grupo $g$, as perturbações Gumbel para cada coordenada de especialista são geradas via uma cópula Gaussiana permutável. Os tokens compartilham uma variável latente $\zeta_{ge}$, induzindo correlação positiva em suas perturbações aleatórias. Isso aumenta a probabilidade de que tokens relacionados selecionem os mesmos especialistas (coerência).
3. Acoplamento Antitético Entre Grupos: Pares disjuntos de grupos $(g, g')$ são pareados. Uma força de oposição ajustável $\alpha_{g,g'} \in [0, 1]$ é introduzida. O latente compartilhado do segundo grupo é construído como uma combinação linear do latente do primeiro grupo e um termo de ruído independente. Quando $\alpha_{g,g'} = 1$, os grupos recebem impulsos compartilhados totalmente opostos; quando $\alpha_{g,g'} = 0$, eles são independentes. Isso introduz dependência negativa entre grupos para neutralizar a volatilidade da carga (load burstiness).
4. Preservação de Marginais: A transformação dos latentes compartilhados para as variáveis Gumbel garante que, para qualquer token $t$, o vetor de perturbações $(\gamma_{t1}, \dots, \gamma_{tE})$ permaneça com distribuição Gumbel i.i.d., idêntica ao roteamento independente.

Principais Contribuições

1. Arquitetura H-CGA: Uma camada de cópula hierárquica que controla a dependência cruzada entre tokens em ambas as direções (positiva intra-grupo, negativa ajustável entre grupos) sem modificar os logits do roteador, os pesos dos especialistas ou qualquer lei de roteamento por token.
2. Invariância Teórica: O autor prova (Teorema 1) que toda a hierarquia preserva a distribuição condicional da lista Top-K ordenada e dos pesos de mistura de cada token. Consequentemente, as contagens esperadas condicionais de inclusão de especialistas (Corolário 1) são idênticas às do roteamento independente.
3. Caracterização do Trade-off Coerência–Dispersão: O artigo caracteriza o trade-off de variância (Proposição 1):
   • O acoplamento positivo intra-grupo aumenta a variância das cargas de especialistas realizadas em relação ao roteamento independente (um custo de coerência local).
   • A oposição entre grupos não negativa reduz essa variância em relação ao acoplamento "plano" (grupos independentes) no mesmo nível de força intra-grupo, embora não necessariamente reduza a variância abaixo da linha de base do roteamento independente.
4. Adaptação Apenas de Roteamento: O autor formula um método de adaptação onde um controlador pequeno define os seletores de dependência ($\rho_g, \alpha_{g,g'}$) com base em características congeladas. Este controlador é treinado usando um estimador de função de pontuação (REINFORCE), exigindo apenas passagens de ida (forward passes) através do modelo base congelado e gradientes confinados ao controlador.

Resultos e Estudo Piloto
O artigo apresenta um pequeno estudo piloto em um decodificador MoE de 15,8 milhões de parâmetros treinado em TinyStories, com o modelo base congelado durante a avaliação.

• Verificações de Invariância: Testes sintéticos confirmaram que a frequência empírica de listas Top-K ordenadas e as probabilidades de inclusão de especialistas sob H-CGA coincidiram com o roteamento independente dentro das margens de erro de Monte Carlo, validando a invariância teórica.
• Efeitos Estatísticos:
  • O aumento do acoplamento intra-grupo ($\rho$) aumentou significativamente a similaridade de Jaccard dos conjuntos de especialistas dentro das janelas e reduziu o número de especialistas distintos usados por janela.
  • O aumento da oposição entre grupos ($\alpha$) reduziu significavelmente a similaridade de Jaccard nas fronteiras dos grupos, confirmando a capacidade de desacoplar as demandas dos grupos.
• Aprendizado: Um controlador escalar aprendido (3 parâmetros) foi treinado de forma estável, mas não mostrou melhoria na entropia cruzada de validação em comparação com a linha de base independente. As forças de acoplamento aprendidas foram baixas, consistente com a observação do artigo de que a preservação marginal limita o sinal disponível para otimização quando apenas a perda por token é usada.
• Variância de Carga: O piloto observou mudanças mínimas no coeficiente de variação (CV) da carga agregada, observando que esta métrica de resumo não testa diretamente os limites de variância condicional propostos na Proposição 1.

Significância e Alegações
O artigo afirma que o H-CGA expõe um novo grau de liberdade no roteamento de MoE: a estrutura de dependência conjunta do ruído de roteamento. Sua principal significância é teórica e mecânica:

• Segurança: Fornece uma "garantia de segurança" de que a modificação da dependência entre tokens não altera o comportamento de roteamento marginal de qualquer token, garantindo que as preferências aprendidas pelo modelo base não sejam afetadas.
• Controle: Oferece dois "seletores" complementares para trocar a coerência do conjunto de especialistas local pela dispersão de carga global.
• Eficiência: Permite a adaptação apenas de roteamento onde o modelo base é avaliado apenas na direção de ida, evitando o custo computacional de atualizar os pesos dos especialistas.

Limitações e Modéstia
O autor é explícito sobre o escopo e as limitações:

• Garantia Local da Camada: Os resultados de invariância são condicionais aos estados ocultos e logits que entram em uma camada de roteamento específica. Eles não garantem invariância de ponta a ponta em uma rede multicamadas, pois o acoplamento pode alterar os estados ocultos subsequentes.
• Ausência de Ganhos em Nível de Tarefa: O piloto não estabelece que a adaptação apenas de roteamento melhora o desempenho em tarefas de jusante (downstream). O autor observa que a preservação marginal implica que objetivos dependentes exclusivamente de distribuições de tokens únicos não podem identificar o controlador; sinais úteis devem surgir de efeitos conjuntos (ex: interações de atenção), que podem ser fracos.
• Variância vs. Transbordamento (Overflow): Embora o método caracterize teoricamente a variância da carga, o piloto não mediu diretamente a variância condicional da carga ou os transbordamentos de capacidade. O artigo não afirma resolver o transbordamento de capacidade, mas sugere que o seletor de oposição é um mecanismo para abordá-lo.
• Sobrecarga (Overhead): O método requer a geração e coordenação de perturbações Gumbel para todos os especialistas, o que pode introduzir sobrecarga computacional em comparação ao roteamento padrão.

Em resumo, o artigo estabelece um framework matematicamente rigoroso para controlar a dependência de roteamento entre tokens em MoEs congelados, provando que a coerência e a dispersão de carga podem ser manipuladas sem alterar as leis de roteamento por token, mantendo-se cauteloso quanto aos benefícios práticos imediatos para o desempenho da tarefa.