Resumo Técnico: Roteamento Adaptativo de Índice Invertido para Misturas de Especialistas Granulares (AIR-MoE)

1. Declaração do Problema

Os modelos de Mistura de Especialistas (MoE) tornaram-se um pilar dos Modelos de Linguagem de Grande Escala (LLMs) escaláveis, ativando apenas um subconjunto de especialistas por token. Leis de escala empíricas recentes sugerem que o desempenho melhora com especialistas granulares — utilizando muitos especialistas pequenos em vez de poucos grandes. No entanto, esse regime introduz um gargalo computacional significativo: o custo de roteamento. À medida que o número de especialistas ($E$) aumenta, o roteador deve calcular produtos internos entre cada representação de token e cada centróide de especialista para identificar os top-$K$ especialistas. Essa sobrecarga pode dominar a computação total, tornando o treinamento de modelos altamente granulares proibitivo.

As soluções existentes para reduzir os custos de roteamento geralmente impõem restrições estruturais que limitam a expressividade do modelo:

• MoEs Hierárquicos dividem os especialistas em grupos fixos e não sobrepostos, restringindo o roteamento dentro desses grupos.
• Restrições de Chave de Produto/Estruturais (por exemplo, He, 2024) dividem os centróides de especialistas em componentes para permitir recuperação exata, mas restringem o espaço de representações de especialistas admissíveis (por exemplo, forçando uma estrutura de produto cartesiano).
• Abordagens de Hashing Fixo carecem de aprendibilidade e demonstraram ser fundamentalmente limitadas.

O desafio é aproximar o roteamento top-$K$ de forma eficiente sem impor restrições estruturais aos parâmetros dos especialistas ou exigir a pontuação completa de todos os especialistas.

2. Metodologia: AIR-MoE

Os autores propõem o Roteamento Adaptativo de Índice Invertido para MoE (AIR-MoE), uma arquitetura de roteamento inspirada em índices de arquivo invertido usados em recuperação de informação e busca aproximada de vizinhos mais próximos. O método aproxima os verdadeiros especialistas top-$K$ por meio de um processo de recuperação grosseiro-fino em duas etapas, evitando a pontuação completa de todos os especialistas.

2.1 Recuperação em Duas Etapas

Para cada token de entrada $h$, o roteamento procede em paralelo:

1. Pré-seleção Grosseira: O token é mapeado para a palavra-código $c_g$ mais próxima em um código aprendido $\mathcal{C} = \{c_1, \dots, c_G\}$ usando similaridade de cosseno. Cada palavra-código $c_g$ está associada a uma pré-seleção $L_g$ pré-computada de $M$ especialistas candidatos. Essas pré-seleções são essencialmente "listas de postagem" em um índice invertido, armazenando os índices dos especialistas com os maiores produtos internos em relação ao centróide da palavra-código.
2. Pontuação Fina: O roteador calcula produtos internos exatos apenas entre o token $h$ e os especialistas dentro da pré-seleção recuperada $L_g$. Os top-$K$ especialistas são então selecionados a partir desse conjunto de candidatos reduzido.

Essa hierarquia garante que o conjunto de especialistas ativos $T$ seja um subconjunto da pré-seleção $L$, que é um subconjunto do conjunto completo de especialistas $[E]$.

2.2 Treinamento e Otimização

Um desafio central é que a atribuição de tokens a palavras-código é não diferenciável. O AIR-MoE aborda isso por meio de uma estratégia de otimização de dois níveis:

• Treinamento Baseado em Gradiente: Representações de tokens e centróides de especialistas são atualizados via retropropagação padrão a partir da perda de modelagem de linguagem a jusante.
• Treinamento de Código Livre de Gradiente: O código $\mathcal{C}$ e suas pré-seleções associadas são atualizados independentemente usando um algoritmo de k-médias esférico adaptativo. Este algoritmo:
  • Usa uma Média Móvel Exponencial (EMA) para rastrear a distribuição das representações de tokens.
  • Opera na esfera unitária para lidar efetivamente com a geometria de cosseno de alta dimensão.
  • Inclui um mecanismo de reinicialização para palavras-código "mortas" (aquelas com baixas contagens de atribuição) para evitar colapso.
  • Atualiza as pré-seleções $L_g$ apenas após etapas do otimizador, garantindo que o índice permaneça consistente com o estado atual dos centróides de especialistas.

Crucialmente, o AIR-MoE não usa o truque de estimativa direta; os gradientes não fluem através da etapa de quantização para o código.

2.3 Garantias Teóricas

