Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer; Azalia Mirhoseini; Krzysztof Maziarz; Andy Davis; Quoc Le; Geoffrey Hinton; Jeff Dean

Artigos · arXiv 1701.06538

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean

Publicado em 2017-01-23cs.LGcs.CLcs.NEstat.ML arXiv PDF

Ler com tradução e IA

Resumo em três linhas

🎯 Redes neurais densas sofrem com explosões quadráticas de custo ao expandir a capacidade do modelo para absorver conjuntos massivos de dados.
🔧 Introduz uma camada Sparsely-Gated Mixture-of-Experts (MoE) usando chaveamento Noisy Top-K e perdas auxiliares de balanceamento de carga entre shards distribuídos de GPUs.
📈 Escala modelos de linguagem e tradução até 137 bilhões de parâmetros, superando a melhor perplexidade de LM publicada em 24% a 39% com menor custo computacional.

Resumo

A capacidade de uma rede neural de absorver informações é limitada pelo seu número de parâmetros. A computação condicional, na qual partes da rede ficam ativas com base em cada exemplo, foi proposta na teoria como uma forma de aumentar drasticamente a capacidade do modelo sem um aumento proporcional na computação. Na prática, no entanto, existem desafios algorítmicos e de desempenho significativos. Neste trabalho, abordamos esses desafios e finalmente concretizamos a promessa da computação condicional, alcançando melhorias superiores a 1000x na capacidade do modelo com apenas perdas menores na eficiência computacional em clusters modernos de GPUs. Introduzimos a camada Sparsely-Gated Mixture-of-Experts (MoE), composta por até milhares de sub-redes feed-forward. Uma rede de chaveamento (gating network) treinável determina uma combinação esparsa desses especialistas a ser usada para cada exemplo. Aplicamos a MoE às tarefas de modelagem de linguagem e tradução automática, onde a capacidade do modelo é crítica para absorver as vastas quantidades de conhecimento disponíveis nos corpora de treinamento. Apresentamos arquiteturas de modelos nas quais uma MoE com até 137 bilhões de parâmetros é aplicada convolucionalmente entre camadas LSTM empilhadas. Em grandes benchmarks de modelagem de linguagem e tradução automática, esses modelos alcançam resultados significativamente melhores do que o estado da arte com menor custo computacional.

Resumo original (inglês)

The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation. In practice, however, there are significant algorithmic and performance challenges. In this work, we address these challenges and finally realize the promise of conditional computation, achieving greater than 1000x improvements in model capacity with only minor losses in computational efficiency on modern GPU clusters. We introduce a Sparsely-Gated Mixture-of-Experts layer (MoE), consisting of up to thousands of feed-forward sub-networks. A trainable gating network determines a sparse combination of these experts to use for each example. We apply the MoE to the tasks of language modeling and machine translation, where model capacity is critical for absorbing the vast quantities of knowledge available in the training corpora. We present model architectures in which a MoE with up to 137 billion parameters is applied convolutionally between stacked LSTM layers. On large language modeling and machine translation benchmarks, these models achieve significantly better results than state-of-the-art at lower computational cost.

Guia de leitura

Problema e motivação

Modelos de deep learning se beneficiam enormemente da capacidade de parâmetros quando treinados em grandes conjuntos de dados. No entanto, as arquiteturas padrão ativam toda a rede para cada token ou exemplo de entrada. Isso leva a um aumento inviável no custo computacional à medida que o tamanho da rede e o volume do dataset escalam.

Embora a computação condicional (ativar apenas subconjuntos de um modelo por exemplo) tenha sido proposta teoricamente, as implementações anteriores enfrentaram dificuldades em aceleradores modernos devido a:

GPUs favorecerem aritmética uniforme em detrimento de ramificações (branching).
O "problema do lote decrescente" (shrinking batch problem), em que rotear subconjuntos de exemplos deixa sub-redes individuais com lotes minúsculos e ineficientes.
Gargalos de largura de banda na rede entre dispositivos.
Roteamento desbalanceado, em que algumas poucas sub-redes dominam enquanto outras ficam subutilizadas.

