Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

1. O Problema

Os Modelos de Linguagem de Difusão (DLMs) emergiram como uma alternativa promissora aos modelos autoregressivos, permitindo a geração de texto paralela e não autoregressiva. No entanto, ao adotar arquiteturas de Mistura de Especialistas (MoE) para escalar para bilhões de parâmetros, os DLMs existentes herdaram diretamente o mecanismo de roteamento Token-Choice (TC) dos modelos autoregressivos.

O roteamento TC apresenta duas falhas críticas no contexto de difusão:

• Desequilíbrio de Carga (Load Imbalance): Como cada token escolhe independentemente seus especialistas preferidos, alguns especialistas ficam sobrecarregados enquanto outros permanecem ociosos. Isso exige uma função de perda auxiliar para balanceamento, que consome recursos computacionais adicionais e introduz sinais de gradiente conflitantes. Mesmo assim, o equilíbrio permanece frágil, causando gargalos de "straggler" (um GPU lento atrasa todo o lote) e reduzindo a eficiência do hardware.
• Alocação Rígida de Computação: O roteamento TC trata a alocação de especialistas como uma consequência emergente das escolhas dos tokens, sem permitir o controle direto de quanto cálculo é dedicado a diferentes etapas do processo de difusão.

2. Metodologia

Os autores propõem a adoção do roteamento Expert-Choice (EC) e a introdução de uma estratégia de agendamento de capacidade dependente do tempo (timestep-dependent).

• Roteamento Expert-Choice (EC): Inverte a lógica do TC. Em vez de os tokens escolherem especialistas, cada especialista seleciona um número fixo e determinístico de tokens (os que têm maior pontuação de gate).
  • Vantagem: Garante um balanceamento de carga perfeito por design, eliminando a necessidade de perdas auxiliares e removendo o problema de "straggler" em GPUs.
• Capacidade Dependente do Passo de Tempo (Timestep-Adaptive Capacity): Aproveitando que a capacidade do especialista no EC é uma variável de design externa (e não emergente), os autores propõem variar a capacidade ($k$) em função da razão de mascaramento ($r$) em cada passo de difusão.
  • Eles testam várias estratégias de agendamento (linear, cosseno, gaussiano e suas reversas) para alocar mais especialistas em passos onde o retorno marginal do cálculo é maior.
  • A hipótese central é que passos com baixa razão de mascaramento (onde a maioria dos tokens já é visível e o modelo precisa refinar previsões precisas) têm uma eficiência de aprendizado muito superior à de passos com alta razão de mascaramento.

3. Contribuições Principais

1. Superioridade do EC para DLMs: Demonstram que o roteamento EC é estritamente superior ao TC para o treinamento de DLMs MoE, oferecendo melhor balanceamento de carga, maior throughput e convergência mais rápida.
2. Agendamento Adaptativo de Capacidade: Introduzem e validam a estratégia de variar a capacidade dos especialistas ao longo dos passos de difusão. A descoberta chave é que alocar mais capacidade para passos de baixa razão de mascaramento (perto do final do processo de denoising) produz o maior retorno em termos de desempenho.
3. Análise Mecanística: Fornecem uma explicação teórica e empírica: tokens em contextos de baixa razão de mascaramento exibem uma eficiência de aprendizado uma ordem de magnitude maior. Concentrar computação nesses passos maximiza o retorno marginal.
4. Retrofitting de Modelos Existentes: Mostram que modelos DLM pré-treinados com roteamento TC podem ser convertidos para EC substituindo apenas o roteador (mantendo os pesos dos especialistas), resultando em convergência mais rápida e melhor precisão em tarefas de fine-tuning.

4. Resultados Experimentais

• Eficiência de Treinamento: Em experimentos controlados, o modelo com EC alcançou uma perda de treinamento de 3.75 em 10.6 horas, sendo 2.0x mais rápido que as melhores configurações de TC (que levaram ~20.7 horas). O throughput do EC foi de 52.1 TFLOP/s/GPU, superando os variantes de TC em 1.5x a 2.1x.
• Balanceamento de Memória: Enquanto o TC apresentou alta variância no uso de memória entre GPUs (desvio padrão de 3.6 GB, com um GPU usando 70.3 GB), o EC manteve um uso perfeitamente uniforme (desvio padrão de 0.0 GB).
• Desempenho de Agendamento: Sob condições de FLOPs iguais, o agendamento Linear-Reverso (que aumenta a capacidade à medida que a razão de mascaramento diminui) superou consistentemente todas as outras estratégias, incluindo o agendamento estático e os que favorecem alta razão de mascaramento.
• Validação em Escala: Em modelos de 8B parâmetros (1B ativos) pré-treinados no conjunto Nemotron-CC, o EC dinâmico superou o EC estático em perplexidade e precisão em benchmarks downstream (MMLU, ARC-Challenge).
• Retrofitting: Ao converter o modelo pré-treinado LLaDA-MoE (TC) para EC, observou-se convergência mais rápida em fine-tuning para tarefas de código (HumanEval), matemática (GSM8K) e medicina (MedQA), com ganhos de precisão e velocidade de inferência (1.3x a 1.5x mais rápido na decodificação).

5. Significado e Conclusão

Este trabalho estabelece o roteamento Expert-Choice como o paradigma superior para modelos MoE de linguagem baseados em difusão. Ele resolve o problema estrutural de desequilíbrio de carga inerente ao roteamento Token-Choice, eliminando a necessidade de perdas auxiliares e melhorando drasticamente a eficiência do hardware.

Além disso, o trabalho demonstra que a computação em DLMs não deve ser tratada como uma constante arquitetônica fixa, mas sim como uma política adaptativa. A capacidade de ajustar dinamicamente a alocação de especialistas com base no passo de tempo (especificamente focando nos passos de baixa razão de mascaramento) permite extrair o máximo de eficiência do modelo. A capacidade de "retrofitar" modelos existentes apenas trocando o roteador torna essa abordagem altamente prática e imediatamente aplicável ao ecossistema atual de modelos de difusão.