Resumo Técnico: Quando a Fusão de Modelos Quebra o Roteamento: Calibração Livre de Treinamento para MoE

1. Declaração do Problema

A fusão de modelos tornou-se uma estratégia de custo-benefício atraente para consolidar as capacidades de múltiplos Grandes Modelos de Linguagem (LLMs) sem os custos proibitivos de um retreinamento conjunto. Embora as técnicas existentes (ex: Task Arithmetic, TIES-Merging) consigam fundir arquiteturas densas via aritmética linear de parâmetros, elas falham quando aplicadas a arquiteturas de Mistura de Especialistas (Mixture-of-Experts - MoE).

Os autores identificam um modo de falha crítico denominado quebra de roteamento (routing breakdown). Em modelos MoE, um roteador despacha dinamicamente tokens para um subconjunto de especialistas usando uma ativação softmax não linear seguida por uma seleção discreta Top-k. O artigo argumenta que os métodos de fusão linear padrão são mal adequados para este cenário porque:

1. Desajuste de Não-Linearidade: A interpolação linear dos pesos do roteador ($W_m = \sum \lambda_i W_i$) não se traduz em interpolação linear das probabilidades de roteamento devido à convexidade da função softmax.
2. Sensibilidade: Pequenas perturbações nos pesos do roteador são exponencialmente amplificadas pelos mecanismos softmax e Top-k.
3. Restrições de Equilíbrio de Carga: Os objetivos de pré-treinamento que impõem o equilíbrio de carga tornam as atribuições de especialistas altamente sensíveis a mudanças no gating.

Evidência Empírica: Experimentos preliminares fundindo dois modelos OLMoE (ajustados para matemática e código) revelaram que 50,11% das seleções de especialistas desviaram-se dos modelos de origem imediatamente após a fusão. Essa instabilidade se acumula através da profundidade da rede, com 99,98% dos tokens experienciando divergência de roteamento ao longo de 16 camadas. Crucialmente, a avaliação cruzada mostrou que restaurar o roteador original no modelo fundido recuperou o desempenho, enquanto aplicar o roteador fundido ao modelo original o degradou, provando que preservar os parâmetros dos especialistas é fútil sem um comportamento de roteamento consistente.

2. Metodologia: Calibração de Roteador Sensível à Hessiana (HARC)

Para abordar a quebra de roteamento, os autores propõem o HARC, um framework livre de treinamento que realinha o roteador fundido usando informações de curvatura de segunda ordem.

2.1. Formulação Teórica

O objetivo é encontrar uma matriz de pesos de roteador unificada $W_m$ que minimize a divergência de Kullback-Leibler (KL) entre a distribuição de saída do roteador fundido e as distribuições dos roteadores de origem em todas as tarefas.
$$W^*_m = \arg \min_{W_m} \sum_{i=1}^D \mathbb{E}_{x \sim \mathcal{D}_i} [\text{KL}(r_i(x) \parallel r_m(x; W_m))]$$

A otimização direta é intratável devido à não-linearidade do softmax. Os autores derivam uma aproximação quadrática da divergência KL usando uma expansão de Taylor de segunda ordem (Lema 4.1):
$$\text{KL}(r_i \parallel r_m) \approx \frac{1}{2}(z_m - z_i)^\top H_i (z_m - z_i)$$
onde $H_i = \text{diag}(r_i) - r_i r_i^\top$ é a Hessiana da função de partição logarítmica. Esta Hessiana captura a dinâmica competitiva entre especialistas (termos fora da diagonal) e a sensibilidade dos especialistas ativos (termos da diagonal).

2.2. Solução de Forma Fechada

Substituindo a relação linear $z = Wx$ na aproximação quadrática, o problema é transformado em um sistema de mínimos quadrados ponderados. Os pesos ótimos $w^*_m$ (forma vetorizada de $W_m$) satisfazem o sistema linear:
$$\left( \sum_{i=1}^D \Sigma_i \right) w^*_m = \sum_{i=1}^D \Sigma_i w_i$$
onde $\Sigma_i = \mathbb{E}_{x \sim \mathcal{D}_i} [H_i \otimes xx^\top]$ é a matriz de precisão específica da tarefa.

2.3. Computação Eficiente (Livre de Matriz)

A construção explícita e a inversão da matriz de precisão agregada $\Sigma$ são computacionalmente proibitivas ($O(K^2d^2)$ de complexidade de espaço). O HARC emprega duas estratégias principais para garantir a escalabilidade:

