Mixture of A Million Experts (MoME): Conceitos-chave explicados

Entendendo o PEER: O poder dos especialistas paralelos

Um dos maiores desafios no dimensionamento de LLMs está nas demandas computacionais e de memória das camadas de alimentação dentro dos blocos de transformadores. O MoE resolve isso substituindo essas camadas por módulos de especialistas ativados de forma esparsa, cada um especializado em diferentes aspectos da tarefa. Essa abordagem aumenta a eficiência ativando apenas os especialistas relevantes para uma determinada entrada, reduzindo a sobrecarga computacional.

As abordagens atuais do MoE têm limitações, como roteadores fixos que precisam ser reajustados quando novos especialistas são adicionados. Assim, uma nova abordagem de roteamento é introduzida, substituindo o roteador fixo por um índice aprendido.

O PEER (Parameter Efficient Expert Retrieval) reduz o número de parâmetros ativos na camada MoE, afetando o consumo de memória de computação e ativação durante o pré-treinamento e a inferência.

O PEER demonstra que, com a aplicação dos mecanismos corretos de recuperação e roteamento, o MoE pode ser dimensionado para milhões de especialistas, reduzindo o custo e a complexidade do treinamento e da utilização de modelos de idiomas muito grandes.

Fonte

No diagrama acima, a consulta de entrada é inicialmente submetida à recuperação da chave do produto para identificar os kprincipais especialistas do site. Em seguida, esses especialistas selecionados processam a entrada com base em seu conhecimento especializado, e seus resultados são combinados durante a fase de inferência para gerar o resultado final do modelo.

Mistura de milhões de especialistas: Os detalhes técnicos

A principal inovação do PEER é o uso da recuperação da chave do produto. O objetivo continua sendo o mesmo do MoE tradicional: encontrar os k melhores especialistas mais adequados para uma determinada tarefa. No entanto, com um grande número de especialistas (potencialmente superior a um milhão), as técnicas anteriores se tornam computacionalmente caras ou ineficientes.

Considere um cenário com N especialistas, cada um representado por um vetord-dimensional. A computação dos principais k especialistas diretamente envolveria o cálculo da similaridade entre a consulta de entrada e todas as N chaves de especialistas, resultando em uma complexidade de tempo de O(Nd). Quando N é muito grande (por exemplo, N ≥ 10^6), isso se torna proibitivamente caro.

O PEER resolve isso empregando uma estratégia inteligente: em vez de usar N chaves especializadas independentes de dimensão d, ele divide cada chave em dois subconjuntos independentes, cada um com dimensionalidade d/2. Da mesma forma, o vetor de consulta é dividido em duas subconsultas. A operação top-k é então aplicada aos produtos internos entre essas subconsultas e subchaves.

Essa estrutura de produto cartesiano das chaves reduz drasticamente a complexidade computacional de O(Nd) para O((N^.5+ k2)d), tornando viável a identificação eficiente dos k melhores especialistas, mesmo com um grande número de especialistas.

Arquitetura da camada PEER

A camada Parameter-Efficient Expert Retrieval (PEER) é uma arquitetura MoE que usa chaves de produto no roteador e MLPs de neurônio único como especialistas.

Uma camada PEER consiste em três componentes:

1. Uma piscina de N especialistas E
2. Um conjunto correspondente de N chaves de produto K
3. Uma rede de consulta q 

Veja como isso funciona: 

1. Dada uma consulta de entrada xrecupera um subconjunto de k especialistas cujas chaves de produto correspondentes têm os produtos internos mais altos com a consulta q(x).
2. Aplique uma função de ativação como sigmoide ou softmax para obter as pontuações do roteador para esses k especialistas recuperados.
3. Calcule o resultado combinando linearmente os resultados dos especialistas ponderados pelas pontuações do roteador.

Essencialmente, a camada PEER identifica de forma eficiente os especialistas mais relevantes para uma determinada entrada, permitindo a utilização efetiva de um grande número de especialistas e mantendo a capacidade de processamento computacional. Essa inovação é um facilitador essencial para o dimensionamento de modelos de MoE para milhões de especialistas, abrindo caminho para LLMs mais poderosos e eficientes.

