🤖 machine learning

MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training

Este artigo introduz o MESH, um método de otimização Sinkhorn eficiente em memória que incorpora atualizações de momento oculto e pré-condicionamento de bloco opcional para permitir com sucesso o treinamento de modelos Mixture-of-Experts com baixa memória, ao abordar a instabilidade temporal dos gradientes de especialistas roteados que faz com os métodos Sinkhorn sem estado padrão falharem.

Autores originais: Masato Fujitake

Publicado 2026-08-06

📖 4 min de leitura☕ Leitura rápida

CC BY 4.0

Autores originais: Masato Fujitake

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está tentando ensinar um robô gigante e superinteligente a escrever histórias. Para tornar esse robô realmente bom, você precisa de um céreio enorme feito de bilhões de interruptores minúsculos. Mas tem um porém: o cérebro do robô é tão grande que mal cabe na memória do computador. É como tentar rodar um filme de grande orçamento em um celular com uma bateria minúscula; o telefone esquenta, fica lento e pode travar.

Para resolver isso, os cientistas usam "treinadores" especiais chamados otimizadores. O treinador mais famoso, chamado AdamW, é como um treinador muito cuidadoso que mantém um caderno detalhado para cada interruptor no cérebro do robô, lembrando exatamente como dar um empurrãozinho da próxima vez. Isso funciona muito bem, mas esses cadernos ocupam um espaço enorme. Assim, pesquisadores inventaram um treinador mais leve chamado Sinkhorn. Ele é como um treinador que não mantém um caderno; em vez disso, ele apenas olha para o erro atual e faz um ajuste rápido e sem memória. Isso economiza um espaço massivo, mas acontece que, para um tipo específico de cérebro de robô chamado "Mixture-of-Experts" (Mistura de Especialistas ou MoE), esse treinador sem memória se confunde e o robô para de aprender bem. A pergunta que os cientistas estão fazendo é: Podemos manter os benefícios de economia de memória do treinador rápido sem perder a capacidade do robô de aprender?

Este artigo apresenta um novo método chamado MESH (Sinkhorn para Especialistas com Eficiência de Memória) para resolver exatamente esse problema. Os pesquisadores descobriram que o treinador rápido e sem memória falha porque os "especialistas" no cérebro do robô nem sempre chegam a trabalhar. Em um modelo Mixture-of-Experts, o robô escolhe apenas alguns especialistas específicos para lidar com cada frase, kind of like um restaurante onde apenas alguns chefs são chamados para a cozinha para um pedido específico. Como os chefs mudam de um pedido para outro, o treinador rápido vê um sinal confuso e ruidoso e se perde.

O artigo sugere que a solução não é dar ao treinador um caderno completo novamente, mas sim dar a ele uma "memória oculta". Em vez de olhar apenas para o pedido atual, o treinador agora lembra a média dos últimos pedidos antes de tomar uma decisão. O autor chama isso de "momentum oculto". Eles testaram isso em um modelo pequeno de 110 milhões de parâmetros (um "nanowhale") e descobriram que esse truque simples — suavizar o ruído antes de ajustar os pesos — restaura a capacidade de aprendizado do robô.

Aqui está o que eles descobriram e o que eles descartaram:

A Solução Principal: O artigo mostra que a falha acontece porque os "especialistas roteados" (os chefs que só trabalham às vezes) precisam que seus sinais sejam suavizados ao longo do tempo antes que o ajuste sem memória aconteça. Ao tirar a média do histórico recente desses especialistas (usando um buffer oculto que não conta como memória permanente), o novo método MESH funciona muito melhor.
O Que Não Funcionou: Os pesquisadores tentaram muitas outras ideias para corrigir o problema, mas as descartaram. Eles descobriram que simplesmente tornar os ajustes menores, mudar a forma como os especialistas são agrupados ou usar apenas informações de "sinal" (apenas saber se o erro subiu ou desceu, não o quanto) não resolveu o problema. A chave era especificamente o suavizamento temporal, não apenas qualquer dado extra.
O Quão Certos Eles Estão? Os resultados baseiam-se em experimentos com um modelo pequeno específico. Nesses testes, o novo método reduziu a memória necessária para o estado do otimizador em 62,5% (caindo de 0,883 GB para 0,331 GB) e diminuiu o uso de memória de pico do computador em cerca de 12,6%. No entanto, a pontuação final do teste do robô (perda de avaliação) ainda foi ligeiramente pior do que a do treinador pesado que guarda cadernos (AdamW), com pontuações em torno de 3,64 para o novo método comparado a 3,59 para o antigo. O autor sugere que essa pequena diferença pode ser porque outras partes do cérebro do robô (como o vocabulário) ainda precisam do treinador pesado, mas eles estão confiantes de que o "momentum oculto" é a peça que faltava para os especialistas.

