O desenvolvimento de Grandes Modelos de Linguagem (LLMs) construídos a partir de modelos de transformador apenas decodificador desempenhou um papel fundamental na transformação do domínio de Processamento de Linguagem Natural (NLP) e no avanço de diversas aplicações de aprendizado profundo, incluindo aprendizado por reforço, análise de séries temporais, processamento de imagens e muito mais. No entanto, apesar de sua escalabilidade e forte desempenho, os LLMs construídos a partir de modelos de transformador apenas decodificador ainda enfrentam limitações significativas. Embora sejam expressivos, o mecanismo de atenção nos LLMs derivados de transformadores exige recursos computacionais intensivos durante a inferência e o treinamento, necessitando de memória substancial para o comprimento da sequência e operações de ponto flutuante (FLOPs) que aumentam quadraticamente. Essa alta exigência computacional limita o comprimento do contexto dos modelos de transformador, torna as tarefas de geração autoregressiva proporcionalmente caras à medida que o modelo escala e impede o aprendizado a partir de fluxos de dados contínuos e a capacidade de processar sequências de comprimento ilimitado de forma eficiente.

Nos últimos tempos, Modelos de Espaço de Estado (MEE) demonstraram capacidades e desempenho notáveis, competindo com modelos de arquitetura de transformador em benchmarks de modelagem em larga escala, enquanto alcançam complexidade de memória como uma função do comprimento da sequência e tempo linear. Além disso, Mamba, um modelo de Espaço de Estado recentemente lançado, demonstrou desempenho excepcional em uma variedade de tarefas de modelagem de linguagem e processamento de sequências longas. Simultaneamente, modelos de Mixture of Expert (MoE) também demonstraram desempenho impressionante, reduzindo significativamente a latência e os custos computacionais da inferência, embora à custa de uma pegada de memória maior. Com base nos modelos Mamba e MoE, este artigo discutirá BlackMamba, uma arquitetura nova que combina o modelo de Espaço de Estado Mamba com modelos MoE para aproveitar os benefícios oferecidos por ambas as estruturas. Experimentos com BlackMamba demonstraram sua capacidade de superar a estrutura existente do Mamba e os modelos de transformador em termos de FLOPs de treinamento e custo de inferência. O desempenho excepcional da estrutura BlackMamba mostra que ela pode combinar as capacidades das estruturas Mamba e MoE de forma eficaz, oferecendo inferência rápida e barata do MoE com geração de complexidade linear do Mamba.

Este artigo visa cobrir a estrutura BlackMamba em profundidade. Exploramos o mecanismo, a metodologia e a arquitetura da estrutura, juntamente com sua comparação com estruturas de geração de imagem e vídeo de última geração. Vamos começar.

BlackMamba: Uma Introdução ao MoE para Modelos de Espaço de Estado

O progresso dos Grandes Modelos de Linguagem (LLMs), particularmente aqueles baseados em arquiteturas de transformador apenas decodificador, influenciou notavelmente o campo de Processamento de Linguagem Natural (NLP) e expandiu-se para diversas aplicações de aprendizado profundo, incluindo aprendizado por reforço, análise de séries temporais, processamento de imagens e além. No entanto, apesar de sua escalabilidade e desempenho robusto, esses LLMs baseados em transformador apenas decodificador enfrentam desafios notáveis. O mecanismo de atenção, uma característica fundamental dos LLMs baseados em transformador, exige recursos computacionais intensivos para inferência e treinamento. Isso envolve a necessidade de memória que cresce com o comprimento da sequência e operações de ponto flutuante (FLOPs) que aumentam quadraticamente. Tais exigências computacionais intensivas restringem o comprimento do contexto dos modelos, elevam os custos das tarefas de geração autoregressiva à medida que o modelo escala e impedem o aprendizado a partir de fluxos de dados contínuos ou processamento de sequências de comprimento ilimitado de forma eficiente.

