Resumo Técnico: Roteamento Sinkhorn Seletivo para um Mixture of Experts Esparso Melhorado

Declaração do Problema

Modelos de Mistura de Especialistas Esparsos (SMoE) oferecem uma abordagem escalável para aumentar a capacidade do modelo mantendo a eficiência computacional ao ativar apenas um subconjunto de especialistas por token. No entanto, os modelos SMoE convencionais que dependem de roteamento baseado em Softmax frequentemente sofrem de "colapso de roteamento", onde um pequeno subconjunto de especialistas é sobreutilizado enquanto outros permanecem subutilizados.

As estratégias de mitigação existentes baseiam-se tipicamente em:

1. Objetivos Auxiliares: Perdas de balanceamento de carga ou perdas z que penalizam o uso desigual de especialistas. Estes podem introduzir desalinhamento de objetivo, instabilidade no treinamento e aumento da complexidade.
2. Ruído de Gating: Injeção de ruído para encorajar a exploração, o que adiciona parâmetros treináveis ou hiperparâmetros.
3. Roteamento baseado em Sinkhorn: Métodos anteriores que utilizam o algoritmo Sinkhorn para transporte ótimo (OT) frequentemente usam o mapa de transporte apenas para seleção de especialistas (top-k), em vez de atribuição de pesos. Além disso, eles podem reduzir a flexibilidade de roteamento porque a matriz de gating não é diretamente otimizada através de aprendizado baseado em gradiente da mesma forma que os métodos Softmax, ou incorrem em um overhead substancial de treinamento devido à natureza iterativa das atualizações Sinkhorn.

Metodologia

O artigo propõe o Roteamento Sinkhorn Seletivo (SSR), um mecanismo leve que reformula a atribuição de token-a-especialista como um problema de transporte ótimo com regularização de entropia.

Formulação Central

Os autores modelam a atribuição de $m$ tokens a $n$ especialistas como um problema de transporte ótimo de custo máximo com regularização de entropia. O objetivo é encontrar um plano de transporte $\hat{\Pi}$ que maximize a compatibilidade enquanto satisfaz restrições específicas:

• Positividade: $\Pi > 0$.
• Soma das Linhas: Cada token roteia toda a sua massa para os especialistas ($\Pi \mathbf{1}_n = \mathbf{1}_m$).
• Soma das Colunas: Cada especialista recebe uma carga total esperada igual ($\Pi^\top \mathbf{1}_m = \frac{m}{n} \mathbf{1}_n$).

A matriz de custo $C$ é derivada das pontuações de gating $S$. Os autores propõem duas variantes:

• Custo Linear: $C = S$.
• Custo Softmax: $C_{i,:} = \text{softmax}(S_{i,:})$, usado para prevenir overflow numérico durante as iterações Sinkhorn.

Ao contrário dos métodos anteriores baseados em OT que usam o plano de transporte apenas para seleção, o SSR utiliza diretamente os valores em $\hat{\Pi}$ para calcular os pesos de roteamento. Para cada token, os $k$ melhores especialistas são selecionados com base nas entradas mais altas na linha correspondente de $\hat{\Pi}$, e os pesos são normalizados a partir dessas entradas.

Estratégia Seletiva e Injeção de Ruído

Um desafio fundamental identificado é que o roteamento Sinkhorn não atualiza a matriz de peso de gating $W_g$ porque $W_g$ é desacoplada do objetivo de OT. Para resolver isso, o SSR emprega uma estratégia de treinamento híbrida:

• Aplicação Seletiva: Durante o treinamento, cada bloco MoE usa o roteamento Sinkhorn com uma pequena probabilidade $p$ (ex: 0,001) e o gating Softmax padrão caso contrário. Isso permite que o gating Softmax atualize $W_g$ efetivamente enquanto o roteamento Sinkhorn fornece sinais de utilização equilibrada de especialistas.
• Injeção de Ruído: Para evitar o colapso de especialistas e encorajar a exploração, ruído Gaussiano é adicionado à matriz de custo ($\tilde{C} = C + \alpha_{noise} \cdot \epsilon$). A análise teórica (Proposição 4.3) mostra que isso garante que cada especialista tenha uma probabilidade não nula de seleção.

