Resumo Técnico: MobileMoE: Escalonamento de Mistura de Especialistas em Dispositivos Móveis

Declaração do Problema

Embora arquiteturas de Mistura de Especialistas (MoE) se tenham tornado o padrão para escalar modelos de linguagem de grande porte (LLMs) para centenas de bilhões de parâmetros em servidores na nuvem, o seu potencial em escalas sub-bilionárias para implementação em dispositivos permanece amplamente inexplorado. Os LLMs existentes em dispositivos (por exemplo, MobileLLM, MobileLLM-Pro) são esmagadoramente densos, exigindo contagens aumentadas de parâmetros ativos e computação de inferência para melhorar a qualidade. Isso cria um gargalo para a IA de borda, onde os dispositivos enfrentam restrições rigorosas de memória (DRAM) e computação (duração da bateria, latência). Embora as capacidades recentes de DRAM de smartphones tenham aumentado (por exemplo, 12–16 GB), uma metodologia principista para escalar arquiteturas MoE sob essas restrições móveis específicas — equilibrando o tamanho total do modelo, parâmetros ativos e custo de inferência — não foi estabelecida.

Metodologia

1. Lei de Escalonamento de MoE em Dispositivos

Os autores formulam uma lei de escalonamento de MoE em dispositivos generalizada para otimizar conjuntamente a arquitetura sob restrições de memória e computação. A função de perda $L$ é modelada como uma função de parâmetros ativos ($N_{act}$), tamanho dos dados de treinamento ($D$), uma transformação monótona da contagem de especialistas ($\hat{E}$) e escolhas arquiteturais ($x$, incluindo granularidade do especialista $g$ e especialistas compartilhados $s$).

O objetivo de otimização minimiza a perda sujeito a:

• Restrições de Computação: FLOPs de treinamento ($F_{train} = 6N_{act}D$) e FLOPs de inferência por token ($F_{inf} = 2N_{act}$).
• Restrições de Memória: Pegada total de memória $M(N_{total}, T)$, que inclui pesos quantizados ($b_w$-bit) e cache KV ($b_{kv}$-bit) no comprimento de contexto $T$.

2. Design Arquitetural via Divisão e Conquista

Usando a lei de escalonamento, os autores identificam um "ponto ideal" ótimo para MoE em dispositivos através de três estudos de ablação controlados:

• Número de Especialistas ($E$): Esparsidade moderada é ótima. Aumentar $E$ reduz a perda, mas com retornos decrescentes além de $E=8$. $E=8$ é selecionado como o ponto ideal prático para dispositivos com restrição de memória.
• Granularidade do Especialista ($g$): Especialistas de granularidade fina (dividindo cada especialista em $g$ sub-especialistas) permitem caminhos de roteamento mais diversos. O estudo encontra $g=8$ como a granularidade ótima de computação, oferecendo redução significativa de perda sobre especialistas de granularidade grossa sem incorrer em sobrecarga excessiva.
• Especialista Compartilhado ($s$): Incorporar um especialista compartilhado, sempre ativo (um caminho denso), complementa os especialistas roteados. A configuração com um especialista compartilhado ($s=\checkmark$) alcança menor perda do que a contraparte apenas roteada sob computação fixa.

Arquitetura Resultante: A família MobileMoE derivada consiste em três escalas (S, M, L) com parâmetros ativos de 0,3B, 0,5B e 0,9B, e parâmetros totais de 1,3B, 2,8B e 5,3B, respectivamente. O design usa $E=8$, $g=8$ (64 especialistas roteados de granularidade fina) e um especialista compartilhado com roteamento top-4.

3. Receita de Treinamento em Quatro Estágios

Para realizar as vantagens arquitetônicas, um pipeline de treinamento abrangente é empregado em conjuntos de dados de código aberto:

1. Pré-treinamento: ~6T tokens com uma mistura pesada na web para garantir cobertura linguística ampla e especialização de especialistas. Técnicas incluem balanceamento sem perda auxiliar, perda z do roteador e kernels MLP agrupados (GMM) para eficiência.
2. Treinamento Intermediário: ~500B tokens focando em dados de domínio específico de alta qualidade (matemática, código, conhecimento) para estender o contexto para 8k e afiar a especialização de especialistas.
3. Ajuste Fino de Instrução (SFT): ~80M amostras em diversos domínios usando despacho de tokens sem perdas para preservar sinais de instrução estruturados.
4. Treinamento Consciente de Quantização (QAT): Quantização de pesos INT4 (tamanho de grupo 32) com ativações dinâmicas INT8. Crucialmente, os pesos do roteador são mantidos em FP32 para manter a estabilidade do roteamento, adicionando sobrecarga de memória negligenciável.