Em resumo, o artigo argumenta que você não precisa de um caderno completo para treinar esses cérebros de robôs especiais; você só precisa de uma maneira de lembrar o passado recente sem anotá-lo permanentemente. Essa abordagem de "momentum oculto", chamada MESH, permite que você chegue muito perto do ideal, economizando muita memória enquanto mantém o robô inteligente o suficiente para aprender.

Resumo Técnico: MESH – Otimização Sinkhorn com Eficiência de Memória para Treinamento de Mixture-of-Experts

1. Definição do Problema

O artigo aborda o desafio de otimizar modelos de linguagem Mixture-of-Experts (MoE) com eficiência de memória. Embora o AdamW seja o otimizador padrão para pré-treinamento devido à sua combinação de suavização do primeiro momento temporal e escalonamento adaptativo por coordenada, ele incorre em um alto custo de memória ao armazenar dois tensores de estado completos por parâmetro.

Alternativas mais eficientes em termos de memória, como otimizadores de matriz estilo Sinkhorn ou híbridos SAGE/Sinkhorn, reduzem esse custo removendo os estados de momento para matrizes densas. No entanto, o autor demonstra que aplicar diretamente esses métodos sem estado (stateless) a especialistas (experts) roteados de MoE falha. Em um cenário controlado de "nanowhale" MoE estilo DeepSeek com 110 milhões de parâmetros, um híbrido SAGE/Sinkhorn padrão reduziu significativamente o estado do otimizador (de 0,883 GB para 0,331 GB), mas resultou em uma degradação severa na perda de avaliação (3,8265 vs. a linha de base do AdamW de 3,58–3,64).

A falha central é identificada como um problema de ordenação específico para especialistas roteados:

Gradientes Roteados são Condicionais: Ao contrário dos gradientes de MLP densos, os gradientes de especialistas roteados são amostras condicionais baseadas em decisões de roteamento e pesos de portão (gate weights), que variam temporalmente.
Não Linearidade do Sinkhorn: A normalização Sinkhorn é uma operação não linear. Aplicar o Sinkhorn a um gradiente roteado instantâneo e de alta variância ( $Sinkhorn(G_t)$ ) comporta-se de forma diferente de aplicá-lo a um sinal temporamente suavizado ( $Sinkhorn(E[G_t])$ ).
Amplificação de Ruído: O Sinkhorn sem estado normaliza o ruído induzido pelo roteamento em cada passo, enquanto a falha dos métodos sem estado decorre da falta de suavização temporal antes que a normalização da matriz ocorra.

2. Metodologia: MESH e MESH-B

O autor propõe o MESH (Memory-Efficient Sinkhorn for Experts), um otimizador híbrido projetado especificamente para arquiteturas MoE.

Mecanismo Central: Momento Oculto (Hidden Momentum)

O MESH substitui a atualização Sinkhorn sem estado para especialistas MoE por uma abordagem de momento oculto. Em vez de armazenar um tensor de primeiro momento explícito no dicionário de estado do otimizador (o que anularia o objetivo de economia de memória), o MESH mantém o primeiro momento dentro do ciclo de vida do buffer de gradiente:

Acumulação: Um acumulador de momento de especialista não normalizado é mantido: $H_{e,t} = \beta_1 H_{e,t-1} + G_{e,t}$ .
Ciclo de Vida do Buffer: O acumulador $H_{e,t}$ é armazenado no buffer de gradiente em vez do dicionário de estado do otimizador.
Ciclo de Atualização: Antes do passo de backward, o buffer é reanexado como o gradiente ( $p.grad \leftarrow \beta_1 H_{e,t}$ ). O passo de backward adiciona o gradiente atual, efetivamente atualizando o acumulador para o próximo passo.
Normalização: A atualização Sinkhorn é computada no sinal suavizado: $U_{e,t} = Sinkhorn(H_{e,t})$ .

Esta abordagem restaura o sinal do primeiro momento temporal necessário para especialistas roteados sem armazenar um tensor de estado de otimizador explícito para eles.

Extensão Opcional: MESH-B

O artigo também avalia o MESH-B, que adiciona um multiplicador de inverso de RMS por bloco/neurônio grosseiro à direção Sinkhorn. Isso atua como um pré-condicionador de nível de bloco. As estatísticas de bloco são computadas a partir da fonte suavizada ( $H_{e,t}$ ) em vez de um fluxo de gradiente bruto separado para manter a simplicidade da implementação.