1. Regularização Diagonal: Para lidar com a deficiência de posto na Hessiana (devido ao gating esparso Top-k), os autores introduzem um termo de regularização que interpola entre a matriz de precisão total e sua aproximação diagonal. Isso estabiliza a solução sem assumir independência entre os especialistas.
2. Gradiente Conjugado (CG) Livre de Matriz: Em vez de materializar $\Sigma$, o algoritmo utiliza um operador livre de matriz para computar produtos matriz-vetor $\Sigma w$ on-the-fly usando identidades de produto de Kronecker. Isso reduz a complexidade de espaço para $O(Kd)$ e permite o uso do método de Gradiente Conjugado para resolver o sistema de forma eficiente.

2.4. Algoritmo

O HARC opera como uma etapa de pós-processamento em um modelo fundido não calibrado. Para cada camada MoE:

1. Realizar passagem direta (forward pass) dos dados de calibração pelo modelo fundido até a camada atual para gerar características de entrada.
2. Computar saídas de roteamento e Hessianas para os modelos de origem.
3. Acumular estatísticas para formar o sistema linear.
4. Resolver os pesos do roteador calibrado usando o solver CG livre de matriz.
5. Atualizar os pesos do roteador no modelo fundido.

3. Principais Contribuições

1. Identificação da Quebra de Roteamento: O artigo estabelece que a instabilidade de roteamento, e não a interferência de parâmetros de especialistas, é a causa primária da degradação de desempenho na fusão de MoE.
2. Framework HARC: Um novo método de calibração livre de treinamento que aproveita a informação da Hessiana de segunda ordem para realinhar as distribuições de roteamento, oferecendo uma solução de forma fechada para o problema de fusão não linear.
3. Otimização Escalável: A derivação de um solver de Gradiente Conjugado livre de matriz que permite que o HARC escale para grandes arquiteturas MoE sem a materialização explícita da Hessiana.
4. Validação Empírica: Demonstração de que o HARC mitiga efetivamente a quebra de roteamento em diversas bases de fusão (Weight Averaging, TIES-Merging, WUDI, etc.) e tarefas (Matemática, Código).

4. Resultados Experimentais

Os experimentos foram conduzidos no modelo OLMoE-1B-7B-0125, fundindo variantes ajustadas para raciocínio matemático (GSM8K, MATH500) e geração de código (HumanEval+, MBPP+).

• Melhoria Consistente: O HARC melhorou consistentemente o desempenho em todas as bases de fusão testadas. Por exemplo, ao ser aplicado à forte base WUDI-Merging, a pontuação geral aumentou de 38,08 para 38,96.
• Mitigação da Sensibilidade de Domínio: A fusão ingênua degradou severamente o desempenho de geração de código (ex: HumanEval+ caiu de 33,50 para 20,27 com Weight Averaging). O HARC recuperou parcialmente essa perda, demonstrando sua capacidade de preservar capacidades específicas de domínio ao corrigir a seleção de especialistas.
• Estudos de Ablação:
  • Estrutura da Hessiana: O uso da Hessiana total superou as aproximações diagonal ou identidade, confirmando que modelar a dinâmica competitiva fora da diagonal é crucial.
  • Regularização: O fator de amortecimento $\alpha$ (suprimindo termos fora da diagonal) provou ser mais estável do que o escalonamento diagonal $\gamma$.
  • Composição de Dados: A calibração usando respostas de origem filtradas como corretas (SrcCrt) produziu os melhores resultados. Surpreendentemente, usar apenas prompts teve um desempenho comparável ao de respostas completas, sugerando que as decisões iniciais de roteamento são primordiais.
• Eficiência de Dados: O HARC atingiu um desempenho próximo ao ótimo usando apenas 40% das amostras de calibração, indicando uma convergência rápida das estatísticas de segunda ordem.
• Análise por Camada: A visualização da divergência KL mostrou que os erros de roteamento nas bases acumulam-se exponencialmente nas camadas mais profundas. O HARC suprimiu com sucesso essa tendência, mantendo o alinhamento em camadas profundas críticas.

5. Significância e Alegações

O artigo alega que preservar a lógica de roteamento é tão essencial quanto alinhar os pesos dos parâmetros em arquiteturas esparsas. Ao abordar o modo de falha específico da quebra de roteamento, o HARC oferece um caminho robusto e eficiente para consolidar modelos MoE sem retreinamento.

Os autores enfatizam que seu método:

• Não requer retreinamento ou retropropagação de gradiente através dos especialistas.
• Generaliza-se através de diferentes estratégias de fusão (aritmética linear, otimização) como um passo de pós-processamento leve.
• Reduz os custos computacionais e as barreiras para implantar modelos capazes, permitindo o reuso confiável de múltiplos modelos de especialistas.

O trabalho conclui que, embora as hipóteses de fusão linear existentes (Conectividade de Modo Linear) funcionem para camadas densas, elas são frágeis para roteadores MoE. O HARC preenche essa lacuna ao considerar explicitamente a dinâmica de gating não linear inerente às arquiteturas de Mistura de Especialistas.