Affinity Is Not Enough: Recovering the Free Energy Principle in Mixture-of-Experts

1. Declaração do Problema

Modelos de Mistura de Especialistas (MoE) esparsos (por exemplo, DeepSeek-V3, Mixtral) atualmente dependem de roteamento baseado em afinidade sem estado. O portão de roteamento seleciona especialistas com base exclusivamente no embedding do token atual ($x_t$), ignorando o contexto temporal, o histórico de confiabilidade dos especialistas e os requisitos futuros.

O artigo identifica um modo de falha crítico: Transições de Domínio. Quando uma sequência muda de uma distribuição (Domínio A) para outra (Domínio B), o token no passo de transição ($t$) ainda pertence ao Domínio A, mas o especialista ótimo para o próximo passo ($t+1$) pertence ao Domínio B.

• O Problema: Roteadores padrão, vendo apenas $x_t$, atribuem probabilidade próxima de zero ao especialista correto no passo de transição.
• Evidência Empírica: Em experimentos controlados, o roteamento padrão atribui apenas 0,006 ± 0,001 de massa de probabilidade ao especialista correto nos passos de transição.
• Consequência: Para alcançar 99% de cobertura de roteamento com tal baixa confiança, um modelo exigiria um número inviavelmente grande de especialistas ($K \approx 765$), anulando os benefícios de eficiência do MoE.

2. Metodologia: O Framework do Princípio da Energia Livre (FEP)

O autor propõe que o roteamento MoE deve ser reenquadrado através do Princípio da Energia Livre de Friston, que dita que um sistema ótimo sob incerteza deve:

1. Atualizar o estado recursivamente (não do zero).
2. Ponderar sinais por precisão (variância inversa do erro).
3. Minimizar a energia livre futura esperada (agir sobre previsões, não apenas observações atuais).

Para implementar isso, o artigo introduz três modificações leves ao portão de roteamento padrão, fundamentadas na dinâmica de Redes Neurais de Spiking (SNN) (especificamente Leaky Integrate-and-Fire, ou LIF):

A. Memória de Roteamento Temporal ($\beta$)

• Mecanismo: Substitui a entrada sem estado $x_t$ por um estado oculto $h_t$ atualizado via dinâmica LIF:
  $$h_t = \beta \odot h_{t-1} + x_t$$
  onde $\beta$ é um vetor de decaimento aprendível (inicializado em 0,9).
• Função: Acumula contexto através dos tokens. Permite que o roteador "lembre" que o sistema esteve no Domínio A por muito tempo, criando um sinal de "saturação" que indica uma troca iminente, mesmo antes que o token de troca apareça.
• Escala de Tempo: Por token (rápido).

B. Portão Ponderado por Precisão ($\Pi$)

• Mecanismo: Mantém um rastreador online e não paramétrico do Erro Quadrático Médio (MSE) para cada especialista. O portão de roteamento é modulado pela variância inversa (precisão) desses erros:
  $$g_t = \text{softmax}(W \cdot x_t \odot \Pi_t)$$
  onde $\Pi_i = 1 / (\hat{\sigma}^2_i + \epsilon)$.
• Função: Pondera dinamicamente para baixo especialistas que são atualmente pouco confiáveis (alto erro) e para cima os confiáveis. Diferente do balanceamento de carga padrão (que rastreia frequência), isso rastreia precisão.
• Escala de Tempo: Por lote (médio).

C. Roteamento Antecipatório

• Mecanismo: Usa uma rede preditora $f(x_t, h_t)$ para estimar o próximo embedding de token $\hat{x}_{t+1}$. A decisão de roteamento é condicionada a essa previsão:
  $$g_t = \text{softmax}(W \cdot h_t + W_{pred} \cdot \hat{x}_{t+1})$$