Esforços significativos foram feitos nos últimos anos para superar essas limitações, e a atenção foi direcionada para desenvolver alternativas arquiteturais aos modelos de transformador canônicos densos com MEEs e modelos MoE sendo as arquiteturas candidatas mais promissoras. O benefício principal obtido ao favor dos Modelos de Espaço de Estado sobre os modelos de arquitetura de transformador é a complexidade computacional linear em relação ao comprimento da sequência de entrada oferecida pelos MEEs, em oposição à complexidade quadrática oferecida pelos transformadores. Teoricamente, a complexidade computacional linear em relação ao comprimento da sequência de entrada permite que os Modelos de Espaço de Estado processem sequências mais longas do que os modelos de arquitetura de transformador para um determinado orçamento de FLOPs ou operações de ponto flutuante por segundo, e tornam a geração autoregressiva constante em computação sem um cache KV. Modelos de Espaço de Estado recentemente desenvolvidos, incluindo Mamba, RetNet e alguns outros, demonstraram inferência e treinamento de sequências longas eficientes, juntamente com desempenho competitivo em tarefas de modelagem de linguagem para transformadores com propriedades de escalabilidade semelhantes. Por outro lado, arquiteturas de modelos de Mixture of Expert estão ganhando popularidade como uma alternativa aos transformadores densos, pois facilitam uma redução significativa nos FLOPs de inferência e treinamento, essenciais para alcançar uma qualidade comparável a um modelo denso. Modelos MoE (Mixture of Experts) operam ativando apenas uma seleção esparsa dos parâmetros totais durante uma única passagem para a frente. Eles utilizam uma função de roteamento para determinar quais ‘especialistas’ são chamados à ação com base no contexto dado. Essa abordagem cria uma separação entre o custo computacional da inferência e o número total de parâmetros, permitindo um desempenho melhorado dentro de um orçamento de inferência fixo, embora com um aumento no número de parâmetros e uma pegada de memória maior.

Essa evolução na arquitetura oferece benefícios notáveis sobre os transformadores tradicionais e representa uma direção emocionante para desenvolvimentos futuros. Acreditamos que integrar essas melhorias em um modelo Mamba-MoE combinado poderia acelerar significativamente as capacidades e eficiência de modelagem de linguagem além do que os modelos de transformador padrão podem oferecer. Os benefícios antecipados de uma arquitetura Mamba-MoE em comparação com um modelo de transformador denso tradicional incluem:

Mamba: alcança complexidade computacional linear em relação ao comprimento da sequência de entrada para ambas as fases de treinamento e inferência. Isso permite que a geração autoregressiva ocorra em um quadro de tempo constante e com uso de memória constante.

MoE: oferece velocidade de inferência e eficiência computacional de treinamento comparáveis a um modelo denso menor, mantendo um nível de qualidade de modelo que rivaliza com o de um modelo com um número equivalente de parâmetros ao modelo denso.

Com isso dito, é essencial afirmar que os modelos de arquitetura de transformador ainda são os melhores, e demonstraram desempenho consistente e notável em tarefas de modelagem de linguagem e processamento de sequências. No núcleo, a arquitetura de transformador emprega autoatenção que realiza uma comparação all-to-all quadrática do produto escalar entre as embeddings de diferentes tokens em uma sequência, e realiza um mapeamento linear para um vetor de saída. O modelo de transformador consiste em blocos de autoatenção empilhados entre blocos de MLP ou Multilayer Perceptron que consistem em um MLP de duas camadas com uma função de ativação dada.

BlackMamba: Arquitetura e Metodologia

Modelos de Espaço de Estado