Comportamento de Inferência

O artigo argumenta que impor restrições de equilíbrio (como somas de colunas iguais) durante a inferência em uma única entrada é contraproducente, pois força um roteamento uniforme independentemente das pontuações de compatibilidade específicas do input. Portanto, tanto o roteamento Sinkhorn quanto a injeção de ruído são desativados durante a inferência, revertendo para o roteamento Softmax determinístico para garantir previsões consistentes e dependentes do input.

Principais Contribuições

1. Novo Framework de Roteamento: Um método que integra transporte ótimo com regularização de entropia com injeção de ruído estocástico para promover o uso equilibrado de especialistas sem perdas de balanceamento auxiliares.
2. Insights Teóricos: Provas demonstrando que o roteamento baseado em Sinkhorn e a injeção de ruído auxiliam o treinamento ao encorajar a exploração e o balanceamento, mas devem ser desativados na inferência para evitar atribuições distorcidas.
3. Eficiência e Desempenho: Avaliações extensas mostrando que o SSR melhora a eficiência de treinamento, precisão e robustez em comparação com baselines de estado da arte.

Resultados Experimentais

Os autores avaliaram o SSR em modelagem de linguagem (WikiText-103, Enwik-8) e classificação de imagens (ImageNet-1K, ImageNet-A, ImageNet-O, ImageNet-R).

• Modelagem de Linguagem:
  • No WikiText-103, as variantes do SSR superaram consistentemente o Vanilla SMoE e baselines usando perda de balanceamento de carga, perda z ou ruído.
  • SSR-L w/ noise alcançou uma perplexidade de teste (PPL) de 34,367, uma redução de 1,183 em relação ao Vanilla SMoE (35,550), superando o próximo melhor baseline por uma margem significativa.
  • Crucialmente, o SSR incorreu em um overhead de treinamento mínimo (0,33% – 0,65%), enquanto o full Sinkhorn-based SMoE incorreu em 72,47% de overhead.
  • Sob configurações de Momentum, o SSR preservou a evolução característica da norma do Vanilla SMoE enquanto melhorava o desempenho, ao passo que outros métodos de balanceamento degradaram o desempenho.
• Modelagem de Nível de Byte (Enwik-8):
  • SSR-S w/ noise alcançou o melhor Bits-per-Character (BPC) de teste de 1,128, melhorando o Vanilla SMoE em 0,010 BPC.
  • Foi 2,28x mais rápido que o full Sinkhorn-based SMoE.
• Tarefas de Visão:
  • No ImageNet-1K, o SSR-L w/ noise alcançou uma acurácia Top-1 de 77,420, um ganho de +2,368 sobre o Vanilla SMoE.
  • Também mostrou forte robustez no ImageNet-O e ImageNet-R.
• Estudos de Ablação:
  • A probabilidade $p$ de aplicar o roteamento Sinkhorn é crítica; valores muito pequenos (ex: $10^{-3}$ ou $10^{-4}$) foram suficientes para fornecer sinais de balanceamento sem restringir excessivamente o roteador.
  • A injeção de ruído melhorou consistentemente o desempenho em todas as configurações.
  • Desativar os mecanismos de balanceamento na inferência gerou os melhores resultados, confirmando as alegações teóricas.

Significância e Alegações

O artigo afirma que o Roteamento Sinkhorn Seletivo oferece uma solução prática e eficaz para o design de SMoE. Ao substituir perdas auxiliares complexas por uma otimização intermitente e eficiente baseada em Sinkhorn, o SSR alcança:

• Melhor Estabilidade de Treinamento: Através de um balanceamento de especialistas integrado sem desalinhamento de objetivo.
• Maior Precisão: Demonstrada em benchmarks de texto e visão.
• Robustez: Desempenho aprimorado contra corrupção de input e exemplos adversários.
• Eficiência: Overhead computacional mínimo em comparação tanto com métodos de perda auxiliar quanto com o roteamento Sinkhorn completo.

Os autores enfatizam que o método é leve, não requer parâmetros treináveis adicionais (na variante de ruído) e mantém a flexibilidade da seleção de especialistas enquanto garante a utilização equilibrada durante a fase de treinamento.