Ideias-chave

Camada MoE com Sparsely-Gating: Sub-redes feed-forward ("especialistas") são posicionadas entre camadas recorrentes (como LSTMs empilhadas). Apenas uma seleção esparsa top- $k$ de especialistas é avaliada por posição.
Noisy Top-K Gating: Ruído gaussiano ajustável é adicionado aos logits de chaveamento antes da seleção top- $k$ e do softmax para fornecer gradientes contínuos e retropropagáveis, além de incentivar uma exploração equilibrada entre os especialistas.
Paralelismo Combinado de Dados e Modelo: As camadas padrão e o chaveamento operam com paralelismo de dados entre os dispositivos, enquanto os especialistas são particionados com paralelismo de modelo entre as GPUs. Os lotes de todos os dispositivos são combinados para que os especialistas recebam mini-lotes agregados suficientemente grandes.
Restrições de Balanceamento: Perdas auxiliares explícitas sobre a importância e a carga dos especialistas impedem que a rede de chaveamento colapse em soluções degeneradas que sobrecarregam apenas um punhado de especialistas.

Método (com as equações principais)

Uma camada MoE consiste em $n$ redes especialistas feed-forward $E_{1}, \dots, E_{n}$ e uma rede de chaveamento $G (x)$ . A saída da camada é: $y = i = 1 \sum n G (x)_{i} E_{i} (x)$ A computação é evitada sempre que $G (x)_{i} = 0$ .

Noisy Top-K Gating