Modelos de Espaço de Estado pertencem ao grupo de modelos de sequência com complexidade linear em relação ao comprimento da sequência de entrada. A arquitetura dos Modelos de Espaço de Estado se alinha mais com Redes Neurais Recorrentes e Redes Neurais Convolucionais do que com arquiteturas baseadas em atenção, e é inspirada em um sistema dinâmico contínuo que mapeia uma função unidimensional por meio de um espaço latente implícito. Um sistema dinâmico linear torna computações paralelas eficientes usando um scan associativo ou convolucional. Na prática, a natureza recorrente dos Modelos de Espaço de Estado foi a razão pela qual ainda não foi adotada em hardware de IA altamente paralelo, como GPUs. No entanto, o surgimento de MEEs como RWKV e Mamba usaram kernels de scan paralelo para mapear operações recorrentes de forma eficiente para GPUs, facilitando o treinamento de novas arquiteturas com eficiência comparável à dos modelos de transformador.

A complexidade quadrática inerente em relação ao comprimento da sequência dentro dos transformadores é uma limitação bem conhecida que impede o raciocínio e a compreensão sobre contextos muito longos. Inovações recentes introduziram a ideia de estender o comprimento do contexto, permitindo que os transformadores sejam treinados em uma escala viável antes de serem aplicados a contextos muito mais longos durante a inferência. Apesar desses avanços, o processo de inferência ainda exige uma quantidade considerável de recursos computacionais e memória, especialmente para manter o cache Key-Value (KV), tornando-o uma empreitada intensiva em recursos. Esforços de pesquisa recentes se concentraram em melhorar as capacidades expressivas dos modelos de espaço de estado, incorporando mecanismos de gateway dependentes de entrada, semelhantes às matrizes QKV encontradas nos mecanismos de atenção.

Esses esforços visam preservar a progressão linear inerente da recursão do espaço de estado, permitindo execução eficiente por meio de um processo de scan seletivo ou convolucional. Essa abordagem reduz significativamente a disparidade de desempenho com os transformadores em aplicações práticas. Entre esses avanços, Mamba se destaca como um modelo de espaço de estado que reflete os objetivos da pesquisa anterior, mostrando níveis de desempenho impressionantes comparáveis aos dos transformadores em escalas de até 2,8 bilhões de parâmetros. Isso é alcançado aplicando gateway dependente de entrada à recursão do modelo de espaço de estado (MEE), garantindo computação eficiente por meio do uso de kernels de scan seletivo personalizados.

Modelos de Mixture of Expert

Modelos de Mixture of Expert (MoE) alcançam uma separação entre o custo de inferência e o número total de parâmetros, ativando parâmetros seletivamente durante a passagem para a frente. Em vez de usar todos os parâmetros, esses modelos direcionam tokens para especialistas Multilayer Perceptron (MLP) específicos. Idealmente, cada especialista é personalizado para processar um tipo específico de entrada, com um mecanismo de roteamento, essencialmente uma rede neural compacta, determinando o especialista mais adequado para cada token. Essa abordagem visa preservar o poder expressivo abrangente de um modelo com um número equivalente de parâmetros em uma configuração densa, mas com demandas computacionais consideravelmente reduzidas. Tipicamente, o roteador é um mapeamento das camadas lineares de tokens para índices de especialista, com cada especialista sendo um MLP padrão. No entanto, os desenvolvedores ainda precisam descobrir o método de treinamento ideal para o roteador, pois o problema de atribuição de especialista é não diferenciável, e os modelos de Mixture of Expert frequentemente lutam com balanceamento de carga e estabilidade de treinamento entre diferentes especialistas para eficiência de hardware.

Arquitetura

No núcleo, BlackMamba emprega um modelo de transformador padrão consistindo em blocos de MLP e blocos de atenção adicionados em sequência ao longo de um fluxo residual. Agora, a maioria dos modelos de Mixture of Expert simplesmente substitui os blocos de MLP por uma camada de especialista roteada. Por outro lado, a estrutura BlackMamba não apenas substitui o bloco de MLP no transformador por uma camada de especialista roteada, mas também substitui a camada de atenção por uma camada de modelo de espaço de estado Mamba. A arquitetura da estrutura BlackMamba é demonstrada na figura a seguir.