O artigo deriva um limite inferior para o Recall de Massa (a massa de probabilidade total dos verdadeiros especialistas top-$K$ retidos na pré-seleção). O limite depende de:

1. O erro de quantização $\epsilon(h) = \|h - c(h)\|$: Um erro menor leva a uma melhor retenção.
2. A massa da pré-seleção $\rho_M(c)$: A massa de probabilidade capturada pelos top-$M$ especialistas em relação à palavra-código.
   Teoricamente, o AIR-MoE retém a maior parte da massa de roteamento quando os tokens são bem quantizados e o tamanho da pré-seleção $M$ é suficientemente grande.

3. Contribuições Principais

1. Arquitetura AIR-MoE: Um mecanismo de roteamento inspirado em índice invertido que aproxima o roteamento top-$K$ para MoEs granulares sem impor restrições estruturais (por exemplo, agrupamento ou restrições paramétricas) às representações dos especialistas.
2. Otimização de Dois Níveis: Uma estratégia de treinamento que desacopla a adaptação livre de gradiente do código (via k-médias esférico adaptativo) do treinamento baseado em gradiente dos parâmetros do modelo, evitando a necessidade de estimativa direta.
3. Análise Teórica: Um limite inferior derivado para o recall de massa, fornecendo intuição sobre as condições sob as quais a aproximação preserva a qualidade do roteamento.
4. Desempenho Empírico: Demonstração de que o AIR-MoE alcança um compromisso superior entre perplexidade e FLOPs em comparação com métodos existentes de roteamento granular eficiente.

4. Resultados Experimentais

Os autores avaliaram o AIR-MoE em WikiText-103, C5 e OpenWebText2 em escalas de modelo pequena, média e grande (até ~0,45B de parâmetros) com 65.536 especialistas.

• Desempenho vs. Baselines: O AIR-MoE superou consistentemente as baselines granulares eficientes existentes, incluindo PEER (He, 2024) e MoE Hierárquico.
  • Alcançou melhorias de perplexidade de até 10% sobre a baseline PEER em custos de FLOP semelhantes.
  • Embora o roteamento granular padrão (pontuando todos os especialistas) tenha alcançado a menor perplexidade, ele incorreu em custos computacionais proibitivos (por exemplo, ~467,5P FLOPs vs. 324,1P para AIR-MoE em WikiText-103 Pequeno).
• Eficiência: O AIR-MoE reduziu significativamente os FLOPs de treinamento em comparação com o roteamento granular exato, mantendo uma perplexidade competitiva.
• Estudos de Ablação:
  • Remover a normalização de centróides ou usar distância euclidiana em vez de k-médias esférico degradou ligeiramente o desempenho.
  • Usar um código estático (não adaptativo) resultou no pior desempenho, confirmando a necessidade do código adaptativo.
  • O roteamento de escolha de especialista levou a alta sobreposição, mas grave fome de especialistas (78,6% de especialistas mortos), indicando que a sobreposição sozinha não é uma métrica suficiente para a qualidade do roteamento.
• Uso de Especialistas: Ao contrário dos roteadores hierárquicos ou PEER, que impõem estruturas rígidas levando a artefatos de roteamento, o AIR-MoE permitiu uso flexível de especialistas sem restrições estruturais.

5. Significado e Alegações

O artigo afirma que o AIR-MoE fornece um mecanismo eficaz para roteamento em arquiteturas de mistura de especialistas esparsas, permitindo especificamente o uso de especialistas granulares sem a penalidade computacional do roteamento completo.

• Substituição Direta: O método não requer modificações na arquitetura do modelo subjacente ou na função de perda, servindo como uma substituição direta para roteadores padrão.
• Flexibilidade: Ao evitar restrições estruturais nos parâmetros dos especialistas, o AIR-MoE permite que o modelo aprenda dependências token-especialista de forma mais flexível do que métodos que exigem agrupamentos fixos ou estruturas de chave de produto.
• Escalabilidade: A abordagem escala favoravelmente em regimes de grande lote efetivo porque o custo de pré-seleção grosseira é amortizado e independente do número de tokens roteados.

Os autores reconhecem limitações, observando que, embora as reduções de FLOP sejam significativas, os aumentos reais de velocidade em tempo de parede dependem de suporte de hardware e implementações otimizadas para mitigar o tráfego de memória e as sobrecargas de indexação inerentes aos MoEs granulares. Eles também observam que a implementação atual não emprega redução de dimensionalidade, deixando isso como uma direção ortogonal para trabalhos futuros.