Temporally Extended Mixture-of-Experts Models

Título: Modelos de Mistura de Especialistas Temporalmente Estendidos

Autores: Zeyu Shen e Peter Henderson (Princeton University)

1. O Problema: O Custo de Troca em MoEs Modernos

Os Grandes Modelos de Linguagem (LLMs) modernos, como o GPT-OSS e o Qwen, utilizam amplamente camadas de Mistura de Especialistas (MoE). Nessas arquiteturas, apenas um subconjunto esparso de "especialistas" (sub-redes) é ativado para cada token gerado, permitindo escalar o número total de parâmetros sem aumentar o custo computacional de inferência por token.

No entanto, o artigo identifica um gargalo crítico de memória e latência:

• Troca Constante: Em MoEs padrão, o conjunto de especialistas ativos muda em quase cada token gerado (taxa de troca > 50-90%).
• Limitação de Memória: Quando o número total de especialistas excede a memória da GPU, os pesos devem ser descarregados para a memória do host (RAM) ou disco e carregados sob demanda.
• Ineficiência de Otimizações: A troca frequente torna ineficazes otimizações como prefetching (pré-busca) e offloading (descarregamento), pois o sistema não consegue prever quais especialistas serão necessários nos próximos passos. Isso gera latência significativa e reduz o throughput.

2. Metodologia: Enquadramento no Framework de Opções

Os autores propõem tratar a seleção de especialistas não como uma decisão token-a-token, mas como uma decisão temporalmente estendida, utilizando o Framework de Opções da Aprendizagem por Reforço (RL).

Conceitos Chave:

• Opções (Options): Um conjunto de especialistas permitidos (uma "máscara de especialistas") atua como uma "opção" de alto nível que persiste por vários tokens.
• Controlador (Controller): Um módulo leve é adicionado a cada camada do modelo. Ele decide:
  1. Quando trocar: Se a opção atual deve ser encerrada (terminação).
  2. Para qual trocar: Qual novo conjunto de especialistas deve ser carregado.
• Custo de Deliberação (Deliberation Cost - $\eta$): Um termo explícito na função de recompensa que penaliza a troca de opções. Isso força o controlador a aprender a manter a mesma máscara de especialistas por mais tempo, a menos que a melhoria esperada na qualidade da geração justifique o custo de latência da troca.

Arquitetura e Treinamento:

• Base: O método é aplicado ao modelo gpt-oss-20b (24 camadas, 32 especialistas por camada, top-4).
• Controlador Leve: Utiliza uma arquitetura baseada em Option-Critic.
  • Embedding: Usa DeepSets para codificar o conjunto atual de especialistas.
  • Heads: Cabeças de terminação (Bernoulli), seleção (Plackett-Luce) e valor (V/Q).
• Treinamento (Fine-tuning):
  • Utiliza LoRA (Low-Rank Adaptation) para ajustar os parâmetros do modelo e do controlador com baixo custo.
  • Recompensa: Baseada em Self-Distillation (KL reverso) entre o modelo estudante (com controlador) e o modelo professor (original congelado). Isso garante que a qualidade da geração seja mantida.
  • Algoritmo: Adaptação do algoritmo A2OC (Actor-Critic com Opções) com custo de deliberação.

3. Contribuições Principais

1. Mudança de Paradigma: Propõe a filosofia de "MoEs Temporalmente Estendidos", identificando que a falta de continuidade temporal nas rotas atuais impede otimizações de memória.
2. Formalização Matemática: Modela o carregamento dinâmico de especialistas como um Processo de Decisão Semi-Markoviano (s-MDP), onde máscaras de especialistas são opções e a latência de carregamento é o custo de deliberação.
3. Viabilidade em Pós-Treinamento: Demonstra que modelos pré-treinados existentes podem ser convertidos em MoEs temporalmente estendidos com treinamento leve (apenas LoRA e um controlador), sem necessidade de pré-treinamento massivo do zero.
4. Trade-off Controlável: Introduz um hiperparâmetro ($\eta$) que permite aos pesquisadores e engenheiros ajustar o equilíbrio entre a taxa de troca (eficiência de memória) e a capacidade do modelo (precisão).

4. Resultados Experimentais

Os experimentos foram realizados no gpt-oss-20b com orçamentos de especialistas ($\hat{k}$) de 8 e 16 (de um total de 32).

• Redução Drástica na Taxa de Troca:
  • Base: > 50% de troca por token (mudança quase a cada token).
  • Método Proposto: Redução para < 5% (e até 1% em configurações mais agressivas), mantendo a mesma máscara de especialistas por longos períodos.
• Desempenho (Precisão):
  • Com $\hat{k}=16$ e custo de deliberação baixo ($\eta=0.02$), o modelo atingiu ~90% da precisão do modelo base em benchmarks como MATH, MMLU e MMMLU.
  • Exemplo: No MATH, o modelo base teve 71.5%, enquanto o método proposto alcançou 64.0% (vs. 53.5% para pruning por frequência e 15% para seleção aleatória).
• Estabilidade: O método evitou o colapso de repetição (catastrophic repetition), um problema comum em modelos com roteamento restrito, mantendo a coerência do texto.
• Visualização: As figuras mostram que, ao contrário do modelo base (onde especialistas mudam a cada token), o modelo com controlador mantém o mesmo conjunto de especialistas ativo por centenas de tokens.

5. Significado e Impacto Futuro

Este trabalho abre um caminho principiado para a serviço eficiente em memória e aprendizado contínuo de modelos MoE cada vez maiores:

1. Inferência com Memória Reduzida: Ao manter um conjunto fixo de especialistas por longos períodos, é possível servir modelos com muitos especialistas usando apenas uma fração da VRAM (ex: carregar apenas 16 especialistas em vez de 128), reduzindo drasticamente a latência de E/S.
2. Treinamento por "Chunks": Permite dividir o treinamento em blocos contíguos onde apenas um subconjunto de especialistas é ativo, permitindo o treinamento de modelos maiores do que a memória disponível.
3. Aprendizado Contínuo: Facilita a adição de novos especialistas ao pool sem aumentar o custo de inferência por token, pois o controlador pode aprender a ativá-los apenas quando necessário.

Em resumo, o artigo demonstra que a introdução de estrutura temporal nas decisões de roteamento de especialistas, guiada por custos de troca explícitos, resolve um dos maiores obstáculos para a escalabilidade prática de modelos MoE massivos.