4. Inferência Eficiente em Dispositivos

Para fechar a lacuna até a implementação, os autores implementam um kernel MoE fundido personalizado no ExecuTorch. Este kernel:

• Converte despacho de especialistas esparsos em GEMMs agrupados densos reordenando tokens.
• Funde todas as sub-operações (roteamento, despacho, projeções, ativação, dispersão) em um único operador para amortizar a sobrecarga de lançamento de kernel e quantização.
• Permite inferência eficiente em CPUs de smartphones comerciais (Snapdragon 8 Elite, Apple A18 Pro) e GPUs.

Resultados Chave

Desempenho em Benchmarks

Em 14 benchmarks fundamentais (senso comum, conhecimento, ciência, leitura, raciocínio), o MobileMoE estabelece uma nova fronteira de Pareto:

• Eficiência vs. Modelos Densos: MobileMoE-S/M igualam ou superam LLMs densos líderes em dispositivos (por exemplo, MobileLLM-Pro, Llama 3.2 1B) usando 2–4× menos FLOPs de inferência em pegadas de memória comparáveis.
• Eficiência vs. Modelos MoE: MobileMoE-M iguala a precisão do estado da arte OLMoE-1B-7B com ~60% menos parâmetros ativos e totais. MobileMoE-L supera o OLMoE-1B-7B com 30% menos parâmetros ativos e uma pegada de memória 23% menor.
• Eficiência de Token: MobileMoE-L alcança desempenho superior com apenas ~6T tokens de pré-treinamento, superando bases densas treinadas em 1,5–2× mais tokens (por exemplo, Llama 3.2 1B em 9T).

Perfil de Tempo de Execução em Dispositivos

Implantado em Samsung Galaxy S25 e iPhone 16 Pro com pesos INT4:

• Latência: MobileMoE-S alcança 1,8–3,8× mais rápido no preenchimento e 2,2–3,4× mais rápido na decodificação comparado à base densa MobileLLM-Pro em memória de pesos INT4 comparável.
• Memória: Em contexto de 8k, MobileMoE-S consome até 22% menos RAM de pico do que a base densa. MobileMoE-L cabe dentro de 5 GB de DRAM em contexto de 8k, confirmando a viabilidade para smartphones modernos.
• Precisão: MobileMoE-L entrega precisão substancialmente maior (Média 57,8 vs. 45,5 para MobileLLM-Pro) com apenas um custo de tempo de execução moderado, que diminui à medida que o comprimento do contexto aumenta.

Significado e Alegações

O artigo alega apresentar a primeira família MoE ativa sub-bilionária otimizada especificamente para implantação em borda, derivada de uma nova lei de escalonamento em dispositivos. Seu significado reside em três contribuições principais:

1. Fundação Teórica: Fornece um framework principista para design de MoE sob restrições conjuntas de memória e computação, identificando esparsidade moderada ($E=8$), especialistas de granularidade fina ($g=8$) e especialistas compartilhados como a configuração ótima para dispositivos móveis.
2. Eficiência de Treinamento: Demonstra que arquiteturas MoE podem alcançar desempenho em dispositivos no estado da arte com significativamente menos tokens de treinamento (~6T) comparado a contrapartes densas, validando a eficiência de aprendizado de MoE em escalas sub-bilionárias.
3. Implantação Prática: Fecha a "última milha" para implantação móvel ao fornecer a primeira inferência MoE eficiente em smartphones comerciais. Os resultados confirmam que MoE não é apenas uma arquitetura de escala de nuvem, mas um caminho prático e eficiente para a próxima geração de IA de borda, habilitando aplicações de baixa latência, preservação de privacidade e custo efetivo em smartphones e wearables.

Os autores concluem que, enquanto os modelos existentes em dispositivos permanecem predominantemente densos, o MobileMoE demonstra que MoE oferece uma alternativa superior para a borda, estabelecendo uma nova fronteira de Pareto que equilibra precisão, memória e latência.