Estratégia de Atribuição de Papéis

O MESH mantém uma estratégia híbrida para diferentes classes de parâmetros:

Embeddings / LM Head: Tratados por SAGE (atualizações de estado leve).
Matrizes Densas Não-Especialistas: Tratadas por Sinkhorn sem estado.
Especialistas Roteados/Compartilhados: Tratados por MESH (Sinkhorn com momento oculto).
Normas / Biases: Tratados por SAGE.

3. Principais Contribuições

Identificação do Ponto de Falha: O artigo identifica as matrizes de especialistas MoE roteados como o ponto de falha específico da otimização SAGE/Sinkhorn direta no pré-treinamento, distinguindo-as de matrizes densas e embeddings.
Algoritmo MESH: Introdução de uma atualização Sinkhorn com momento oculto que restaura um sinal de primeiro momento para os especialistas sem adicionar estado de otimizador explícito, utilizando o ciclo de vida do buffer de gradiente.
Análise de Ablação: Ablações negativas sistemáticas mostrando que intervenções de apenas segundo momento, apenas sinal/momento, apenas confiança e apenas escala não explicam a lacuna de desempenho. O ingrediente causal primário é identificado como suavização temporal antes da normalização da matriz.
Avaliação MESH-B: Avaliação de uma variante pré-condicionada por bloco, mostrando que, embora possa melhorar a fronteira de memória-qualidade em sementes específicas, não é universalmente necessária comparada ao suavimento temporal fornecido pelo M que MESH.

4. Resultados Experimentais

Os experimentos foram conduzidos em um modelo MoE estilo DeepSeek de 110 milhões de parâmetros treinado no FineWeb-Edu por 5.000 passos.

Eficiência de Memória:
- AdamW: 0,883 GB de estado do otimizador.
- MESH / MESH-B: 0,331 GB de estado do otimizador (uma redução de 62,5% em relação ao AdamW).
- Alocação Máxima de CUDA: Reduzida em aproximadamente 12,6% (de ~4,37 GB para ~3,82 GB).
Desempenho (Perda de Avaliação):
- Linha de Base AdamW: ~3,5877 (média entre as sementes 100/512).
- SAGE/Sinkhorn sem estado: 3,8265 (Degradação significativa).
- MESH: 3,6430 (Próximo ao AdamW, mas ligeiramente superior).
- MESH-B: 3,6365 (Ligeiramente melhor que o MESH, mas ainda atrás do AdamW).
Achados Diagnósticos (Semente 42):
- Variantes de estado total (MESH-FULL, MESH-B-FULL) que armazenam explicitamente o primeiro momento recuperam a qualidade do tipo AdamW (ex: MESH-B-FULL com V(H) atingiu 3,6102) com aproximadamente metade do estado do otimizador do AdamW.
- Isso confirma que a suavização temporal é o ingrediente crucial que faltava, enquanto a implementação de momento oculto consegue aproximar isso sem o custo total do estado.

5. Significância e Alegações

O artigo alega modestamente que o MESH fornece um ponto de fronteira memória-qualidade substancialmente menor para o treinamento de MoE.

Conclusão Primária: Especialistas MoE exigem suavização temporal antes da normalização da matriz para funcionarem efetivamente com otimizadores estilo Sinkhorn. Esta suavização é o fator causal dominante para fechar a lacuna de desempenho com o AdamW.
Conclusão Secundária: O pré-condicionamento de bloco/neurônio (MESH-B) pode melhorar os resultados em configurações específicas, mas não é estabelecido como um componente universalmente necessário.
Limitações: O autor reconhece uma lacuna de desempenho residual entre MESH/MESH-B e AdamW. Ele hipotetiza que isso se deve às matrizes de papéis não-especialistas (especificamente matrizes de vocabulário e matrizes densas não-especialistas), que permanecem no SAGE ou Sinkhorn sem estado em sua configuração, enquanto o AdamW as adapta com momentos completos.
Escopo: O estudo é intencionalmente estreito, focando em uma arquitetura de 110M e execuções de 5.000 passos. O autor não afirma que o MESH supera o AdamW em todas as métricas, mas sim que ele recupera a maior parte da qualidade perdida pelos métodos sem estado enquanto reduz drasticamente o overhead de memória.

Em resumo, o MESH demonstra que a falha dos otimizadores eficientes em termos de memória no treinamento de MoE não é inerente à própria estrutura MoE, mas sim à falta de suavização temporal para os gradientes roteados. Ao implementar o "momento oculto", o MESH resolve esse problema de ordenação, oferecendo um caminho viável para o pré-treinamento de MoE com eficiência de memória.