Resumo Técnico: DOT-MoE: Transporte Ótimo Diferenciável para MoEficação

1. Definição do Problema

O escalonamento rápido dos Grandes Modelos de Linguagem (LLMs) criou um gargalo significativo na eficiência da inferência. Embora as arquiteturas de Mistura de Especialistas (Mixture-of-Experts - MoE) decouplam o tamanho do modelo do custo de inferência ao ativar apenas um subconjunto de parâmetros por token, treinar MoEs do zero é notoriamente instável, faminto por dados e computacionalmente caro.

Consequentemente, a MoEficação — a conversão de modelos densos pré-treinados em MoEs esparsos — emergiu como uma alternativa viável. No entanto, os métodos existentes para converter Redes Feed-Forward (FFNs) densas em especialistas esparsos dependem fortemente de estratégias heurísticas (ex: divisão aleatória, agrupamento baseado em pesos ou agrupamento baseado em ativações). Essas abordagens sofrem de três limitações primárias:

1. Separação de Preocupações: Elas tipicamente tratam a atribuição de neurônios (particionamento da FFN) e o treinamento do roteador como estágios separados e sequenciais, impedindo que ambos se co-adaptem.
2. Otimização por Proxy: Elas otimizam com base em estatísticas intermediárias (pesos ou ativações) em vez da contribuição real de saída dos neurônios, levando frequentemente a uma reconstrução subótima do comportamento do modelo denso.
3. Falta de Diferenciabilidade: A natureza discreta da atribuição de neurônios impede o fluxo de gradiente de ponta a ponta, tornando difícil otimizar conjuntamente a decomposição estrutural e a política de roteamento.

2. Metodologia: DOT-MoE

Os autores propõem o DOT-MoE, um framework que formula a decomposição de camadas densas em especialistas como um problema de Transporte Ótimo Diferenciável (DOT). Esta abordagem permite o aprendizado de ponta a ponta tanto da decomposição dos especialistas quanto do mecanismo de roteamento.

2.1. Formulação de Transporte Ótimo

A ideia central é visualizar a atribuição de neurônios como um problema de transporte equilibrado onde a "massa" (neurônios) deve ser transportada de uma distribuição de origem (os neurônios da FFN densa) para uma distribuição de destino (os especialistas) sob estritas restrições de capacidade.

• Origem: $dffn$ neurônios, cada um com unidade de massa.
• Destino: $E$ especialistas, cada um exigindo exatamente $s$ neurônios ($s \cdot E = dffn$).
• Função de Custo: Em vez de uma matriz de custo fixa, o método utiliza uma matriz de afinidade aprendível $A$, onde $A_{i,e}$ representa a afinidade de atribuir o neurônio $i$ ao especialista $e$. O objetivo é maximizar a afinidade total enquanto satisfaz as restrições marginais.

2.2. Relaxação Diferenciável via Sinkhorn-Knopp

Para tornar o problema de atribuição discreta diferenciável, o DOT-MoE emprega Regularização Entrópica.

• O problema de transporte ótimo é relaxado usando o algoritmo Sinkhorn-Knopp, que normaliza iterativamente as linhas e colunas da matriz de plano de transporte.
• Isso gera uma matriz de atribuição suave $M_{soft} \in [0, 1]^{dffn \times E}$ que satisfaz as restrições de capacidade em expectativa.
• À medida que o parâmetro de temperatura $\tau \to 0$, a solução se aproxima do ótimo discreto.

2.3. Otimização Conjunta com Estimadores de Passagem Direta (STE)

Um desafio crítico é que a implantação requer atribuições discretas, e o roteamento envolve seleção top-k discreta, ambos não diferenciáveis. O DOT-MoE aborda isso usando Estimadores de Passagem Direta (Straight-Through Estimators - STE):

• Passagem para Frente (Forward Pass): Utiliza decisões rígidas (arredondamento guloso de $M_{soft}$ para atribuição; seleção top-k para roteamento).
• Passagem para Trás (Backward Pass): Roteia gradientes através dos equivalentes suaves ($M_{soft}$ e os logits do roteador $P$).
• Isso permite que os pesos do roteador ($W_r$) e a matriz de afinidade de atribuição ($A$) se co-adaptem durante o treinamento.

2.4. Objetivo de Treinamento