Vantagens da abordagem PEER

O PEER, em conjunto com a arquitetura do MoME, oferece várias vantagens atraentes em relação às abordagens tradicionais do MoE, ampliando os limites dos recursos do LLM:

• Eficiência aprimorada: Ao utilizar um grande número de especialistas especializados, o PEER permite a recuperação mais precisa e eficiente de informações relevantes, levando a um melhor desempenho geral do modelo.
• Aprendizagem ao longo da vida: O PEER facilita o aprendizado contínuo e o crescimento do modelo. Novos especialistas podem ser adicionados de forma incremental, expandindo a base de conhecimento do modelo sem a necessidade de retreinamento completo, preservando assim o conhecimento adquirido anteriormente.
• Escalabilidade: O PEER aborda as limitações de dimensionamento das arquiteturas MoE tradicionais, permitindo o desenvolvimento de modelos de trilhões de parâmetros, como o GPT-MoE-1.8T. Isso abre novas possibilidades para a criação de LLMs ainda mais potentes e capazes.

Aplicações em potencial do MoME

O Mixture-of-Experts já é um paradigma de modelo amplamente utilizado no setor, com empresas como o YouTube integrando-o ao seu sistema de recomendação. O futuro do MoME parece semelhante e já foi sugerido no GTC 2024 da Nvidia, quando eles falaram sobre o modelo de 1,8 trilhão da GPT.

Fonte

O Mixture of Million Experts (MoME), com sua arquitetura PEER, é particularmente promissor para tarefas complexas de NLP que exigem uma ampla base de conhecimento e recuperação rápida de respostas. Ele aborda os desafios de dimensionamento inerentes ao treinamento e ao fornecimento de modelos de linguagem muito grandes, abrindo novas possibilidades para seu uso em domínios como visão computacional, geração de conteúdo, sistemas de recomendação e computação inteligente.

Desafios e limitações

O dimensionamento de um modelo para até um milhão de especialistas parece promissor do ponto de vista da eficiência, mas o gerenciamento de uma rede tão grande apresenta seus desafios. Vamos explorar alguns deles: 

1. Complexidade e eficiência computacional:
• Custo de recuperação de especialistas: Computar os principais k especialistas de um grupo de um milhão pode ser computacionalmente caro. Mesmo com otimizações como a recuperação de chaves de produtos, o gerenciamento eficiente do processo de recuperação continua sendo um desafio.
• Treinamento em geral: O treinamento de um modelo tão grande com um grande número de especialistas requer recursos computacionais substanciais, o que pode não ser viável para todos os laboratórios de pesquisa ou empresas.
2. Restrições de memória:
• Memória de ativação: À medida que o número de especialistas aumenta, a memória necessária para armazenar ativações e resultados intermediários durante o treinamento e a inferência pode se tornar proibitiva.
• Armazenamento de parâmetros: O armazenamento de um milhão de especialistas exige uma capacidade de memória significativa, o que pode limitar a implementação em dispositivos com recursos limitados. dispositivos com recursos limitados.
3. Regularização e estabilidade:
• Ajuste excessivo: Com um número tão grande de parâmetros, você corre o risco de se ajustar demais aos dados de treinamento. Técnicas de regularização eficazes são essenciais, mas podem ser difíceis de projetar e implementar.
• Estabilidade: A estabilidade do treinamento pode ser um problema, pois o gerenciamento de gradientes e atualizações para um milhão de especialistas pode levar a instabilidades numéricas e problemas de convergência.

Conclusão

Neste artigo, exploramos a técnica Mixture of Million Experts (MoME), uma abordagem dimensionável para grandes modelos de linguagem.

O MoME aproveita as redes de especialistas especializados e o mecanismo de roteamento PEER para melhorar a eficiência e o desempenho. 

Discutimos seus principais componentes, benefícios e possíveis aplicações. Para se aprofundar no assunto, consulte o documento de pesquisa para obter detalhes técnicos e resultados de benchmark.

Para obter mais informações, consulte o artigo de pesquisa para obter detalhes técnicos e resultados de benchmark.