Treinamento e Conjunto de Dados

O modelo BlackMamba é treinado em mais de 300 bilhões de tokens em um conjunto de dados personalizado e usa a função de ativação SwiGLU para os multilayer perceptrons dos especialistas. A estrutura treina com 8 especialistas, um número que os desenvolvedores encontraram ser o equilíbrio correto e a troca entre a pegada de memória e o custo de inferência do modelo. O conjunto de dados personalizado usado para treinar a estrutura BlackMamba consiste em uma mistura de conjuntos de dados de código aberto existentes, incluindo Starcoder, SlimPajama, Pile e mais. A tabela a seguir demonstra os pesos de cada um dos conjuntos de dados usados para treinar a estrutura BlackMamba. No total, há 1,8 trilhão de tokens no conjunto de dados.

BlackMamba: Resultados

Para garantir uma comparação justa entre Mamba e BlackMamba, os desenvolvedores treinaram ambos os modelos com os mesmos parâmetros de treinamento nos mesmos dados de treinamento. A estrutura BlackMamba é capaz de superar ambos os modelos Mamba e os modelos de transformador para o mesmo tamanho de modelo na passagem para a frente, tanto em termos de FLOPs de treinamento quanto de custo de inferência. A figura a seguir demonstra o tempo necessário para gerar uma sequência de determinado comprimento de forma autoregressiva a partir de um prompt inicial de um token como uma função do comprimento da sequência.

Além disso, os benefícios de latência tanto dos modelos de Mixture of Expert quanto dos modelos Mamba são combinados na estrutura BlackMamba, resultando em tempos de inferência significativamente mais rápidos em comparação com os modelos de transformador, os modelos Mamba puros e os modelos MoE. Além disso, a vantagem de inferência da estrutura BlackMamba é diretamente proporcional ao comprimento das sequências, tornando BlackMamba extremamente eficaz na geração de sequências longas. Em seguida, a figura a seguir ilustra o número de tokens atribuídos aos modelos BlackMamba com 340 milhões e 640 milhões de parâmetros, respectivamente. Como pode ser visto, a maioria das camadas demonstra um alto nível de equilíbrio de especialista como resultado do algoritmo Sinkhorn aprimorado implementado pelos modelos BlackMamba.

A tabela a seguir cobre as pontuações de avaliação da estrutura BlackMamba em comparação com uma variedade de modelos de linguagem pré-treinados de código aberto. Como pode ser observado, a estrutura BlackMamba é capaz de competir e superar a maioria das estruturas em todas as linhas de base. Além disso, é digno de nota que os modelos que superam BlackMamba têm um número significativamente maior de parâmetros, e a lacuna no desempenho é mínima, indicando a capacidade da estrutura BlackMamba de alcançar um desempenho comparável com menos parâmetros.

Pensamentos Finais

Neste artigo, discutimos BlackMamba, uma arquitetura nova que combina o modelo de espaço de estado Mamba com modelos MoE para aproveitar os benefícios oferecidos por ambas as estruturas. Experimentos com BlackMamba demonstraram sua capacidade de superar a estrutura existente do Mamba e os modelos de transformador em termos de FLOPs de treinamento e custo de inferência. O desempenho excepcional da estrutura BlackMamba demonstra que ela pode combinar as capacidades das estruturas Mamba e MoE de forma eficaz, oferecendo inferência rápida e barata do MoE com geração de complexidade linear do Mamba. Discutimos como a arquitetura da estrutura BlackMamba é capaz de superar modelos de linguagem treinados fortes, a estrutura existente do Mamba e os modelos de Mixture of Expert em termos de FLOPs de treinamento e custo de inferência. Além disso, a estrutura BlackMamba também herda a geração de FLOPs e o treinamento reduzido de ambos os modelos de Mixture of Expert e Mamba simultaneamente.