Teacher-Guided Routing for Sparse Vision Mixture-of-Experts

1. O Problema: Instabilidade no Treinamento de MoE Esparsos

O artigo aborda um desafio crítico no treinamento de arquiteturas Mistura de Especialistas (MoE) esparsas, especificamente em modelos de visão computacional (VMoE). Embora o MoE permita escalar a capacidade do modelo ativando apenas um subconjunto de especialistas para cada entrada (reduzindo o custo computacional), o treinamento dessas redes enfrenta dificuldades de otimização:

• Bloqueio de Gradientes (Gradient Blocking): Em um MoE padrão, o "roteador" (a rede que decide qual especialista usar) recebe gradientes informativos apenas dos especialistas selecionados durante a passagem direta (forward pass). As rotas não selecionadas não enviam feedback.
• Flutuação de Roteamento (Routing Fluctuation): Devido à falta de feedback global e dependente dos dados, o roteador frequentemente altera suas decisões de alocação de especialistas de forma excessiva e instável durante o treinamento. Um mesmo token pode ser atribuído a especialistas diferentes em iterações consecutivas.
• Consequências: Essa instabilidade impede a especialização estável dos especialistas, leva a um sub-treinamento de alguns módulos e dificulta a convergência, especialmente em configurações altamente esparsas ou com muitos especialistas.

2. Metodologia: TGR-MoE (Teacher-Guided Routing)

Os autores propõem o TGR-MoE, um método que estabiliza o aprendizado do roteador utilizando a supervisão derivada de um modelo denso pré-treinado (o "Professor").

Arquitetura e Funcionamento:

1. Construção do Roteador Professor:

   • Um modelo denso pré-treinado (ex: DeiT-III) é mantido congelado.
   • Um roteador auxiliar leve (Teacher Router) é anexado às representações intermediárias desse modelo professor.
   • Este roteador professor é otimizado para produzir distribuições de roteamento estáveis e balanceadas, utilizando duas funções de perda:
     • Perda de Balanceamento de Carga ($L_{load}$): Garante que todos os especialistas sejam utilizados uniformemente.
     • Perda de Entropia ($L_{ent}$): Impede que o roteador colapse em uma distribuição excessivamente uniforme, incentivando atribuições confiantes e distintas.

2. Treinamento Guiado do Aluno (Student Router):

   • O modelo MoE estudante (aluno) é treinado para imitar a distribuição de roteamento do professor.
   • A função de perda total do estudante combina a perda da tarefa principal (ex: classificação) com uma perda de destilação ($L_{distill}$) baseada na Divergência de Kullback-Leibler (KL) entre a distribuição do roteador do aluno e a do professor.
   • Vantagem Chave: O professor fornece um "prior" global e estável. Isso preenche a lacuna de informação deixada pelo bloqueio de gradientes, guiando o roteador do aluno mesmo para especialistas que não foram selecionados na etapa direta.

3. Eficiência:

   • O backbone do professor permanece congelado; apenas o roteador professor e o modelo estudante são otimizados.
   • Sem custo adicional na inferência: O roteador professor é usado apenas durante o treinamento. Na fase de teste, apenas o modelo estudante é utilizado.

3. Contribuições Principais

• Solução para Instabilidade: Introduz um mecanismo simples e eficaz para estabilizar a dinâmica de roteamento em MoEs esparsos, mitigando o problema de flutuação de especialistas.
• Supervisão Externa: Demonstra que a incorporação de um prior de roteamento derivado de um modelo denso pré-treinado é superior a métodos que dependem apenas de gradientes esparsos ou regularizações de balanceamento de carga.
• Escalabilidade: O método permite que o MoE escale de forma mais confiável com o aumento do número de especialistas, mantendo ganhos de desempenho onde os métodos padrão saturam.
• Consistência na Transferência: Mostra que o conhecimento de roteamento aprendido durante o pré-treinamento é preservado mais fielmente durante o fine-tuning quando o TGR-MoE é mantido ativo.

4. Resultados Experimentais

Os autores avaliaram o TGR-MoE em ImageNet-1K (pré-treinamento) e em tarefas de transferência (CIFAR-10/100, Oxford-IIIT Pets), comparando com VMoE padrão, Expert Choice, SoftMoE e regularização z-loss.

• Desempenho de Precisão:
  • O TGR-MoE superou consistentemente o VMoE padrão em todas as escalas de modelo (Tiny, Small, Base).
  • No ImageNet-1K, melhorou a acurácia Top-1 do modelo Tiny de 77,85% para 78,78% e do Small de 82,63% para 83,34%.
  • Foi competitivo ou superior a variantes avançadas como SoftMoE e Expert Choice.
• Estabilidade e Consistência de Roteamento:
  • Análise de consistência mostrou que o TGR-MoE atinge acordos de alocação de especialistas (>70%) muito mais cedo (após ~50 épocas) em comparação ao VMoE, que oscila significativamente.
  • Durante o fine-tuning, o TGR-MoE manteve uma taxa de acordo de 73,75% entre as configurações de roteamento pré-treinadas e ajustadas, enquanto o VMoE padrão caiu para 50,56%.
• Escalabilidade com Número de Especialistas:
  • Ao aumentar o número de especialistas (de 4 a 256), o VMoE padrão viu seus ganhos estagnarem, enquanto o TGR-MoE continuou a apresentar melhorias lineares e estáveis.
• Análise de Ablação:
  • A destilação sozinha já fornece supervisão suficiente para roteamento estável.
  • A melhor estratégia foi aplicar a destilação principalmente na primeira metade do treinamento (para estabilizar o roteamento inicial) e depois focar na otimização da tarefa.

5. Significado e Conclusão

O trabalho demonstra que o roteamento em MoEs esparsos não deve depender exclusivamente de gradientes locais e esparsos. Ao introduzir um guia global estável proveniente de um modelo denso pré-treinado, é possível:

1. Resolver o problema de bloqueio de gradientes no roteador.
2. Permitir que os especialistas se especializem de forma mais rápida e consistente.
3. Alcançar arquiteturas mais escaláveis e robustas sem custo computacional adicional na inferência.

O TGR-MoE estabelece um novo paradigma para o treinamento de modelos esparsos, sugerindo que a "sabedoria" de modelos densos pré-treinados pode ser transferida não apenas para as camadas de características, mas especificamente para a lógica de decisão de alocação de recursos (roteamento), resultando em modelos mais eficientes e precisos.