Tying the Loop -- Tied Expert Layers in Mixture-of-Experts Language Models

Resumo Técnico: Tying the Loop - Camadas de Especialistas Vinculadas em Modelos de Linguagem Mixture-of-Experts

Declaração do Problema

As arquiteturas Mixture-of-Experts (MoE) tornaram-se o padrão para escalar Modelos de Linguagem de Grande Escala (LLMs) ao desacoplar a contagem total de parâmetros do computação por token. Ao ativar apenas um pequeno subconjunto de redes feed-forward (FFNs) de especialistas por token, os modelos MoE alcançam alta capacidade com inferência eficiente. No entanto, um gargalo crítico permanece: a pegada de memória. Embora apenas uma fração dos parâmetros seja ativa durante qualquer passagem de forward, o conjunto total de parâmetros (dominado pelos pesos dos especialistas) deve residir na memória de treinamento e de inferência.

Isso cria uma tensão com a tendência emergente de "eficiência de parâmetros" em modelos de raciocínio e de profundidade em loop, que visam extrair mais computação de cada parâmetro único. As arquiteturas MoE padrão são vistas como ineficientes nesse aspecto porque inflam a memória com parâmetros que permanecem ociosos em qualquer passagem dada. O artigo busca reconciliar a baixa computação por token do MoE com o objetivo de uma maior eficiência de computação por parâmetro único sem sacrificar a qualidade do modelo.

Metodologia: Expert Tying (Vinculação de Especialistas)

Os autores propõem o Expert Tying, uma modificação arquitetônica que compartilha parâmetros de FFN de especialistas entre um grupo de camadas consecutivas do transformer, preservando componentes específicos de cada camada.

Mecanismo Central

• Compartilhamento de Parâmetros: Para um grupo de $g$ camadas consecutivas, os pesos de projeção do gate (porta), up e down dos $N$ especialistas são aliados (compartilhados). Isso reduz os parâmetros únicos de FFN por um fator de $g$.
• Independência Preservada: Crucialmente, o roteador (rede de roteamento), as projeções de auto-atenção (Q, K, V, O) e as camadas de normalização permanecem independentes para cada camada.
• Implementação: O método é implementado como uma simples atribuição de ponteiro no código do modelo (ex: Hugging Face Transformers). Não requer mudanças na infraestrutura de treinamento ou inferência, exceto para garantir que o otimizador acumule corretamente os gradientes das múltiplas frentes de uso dos parâmetros vinculados.
• Escalonamento da Taxa de Aprendizado: Para gerenciar o acúmulo de gradiente de $g$ camadas, a taxa de aprendizado para os pesos de especialistas vinculados é escalada por $1/\sqrt{g}$.

Racional de Design

Os autores hipotetizam que o pool de especialistas de FFN pode ser compartilhado entre camadas próximas porque a atenção por camada garante que o operador efetivo de cada camada permaneça distinto. Mesmo com pesos de especialistas idênticos, atender sobre diferentes misturas de posições de tokens produz transformações diferentes em cada profundidade.

Principais Contribuições

1. Ablação Controlada de Componentes

Antes de escalar para modelos de produção, os autores realizaram um estudo de ablação controlado em uma arquitetura menor para determinar quais componentes podem ser vinculados sem degradar a qualidade.

• Descoberta: A auto-atenção por camada é o principal motor da distinção das camadas. Desvincular a atenção (mantendo-a independente) melhora a perda significativamente mais do que desvincular o roteador.
• Impacto do Roteador: Embora desvincular o roteador aumente a diversidade de roteamento (o acordo cross-loop cai), isso tem um impacto negligenciável na perda final se a atenção for vinculada.
• Conclusão: A estratégia ideal é vincular os pesos dos especialistas, manter a atenção por camada e tratar o roteador como vinculado ou desvinculado com diferença de custo mínima.
• Prelúdio/Coda: As primeiras e últimas camadas da pilha (entrada de embedding e saída do LM-head) devem permanecer desvinculadas. Um prelúdio e uma coda "2+2" desvinculados oferece o maior ganho arquitetônico em comparação com designs totalmente recorrentes.