O framework é treinado via uma Fase de Alinhamento onde os pesos do modelo denso são congelados. O modelo minimiza uma combinação de perdas:

1. Reconstrução de Saída: Medida via divergência KL entre as saídas do professor denso e do aluno MoE esparso.
2. Modelagem de Linguagem: Perda de entropia cruzada sobre o objetivo de modelagem de linguagem.
3. Auxiliares de Estabilidade: Perda $z$ do roteador (para evitar explosão de logits) e perda de balanceamento de carga (para evitar o colapso de especialistas).

3. Principais Contribuições

1. Framework de Transporte Ótimo: O artigo introduz uma formulação inovadora da conversão de denso-para-MoE como um problema de transporte equilibrado, resolvido via iterações Sinkhorn diferenciáveis para impor estritas restrições de capacidade de especialistas.
2. Co-Adaptação de Dois Níveis: Ao contrário de métodos anteriores que fixam as atribuições antes de treinar os roteadores, o DOT-MoE otimiza conjuntamente a decomposição neurônio-especialista e o roteamento de token-especialista usando STEs complementares.
3. Construção Consciente da Saída: Ao modelar diretamente a contribuição do neurônio para a saída final (em vez de otimizar proxies como pesos ou ativações intermediárias), o método alcança maior fidelidade de reconstrução.

4. Resultados Experimentais

Os autores avaliaram o DOT-MoE em três famílias de modelos (LLaMA-2, LLaMA-3, Qwen2.5) e seis benchmarks (ARC-Challenge, Winogrande, HellaSwag, PIQA, SciQ, BoolQ).

• Desempenho vs. Poda (Pruning): No LLaMA-2 7B com um orçamento de 50% de parâmetros, o DOT-MoE alcançou uma perplexidade de 7.99 no WikiText-2, superando significativamente o método de poda estruturada de última geração DISP-LLM (9.84).
• Desempenho vs. MoEficação Existente: Em configurações zero-shot, o DOT-MoE superou substancialmente métodos como CMoE e LLaMA-MoE-v2. Por exemplo, no LLaMA-2 7B, o DOT-MoE alcançou 61.5% de acurácia média comparado aos 44.5% do CMoE.
• Eficiência: O DOT-MoE retém aproximadamente 90% do desempenho do modelo denso original enquanto reduz os parâmetros ativos em 50%.
• Ajuste Fino (Fine-Tuning): Com apenas 1.2B de tokens de ajuste fino contínuo, os modelos DOT-MoE recuperaram um desempenho próximo aos contrapartes densos (ex: 66.6% vs 76.6% no LLaMA-2 7B), superando baselines que exigiram dados similares ou maiores.
• Estudos de Ablação:
  • Granularidade: O desempenho melhora com o aumento da granularidade dos especialistas até um ponto de saturação, após o qual os ganhos diminuem.
  • Velocidade de Inferência: Apesar do maior número de especialistas, o throughput de inferência permanece estável ao usar kernels de MoE fundidos (ex: vLLM), pois a dimensão intermediária total fundida permanece constante.
  • Esparsidade: Modelos treinados em níveis de esparsidade mais altos (75%) generalizam robustamente para regimes de esparsidade ainda mais extremos no tempo de inferência.

5. Significância e Alegações

O artigo afirma que o DOT-MoE representa um avanço significativo na eficiência da inferência de LLMs ao fornecer um método diferenciável e principiado para converter modelos densos em MoEs esparsos.

• Mudança de Paradigma: Ele se afasta de partições heurísticas e estáticas em direção a uma decomposição dinâmica e aprendível que respeita a coerência funcional dos especialistas.
• Praticidade: O método permite a esparsificação pós-hoc eficiente de modelos pré-treinados sem a necessidade de retreinamento completo do zero, retendo alta capacidade enquanto reduz drasticamente os parâmetros ativos.
• Robustez: O framework demonstra que preservar a fidelidade da FFN requer modelar explicitamente a contribuição do neurônio para a saída, um princípio que os métodos baseados em proxy existentes falham em capturar.

Os autores concluem que o DOT-MoE oferece uma inicialização superior para modelos esparsos, diminuindo a lacuna de desempenho entre arquiteturas densas e esparsas com um volume mínimo de dados de treinamento adicionais.