Para combinar esparsidade com chaveamento diferenciável, ruído é injetado nas projeções lineares: $H (x)_{i} = (x \cdot W_{g})_{i} + StandardNormal () \cdot Softplus ((x \cdot W_{noise})_{i})$ $KeepTopK (v, k)_{i} = {v_{i} - \infty se v_{i} estiver entre os top k elementos de v caso contr \overset{a}{ˊ} rio$ $G (x) = Softmax (KeepTopK (H (x), k))$

Perdas de Balanceamento de Carga

Para impor pesos totais de chaveamento iguais, uma perda de importância é definida sobre um lote $X$ usando o coeficiente de variação ( $C V$ ): $Importance (X) = x \in X \sum G (x)$ $L_{importance} (X) = w_{importance} \cdot C V (Importance (X))^{2}$ Para garantir números iguais de exemplos por especialista, o estimador suave de carga usa a CDF normal $Φ$ : $P (x, i) = Φ (\frac{( x \cdot W _{g} ) _{i} - kth_excluding ( H ( x ) , k , i )}{Softplus (( x \cdot W _{noise} ) _{i} )})$ $Load (X)_{i} = x \in X \sum P (x, i)$ $L_{load} (X) = w_{load} \cdot C V (Load (X))^{2}$

Resultados (com os números mais importantes)

1 Billion Word Benchmark: Um modelo MoE-4096-h computacionalmente equivalente alcançou uma perplexidade de teste de 34,1 (uma redução de 24% em comparação com as LSTMs de base computacionalmente equivalentes). Um modelo MoE de alto orçamento atingiu 28,0 de perplexidade de teste em 47 horas em 32 GPUs K40, superando o estado da arte anterior de 34,7 (que levou 59 horas em 32 GPUs K40).
Corpus de 100 Bilhões de Palavras: Escalar para 65.536 especialistas (68 bilhões de parâmetros) reduziu a perplexidade de teste em 39% em comparação com a linha de base densa (28,9 vs. 47,0), mantendo 0,72 TFLOPS/GPU.
Tradução Automática: O GNMT aumentado com MoE alcançou 40,56 BLEU no WMT'14 En $\to$ Fr (+1,34 BLEU sobre o GNMT sem RL) e 26,03 BLEU no WMT'14 En $\to$ De (+1,12 BLEU), exigindo menos da metade das operações por passo temporal.
Tradução Multilíngue: Um único modelo MoE multilíngue superou uma linha de base multilíngue do GNMT em 11 de 12 pares de idiomas (até +5,84 BLEU em Coreano $\to$ Inglês) e reduziu a perplexidade de validação (dev) em 19%.

Limitações e questões em aberto

A esparsidade extrema pode levar a retornos decrescentes ou regressões de desempenho; a perplexidade degradou ao escalar de 65.536 para 131.072 especialistas no dataset de 100B palavras.
O overhead de comunicação distribuída permanece dependente da razão entre dimensões ocultas e de entrada/saída, tornando especialistas de baixa capacidade menos eficientes em largura de banda.
Aplicar MoE de forma recorrente dentro das etapas de recorrência de uma RNN quebra a otimização de desenrolamento de lote (batch-unrolling) através das posições da sequência, exigindo esquemas de recomputação de ativações.

Termos-chave

Sparsely-Gated Mixture-of-Experts Layer (MoE): Uma camada de rede neural contendo múltiplas sub-redes especialistas paralelas, na qual uma rede de chaveamento treinável seleciona apenas um subconjunto esparso de especialistas para cada exemplo de entrada.
Conditional Computation: Um paradigma arquitetural no qual partes da rede neural são dinamicamente ativadas ou desativadas com base em cada exemplo para aumentar a capacidade de parâmetros sem um aumento proporcional na computação.
Noisy Top-K Gating: Um mecanismo de chaveamento que adiciona ruído gaussiano ajustável aos logits antes de reter apenas os k maiores valores e aplicar um softmax, viabilizando a retropropagação e evitando o colapso de roteamento.
Shrinking Batch Problem: A redução no tamanho de lote efetivo recebido por sub-redes individuais quando as entradas são particionadas de forma esparsa entre muitos especialistas, reduzindo a eficiência aritmética do hardware.
Importance Loss: Uma perda auxiliar que penaliza o quadrado do coeficiente de variação da soma dos valores dos gates por especialista ao longo de um lote para incentivar uma alocação equilibrada entre eles.
Load Loss: Um objetivo auxiliar que minimiza a variância na contagem estimada de atribuições ponderadas por probabilidade de tokens entre especialistas para balancear a carga do hardware.
Hierarchical Mixture of Experts: Uma estrutura MoE de dois níveis na qual uma rede de chaveamento primária roteia as entradas para clusters secundários de MoE com o intuito de reduzir o fator de ramificação e o overhead de comunicação distribuída.

Perguntas

Como a rede de chaveamento viabiliza a retropropagação através do roteamento discreto top-k?

Como a seleção top-k é aplicada antes do softmax, os valores não-nulos de chaveamento para os especialistas escolhidos mantêm derivadas contínuas em relação aos pesos de entrada, permitindo que os gradientes fluam de volta através da rede de chaveamento.

Como o problema do lote decrescente (shrinking batch problem) é resolvido entre dispositivos distribuídos?

O modelo combina paralelismo de dados e de modelo: camadas padrão processam mini-lotes independentes entre dispositivos, mas os tokens de todos os dispositivos são agrupados antes de serem roteados para especialistas estacionários hospedados em diferentes GPUs.

Por que o ruído é necessário na rede de chaveamento?

O ruído gaussiano injetado promove a exploração entre os especialistas no início do treinamento e fornece um estimador suave para a perda de balanceamento de carga, evitando o colapso auto-reforçado dos especialistas.

O que acontece quando o modelo é escalado para 131.072 especialistas?

A perplexidade degradou ligeiramente de 28,9 (com 65.536 especialistas) para 29,2 (com 131.072 especialistas), e a eficiência computacional caiu para 0,30 TFLOPS/GPU, provavelmente devido à esparsidade extrema e a tamanhos de lote fixos e não dimensionados.

A explicação desta página foi escrita por IA a partir do artigo e pode conter erros. Consulte o artigo para os detalhes.

Ler com tradução e IA