2. Escalonamento para Arquiteturas de Produção

O método foi aplicado a três arquiteturas MoE de última geração: OLMoE, Qwen3-MoE e DeepSeekMoE.

• Configuração: Grupos de $g=4$ camadas compartilham pesos de especialistas, com 2 camadas desvinculadas no início e no fim.
• Redução de Parâmetros: Esta configuração reduz a contagem total de parâmetros em 29% a 52% (dependendo da arquitetura e do tamanho do grupo) com degradação negligível na perplexidade ou precisão downstream.
• Throughput: A redução de parâmetros únicos leva a uma aceleração de 23,7% no tempo de execução (wall-clock) do treinamento devido à redução da largura de banda de carregamento de pesos e atualizações de estado do otimizador.

3. Expansão de Largura Heterogênea

O artigo introduz um uso inovador dos parâmetros economizados: reinvestir esses parâmetros nas camadas médias vinculadas.

• Em vez de simplesmente comprimir o modelo, os autores alargam as camadas médias vinculadas (aumentando o número de especialistas) para corresponder à contagem total de parâmetros original ("iso-parameter").
• Resultado: Modelos com camadas médias vinculadas e largura expandida superam consistentemente a linha de base não vinculada (untied) no mesmo número de parâmetros. Isso redefine o expert tying não apenas como uma técnica de compressão, mas como um eixo de design de profundidade-vs-largura.

4. Validação em Larga Escala

Os experimentos foram estendidos a um modelo de 7 bilhões de parâmetros (OLMoE-1B-7B).

• O modelo vinculado com $g=4$ alcançou perda e precisão downstream equivalentes à linha de base não vinculada, enquanto usou menos da metade dos parâmetros totais (3.50B vs 7.12B).
• A variante de largura expandida melhorou ainda mais o desempenho, superando o checkpoint oficial do OLMoE-1B-7B treinado com o dobro de tokens, apesar de ter metade dos parâmetros.

Resumo dos Resultados

• Eficiência de Memória: O expert tying reduz a pegada de memória em quase 2× (até 52% de redução de parâmetros) com quase nenhuma degradação na perplexidade.
• Preservação de Qualidade: Com $g=4$, a penalidade de perda é mínima (ex: +0,05 a +0,07 na perda de validação), e a precisão downstream cai no máximo 1,9%.
• Ganhos de Performance: O throughput de treinamento aumenta de 15,7% a 23,7% devido à redução de movimentação de dados e overhead do otimizador.
• Superioridade Iso-Parameter: Quando os parâmetros economizados são reinvestidos para alargar as camadas vinculadas, os modelos resultantes superam os MoEs não vinculados padrão do mesmo tamanho.

Significância e Alegações

O artigo afirma que o Expert Tying reconcilia com sucesso as tendências opostas de escalonamento de MoE (baixa computação por token) e eficiência de parâmetros (alta computação por parâmetro único). Ao explorar a redundância inerente de parâmetros nos caminhos de MoE, o método remove a penalidade de memória sem remover a própria esparsidade.

Os autores enfatizam que sua contribuição não é meramente a configuração (que compartilha semelhanças com trabalhos anteriores como Megrez2), mas o estudo controlado que identifica quais componentes podem ser compartilhados. Eles demonstram que a atenção, e não o roteamento, é o componente crítico para manter a distinção das camadas, permitindo que o maior pool de parâmetros (especialistas) seja compartilhado com o menor custo de qualidade.

O trabalho posiciona o expert tying como uma modificação "drop-in" para bases de código de MoE existentes que avança o treinamento e o escalonamento eficiente de LLMs de próxima geração, ofereando um trade-off altamente favorável entre computação e memória. Os autores mantêm-se modestos quanto às limitações, observando que os experimentos estão atualmente limitados a 7 bilhões de parâmetros e que os ganhos de eficiência são limites inferiores devido à falta de kernels especializados para camadas vinculadas.