• Função: Roteia com base no estado futuro previsto em vez do estado atual. Isso permite que o roteador mude para o especialista do novo domínio antes que os tokens do novo domínio apareçam na entrada.
• Insight Crucial: Um preditor sem estado ($f(x_t)$) falha porque os tokens pré-transição são distribucionalmente idênticos. O preditor deve ser condicionado ao estado acumulado por $\beta$, $h_t$, para detectar o sinal de saturação indicando uma transição.

3. Contribuições Principais

1. Unificação Teórica: Conecta o roteamento MoE ao Princípio da Energia Livre, demonstrando que o MoE padrão carece de estimativa de estado recursiva, ponderação por precisão e minimização de estado futuro.
2. Três Mecanismos Leves: Propõe $\beta$ (memória), $\Pi$ (precisão) e Antecipação (previsão) como componentes necessários, implementáveis com sobrecarga mínima de parâmetros (~200 linhas de código no total).
3. Descoberta de Super-Aditividade: A descoberta mais significativa do artigo é a interação $\beta \times \text{Ant}$.
   • $\beta$ sozinho: ganho de +0,295.
   • Antecipação sozinho: ganho de +0,000 (falha sem memória).
   • Combinado: ganho de +0,741.
   • Efeito de Interação: O ganho combinado excede a soma das partes individuais em +0,446, provando que a memória é um pré-requisito para a antecipação.
4. Explicação Estrutural para DeepSeek-V3: Explica por que o DeepSeek-V3 descarta seu módulo de Previsão de Múltiplos Tokens (MTP) na inferência: um cabeçalho MTP sem estado não pode detectar transições de domínio. O artigo argumenta que o MTP deve ser mantido se condicionado à memória temporal ($\beta$).

4. Resultados Experimentais

Os experimentos foram conduzidos em tarefas controladas de troca de domínio e em um modelo de linguagem MoE em nível de caractere (5 sementes).

• Precisão de Transição: O mecanismo combinado eleva a probabilidade de selecionar o especialista correto no passo de transição de 0,6% para 74,8%, fechando 75% da lacuna até o oráculo.
• Eficiência: O número necessário de especialistas ($K$) para alcançar 99% de cobertura cai de um ~765 inviável para uma constante pequena (2,4) nos passos de transição.
• Validação do Modelo de Linguagem: Em um LM em nível de caractere, o roteamento $\beta$ reduziu os Bits por Caractere (BPC) no passo de transição de 6,56 para 4,01. Adicionar antecipação estabilizou a decisão de roteamento (reduziu a variância), mas não reduziu adicionalmente o BPC nesta escala pequena, sugerindo que o benefício primário é a estabilidade de roteamento e a redução de latência via pré-busca.

5. Significado e Implicações

• Navegação vs. Classificação: O artigo reenquadra o roteamento de um problema estático de classificação ("qual especialista se encaixa neste token?") para um problema dinâmico de navegação ("para onde o sistema está indo e qual caminho é confiável?").
• Redução de Latência: Ao antecipar trocas de domínio um passo à frente, o sistema pode pré-buscar pesos de especialistas da memória (HBM) para unidades de computação (SRAM) enquanto o passo atual está sendo executado, eliminando travamentos de latência de memória em pontos críticos de transição.
• Reavaliação do MTP: O trabalho sugere que módulos de Previsão de Múltiplos Tokens em modelos como DeepSeek-V3 estão subutilizados. Se condicionados ao estado temporal, eles poderiam servir como poderosos mecanismos de roteamento antecipatório, em vez de apenas auxílios de treinamento.
• Plausibilidade Biológica: Os mecanismos espelham processos biológicos (neurônios LIF, ponderação por precisão em codificação preditiva, inferência ativa), sugerindo que as arquiteturas MoE atuais estão perdendo princípios fundamentais de alocação inteligente de recursos.

Em conclusão, o artigo demonstra que a afinidade é insuficiente para um roteamento MoE robusto. Ao integrar memória temporal, ponderação por precisão e antecipação, modelos MoE podem alcançar desempenho de roteamento próximo ao oráculo em transições de domínio, reduzindo drasticamente o custo computacional de manter alta precisão.