Resumo Técnico: EMO - Pré-treinamento de Mixture of Experts para Modularidade Emergente

Declaração do Problema

Os Modelos de Linguagem de Grande Escala (LLMs) são tipicamente implantados como sistemas monolíticos, exigindo que o modelo completo seja carregado mesmo quando uma aplicação requer apenas um subconjunto restrito de capacidades (por exemplo, geração de código ou raciocínio matemático). Esse paradigma incorre em custos computacionais e uso de memória desnecessários. Embora arquiteturas de Mixture-of-Experts (MoE) ofereçam uma solução potencial ao ativar apenas um subconjunto de especialistas por entrada, os MoEs padrão não suportam implantação modular. Nos MoEs existentes, os tokens dentro de uma única entrada ativam especialistas diversos, frequentemente impulsionados por padrões léxicos de baixo nível (por exemplo, preposições, pontuação) em vez de domínios semânticos de alto nível. Consequentemente, restringir a inferência a um subconjunto específico de especialistas para um determinado domínio leva a uma degradação severa de desempenho, tornando os MoEs impraticáveis para ambientes com restrições de memória, onde apenas uma fração do modelo pode ser carregada.

Metodologia: EMO

Os autores introduzem o EMO, uma arquitetura MoE projetada para induzir modularidade como um objetivo de primeira classe durante o pré-treinamento. O objetivo central é treinar um modelo no qual subconjuntos de especialistas possam ser usados independentemente para domínios específicos e compostos em um modelo geral forte, sem depender de rótulos de domínio definidos por humanos.

Mecanismo Chave: Pools Compartilhados ao Nível de Documento

A intuição central é que os tokens dentro do mesmo documento geralmente compartilham um domínio semântico. O EMO aproveita as fronteiras de documento como um sinal de supervisão fraca para impor essa estrutura:

1. Pool de Especialistas Compartilhado: Para cada documento, o roteador seleciona um pool compartilhado de especialistas ($D$) do conjunto total de especialistas roteados.
2. Roteamento Restrito: Todos os tokens dentro desse documento são restringidos a selecionar seus especialistas ativos apenas a partir desse pool compartilhado.
3. Granularidade Dinâmica: Para garantir flexibilidade entre diferentes orçamentos de memória, o tamanho do pool de especialistas ($d$) é tratado como uma variável aleatória amostrada de uma distribuição uniforme $U\{k, n_r\}$ durante o pré-treinamento, onde $k$ é o número de especialistas ativos por token e $n_r$ é o total de especialistas roteados.

Objetivo de Treinamento e Considerações Técnicas

• Objetivo: O modelo é treinado usando perda padrão de modelagem de linguagem autoregressiva ($L_{CE}$) combinada com perdas auxiliares.
• Balanceamento de Carga: Uma contribuição técnica crítica é a adoção de balanceamento de carga global. O balanceamento de carga padrão em micro-lotes entra em conflito com a restrição de roteamento ao nível de documento (forçando tokens em um micro-lote a se espalharem entre especialistas). O EMO agrega estatísticas de roteamento entre grupos paralelos de dados para calcular a perda de balanceamento de carga, garantindo utilização uniforme de especialistas em todo o conjunto de dados, mantendo a consistência dentro dos documentos.
• Arquitetura: Os autores treinaram um modelo com 1 bilhão de parâmetros ativos e 14 bilhões de parâmetros totais, com 128 especialistas no total (127 roteados, 1 compartilhado), em 1 trilhão de tokens. A arquitetura incorpora melhorias sobre baselines anteriores, incluindo pré-norm, remoção da QK-norm e inclusão de um especialista compartilhado.

Contribuições Principais

1. Modularidade Emergente: O EMO demonstra que uma estrutura modular pode emergir diretamente dos dados usando fronteiras de documento como o único sinal, eliminando a necessidade de rótulos de domínio predefinidos ou priores humanos.
2. Uso Seletivo de Especialistas: O modelo permite a implantação de pequenos subconjuntos de especialistas (por exemplo, 12,5% a 25% dos especialistas) que retêm desempenho próximo ao do modelo completo em tarefas específicas de domínio, uma capacidade que os MoEs padrão não possuem.
3. Especialização Semântica: Diferentemente dos MoEs padrão que se especializam em padrões sintáticos de baixo nível, os especialistas do EMO se organizam em grupos coerentes especializados em domínios semânticos de alto nível (por exemplo, matemática, código, biomédica, política).
4. Melhoria da Fronteira de Pareto: Subconjuntos de especialistas do EMO superam modelos com correspondência de memória treinados do zero (tanto densos quanto MoE padrão), estabelecendo uma nova fronteira de Pareto para o trade-off memória-precisão.

Resultados

• Desempenho do Modelo Completo: Como modelo completo, o EMO iguala o desempenho de um MoE padrão treinado nos mesmos dados (1T de tokens), alcançando resultados competitivos em benchmarks como MMLU, MMLU-Pro e GSM8K.
• Desempenho de Subconjuntos:
  • Ao reter 25% dos especialistas, o EMO sofre apenas uma queda absoluta de 1% na precisão do MMLU, enquanto os MoEs padrão degradam-se em mais de 10%.
  • Ao reter 12,5% dos especialistas, o EMO sofre uma queda de 3%, enquanto os MoEs padrão caem 15%.
  • Os subconjuntos do EMO superam consistentemente modelos densos e MoEs padrão treinados do zero com a mesma contagem de parâmetros.
• Robustez: O desempenho do EMO é robusto entre diferentes métodos de seleção de especialistas (baseado em roteador vs. poda Easy-EP) e permanece eficaz mesmo com conjuntos de validação extremamente pequenos para seleção de especialistas.
• Análise: A análise de agrupamento revela que os tokens do EMO do mesmo documento são agrupados nos mesmos clusters semânticos (por exemplo, "saúde, medicina e bem-estar"), enquanto os tokens do MoE padrão estão dispersos entre clusters com base em características de superfície.

Significado e Alegações

O artigo alega que o EMO fornece um caminho para a implantação modular e eficiente em memória de modelos grandes e esparsos. Ao permitir o uso independente e a composição de subconjuntos de especialistas, o EMO aborda os gargalos de memória associados à escalabilidade dos MoEs. Os autores destacam várias implicações:

• Implantação Acessível: Os usuários podem servir apenas os especialistas relevantes para uma tarefa ou domínio específico, reduzindo significativamente os requisitos de memória.
• Controle de Alta Granularidade: A modularidade permite a ativação ou desativação seletiva de capacidades (por exemplo, excluindo especialistas associados a spam ou conteúdo adulto) no momento da inferência.
• Desenvolvimento Modular: A estrutura sugere um paradigma potencial para pré-treinamento modular, onde subconjuntos de especialistas específicos de tarefas poderiam ser treinados independentemente e posteriormente reintegrados.
• Monitorabilidade: Os padrões de roteamento distintos fornecem sinais estruturados para auditoria e depuração do comportamento do modelo.

Os autores concluem que a modularidade pode ser incorporada ao próprio pré-treinamento, afastando-se de sistemas monolíticos em direção a arquiteturas composáveis que são mais fáceis de implantar, adaptar e inspecionar.