A Meta publicou recentemente uma descoberta inovadora: um novo tipo de tecnologia de camada de memória que melhora significativamente a precisão factual de grandes modelos de linguagem (LLMs) e permite uma expansão sem precedentes na escala de parâmetros. Essa tecnologia não apenas desafia os métodos tradicionais de expansão de redes neurais, mas também oferece novas direções para o design de arquiteturas de IA futuras.
O cerne desta pesquisa reside no uso de um mecanismo de busca chave-valor treinável para adicionar parâmetros extras ao modelo sem aumentar a carga computacional (FLOPs). A ideia central é complementar as camadas feedforward de alta densidade computacional com uma camada de memória de ativação esparsa, fornecendo capacidade especializada de armazenamento e recuperação de informações.
Em comparação com as redes densas tradicionais, a camada de memória é mais eficiente no processamento do armazenamento de informações. Por exemplo, modelos de linguagem precisam aprender informações associativas simples, como nomes e datas de nascimento, ou capitais de países. A camada de memória pode fazer isso através de um mecanismo simples de busca chave-valor, o que é mais eficiente do que usar uma rede feedforward.
A principal contribuição desta pesquisa é a expansão da escala da camada de memória para um nível sem precedentes, atingindo 128 bilhões de parâmetros. Os resultados experimentais mostram que, em tarefas subsequentes, os modelos de linguagem equipados com a camada de memória aprimorada superam não apenas os modelos densos com o dobro da carga computacional, mas também os modelos de especialistas mistos, mesmo com cargas e parâmetros computacionais correspondentes. A melhoria de desempenho é particularmente notável em tarefas factuais.
Os pesquisadores da Meta alcançaram isso substituindo uma ou mais redes feedforward (FFN) na rede Transformer por uma camada de memória. Essa substituição mostrou vantagens consistentes em diferentes tamanhos de modelos base (de 134 milhões a 8 bilhões de parâmetros) e capacidades de memória (até 128 bilhões de parâmetros). Os resultados experimentais demonstram que a camada de memória pode aumentar a precisão factual dos modelos de linguagem em mais de 100%, com melhorias significativas também na codificação e conhecimento geral. Em muitos casos, os modelos com camada de memória atingem o desempenho de modelos densos que exigem quatro vezes mais poder computacional.
Os pesquisadores também fizeram várias melhorias na camada de memória para superar os desafios de sua aplicação em larga escala:
Mecanismo de busca de chave produto: para resolver o gargalo de recuperação de chaves de consulta em camadas de memória em larga escala, a pesquisa utilizou chaves de quantização de produto treináveis, evitando a comparação de cada par de chaves de consulta.
Camada de memória paralela: para permitir a paralelização da camada de memória em ambientes multi-GPU, os pesquisadores distribuíram as operações de busca e agregação de embeddings em várias GPUs.
Mecanismo de memória compartilhada: para maximizar o compartilhamento de parâmetros, os pesquisadores usaram um pool de parâmetros de memória compartilhado entre todas as camadas de memória.
Otimização de desempenho e estabilidade: os pesquisadores otimizaram a operação EmbeddingBag usando kernels CUDA personalizados, melhorando significativamente a utilização da largura de banda de memória. Além disso, foi introduzido um mecanismo de gate dependente da entrada com não-linearidade silu para melhorar o desempenho e a estabilidade do treinamento.
Os resultados experimentais também revelaram as seguintes descobertas-chave:
O tamanho da camada de memória tem um impacto significativo no desempenho: à medida que o tamanho da camada de memória aumenta, o desempenho na resposta a perguntas factuais melhora continuamente.
Múltiplas camadas de memória são melhores do que uma única camada de memória: o uso de várias camadas de memória com parâmetros compartilhados pode melhorar o desempenho, mas muitas camadas de memória podem reduzir o desempenho. O número ideal de camadas de memória é três.
A camada de memória aprende fatos mais rapidamente: no início do treinamento, os modelos equipados com a camada de memória melhoram mais rapidamente, indicando que a camada de memória ajuda o modelo a aprender fatos mais rapidamente.
A camada de memória e a camada densa são complementares: os experimentos mostram que a camada de memória esparsa e a camada feedforward densa são essenciais.
Para verificar a eficácia da tecnologia de camada de memória, os pesquisadores realizaram avaliações em vários benchmarks, incluindo:
Resposta a perguntas factuais (NaturalQuestions, TriviaQA)
Resposta a perguntas de múltiplos saltos (HotpotQA)
Conhecimento científico e de senso comum (MMLU, HellaSwag, OBQA, PIQA)
Codificação (HumanEval, MBPP)
Os resultados mostram que os modelos equipados com a camada de memória superam os modelos de referência em todos esses testes, sendo a melhoria mais significativa na resposta a perguntas factuais.
Esta pesquisa da Meta não apenas fornece novas ideias para a expansão de modelos de IA, mas também abre novos caminhos para resolver problemas factuais e melhorar o desempenho do modelo. Os pesquisadores acreditam que a tecnologia de camada de memória é altamente escalável e que no futuro poderá ser amplamente aplicada em várias aplicações de IA. Eles também apontam que, embora a camada de memória ainda enfrente desafios em termos de aceleração de hardware, acreditam que, por meio de pesquisas e otimizações contínuas, seu desempenho pode se igualar ou até mesmo superar o das redes feedforward tradicionais.
Além disso, a equipe de pesquisa da Meta espera melhorar ainda mais o desempenho da camada de memória por meio de novos métodos de aprendizado, reduzir o esquecimento e as alucinações do modelo e permitir o aprendizado contínuo.
A publicação desta pesquisa certamente injetou nova vitalidade no campo da IA, e nos deixa cheios de expectativa para o futuro desenvolvimento da IA.
Artigo: https://arxiv.org/pdf/2412.09764
