o problema
Em 2020 já estava claro que escalar Transformer denso funciona: mais parâmetros, mais dados, mais compute, e a perda cai numa lei de potência previsível. O detalhe incômodo é que num modelo denso essas três coisas andam grudadas. Dobrar os parâmetros dobra a conta de multiplicação de matrizes por token, no treino e na inferência. O custo computacional vira o teto do tamanho.
Mixture-of-Experts existia desde 1991 e prometia desatar esse nó: em vez de aplicar todos os pesos a todo exemplo, o modelo escolhe quais pesos usar para cada entrada. Shazeer e coautores mostraram em 2017 que isso funcionava em linguagem, com camadas MoE entre LSTMs. Mas a adoção travou em três coisas bem concretas, e o paper é honesto ao listá-las: complexidade de implementação, custo de comunicação entre dispositivos e instabilidade de treino. MoE era um paper citado, não uma arquitetura que as pessoas usavam.
a ideia
O Switch Transformer troca a camada feed-forward do Transformer por uma camada com N experts — cada um uma FFN independente — e um roteador que manda cada token para um expert só.
Isso soa como uma simplificação óbvia, mas contraria o que o próprio Shazeer havia argumentado: que era preciso rotear para pelo menos dois experts, senão não haveria gradiente para o roteador aprender comparação. Fedus, Zoph e Shazeer mostram que não. O valor do gate do expert escolhido multiplica a saída, e é por aí que o gradiente entra no roteador. Um expert basta.
O ganho de mandar para um só é triplo e todo prático: o roteador computa menos, a comunicação encolhe, e o buffer que cada expert precisa reservar cai pela metade.
como funciona
O roteador é uma matriz W_r que projeta o token em N logits, seguidos de softmax. Pega-se o argmax. Custo O(d_model × num_experts) — desprezível ao lado da FFN.
O problema é que TPU exige shapes estáticos e o roteamento é dinâmico. A solução é a expert capacity: (tokens por batch / número de experts) × capacity factor. Cada expert processa no máximo esse tanto. Se muitos tokens caem no mesmo expert, os excedentes são dropados — pulam a camada inteira e seguem pela conexão residual. Capacity factor maior alivia o descarte e desperdiça compute em slots vazios. Na prática o paper mantém o descarte abaixo de 1%.
O equilíbrio vem de uma loss auxiliar: α · N · Σ f_i · P_i, onde f_i é a fração de tokens que foi para o expert i e P_i é a fração da massa de probabilidade que o roteador deu a ele. O produto é mínimo quando ambos valem 1/N. f não é diferenciável, P é — é por P que o gradiente flui. α = 10⁻², achado num sweep de 10⁻¹ a 10⁻⁵.
Três truques seguram o treino. Primeiro, precisão seletiva: o corpo da função do roteador roda em float32, e os tensores de dispatch e combine voltam a bfloat16 antes do all-to-all. Treinar tudo em bfloat16 diverge (perplexidade -3,780); a versão seletiva atinge -1,716 na mesma velocidade do bfloat16. Segundo, inicialização com escala 10 vezes menor. Terceiro, no fine-tuning, dropout de 0,4 dentro dos experts e 0,1 no resto — o modelo esparso tem muito mais parâmetros e overfita tarefas pequenas.
O resultado de escala: Switch-Base com 64 experts alcança a qualidade do T5-Base em um sétimo do tempo, com os mesmos FLOPs por token. Switch-Base ainda bate T5-Large, que aplica 3,5 vezes mais FLOPs. E o Switch-C, com 1,6 trilhão de parâmetros e 2.048 experts, chega 4 vezes mais rápido à perplexidade do T5-XXL.
o que isso custou
O ganho por passo não é ganho por segundo. O all-to-all que manda tokens para experts em outros cores é custo puro que o modelo denso não paga, e o paper mede tudo em wall-clock justamente por isso.
Mais grave: a qualidade de pré-treino não desce direito para o fine-tuning nos maiores modelos. O Switch-C, com 1,6 trilhão de parâmetros, tira 87,7 no SQuAD contra 89,6 do Switch-XXL, que tem 4 vezes menos parâmetros e aplica cerca de 10 vezes mais FLOPs por token. Os autores chamam isso de “dependência mal compreendida” entre parâmetros, FLOPs e fine-tuning — e é a admissão mais interessante do trabalho.
A estabilidade continua aberta. As técnicas seguraram Base, Large e o Switch-C, mas o Switch-XXL ainda desestabiliza esporadicamente, e por isso não foi pré-treinado até 1M de passos. Colocar experts nas camadas de atenção melhorou a qualidade e divergiu em bfloat16 — ficou de fora. E o No-Token-Left-Behind, que rerroteia tokens estourados para o segundo melhor expert, não deu benefício empírico nenhum.
onde isso aparece hoje
O top-1 routing com capacity factor e loss auxiliar de balanceamento virou o padrão de fato. Mixtral of Experts e DeepSeek-V3 são descendentes diretos dessa linhagem, mesmo quando voltam a rotear para mais de um expert. A ideia que ficou é a do próprio paper: número de parâmetros é um eixo de escala separado do compute, e não a mesma coisa que ele — uma correção prática às leis de escala que tratavam os dois juntos.
A dívida ancestral está em Outrageously Large Neural Networks, e o baseline contra o qual tudo aqui foi medido é o T5.