Uma introdução à arquitetura Mamba LLM: um novo paradigma em machine learning

Modelos de linguagem são um tipo de modelo de machine learning treinado para aprender uma distribuição de probabilidade sobre a linguagem natural. Sua arquitetura é formada principalmente por várias camadas de redes neurais, como camadas recorrentes, camadas feedforward, camadas de embedding e camadas de atenção. Essas camadas se combinam para processar um texto de entrada e gerar previsões de saída.

No fim de 2023, pesquisadores da Carnegie Mellon e da Princeton University publicaram um artigo científico apresentando uma nova arquitetura para Large Language Models (LLMs), chamada Mamba. Mamba é uma nova arquitetura de modelos de espaço de estados, voltada para modelagem de sequências. Ela foi desenvolvida para contornar limitações dos modelos transformer, especialmente no processamento de sequências longas, e vem apresentando resultados promissores.

Vamos explorar a arquitetura do Mamba LLM e sua importância em machine learning.

O que é Mamba?

Mamba é uma nova arquitetura de LLM que integra o modelo de sequência Structured State Space (S4) para lidar com sequências extensas de dados. Combinando o melhor de modelos recorrentes, convolucionais e de tempo contínuo, o S4 consegue simular dependências de longo prazo de forma eficaz e eficiente. Isso permite lidar com dados amostrados de forma irregular, ter contexto ilimitado e manter eficiência computacional durante o treino e a inferência.

Aproveitando e expandindo o paradigma S4, o Mamba traz melhorias importantes, sobretudo no tratamento de operações variantes no tempo. Sua arquitetura gira em torno de um mecanismo especial de seleção que ajusta os parâmetros do modelo de espaço de estados (SSM) conforme a entrada.

Como resultado, o Mamba consegue filtrar dados menos relevantes ao focar apenas nas informações cruciais dentro das sequências. Segundo a Wikipedia, "o modelo transita de um arcabouço invariante no tempo para um variante no tempo, o que impacta tanto o cálculo quanto a eficiência do sistema".

Principais recursos e inovações

O afastamento do Mamba em relação aos blocos tradicionais de atenção e MLP é o que o diferencia. Essa simplificação resulta em um modelo mais leve, mais rápido e que escala linearmente com o tamanho da sequência — algo que nenhum de seus predecessores havia alcançado.

Os componentes-chave do Mamba incluem:

• Selective-State-Spaces (SSM): Modelos recorrentes que processam informações de forma seletiva conforme a entrada atual são a base dos SSMs do Mamba. Isso permite filtrar dados supérfluos e focar no que é relevante, o que pode trazer processamento mais eficiente.
• Arquitetura simplificada: O Mamba substitui os blocos complexos de atenção e MLP dos Transformers por um único bloco coeso de SSM. O objetivo é acelerar a inferência e reduzir a complexidade computacional.
• Paralelismo otimizado para hardware: O Mamba pode ter desempenho ainda melhor por usar um modo recorrente com um algoritmo paralelo projetado especificamente para eficiência em hardware.

Outro componente essencial é a Invariância Linear no Tempo (LTI); LTI é uma das características centrais dos modelos S4. Essa propriedade indica que os parâmetros do modelo permanecem constantes em todos os passos de tempo, mantendo a consistência da dinâmica do modelo. A LTI é a base da recorrência e das convoluções, tornando a construção de modelos de sequência mais simples e eficaz.

Detalhes da arquitetura Mamba LLM

A arquitetura do Mamba reforça a relevância dos avanços recentes em machine learning. Ela modifica a forma como os modelos processam sequências ao introduzir uma camada de SSM seletivo. Com isso, o Mamba consegue fazer duas coisas muito importantes:

1. Focar no que importa – o Mamba pode priorizar os dados com maior poder preditivo para a tarefa ao atribuir um peso diferente a cada entrada.

2. Adaptar-se dinamicamente às entradas – como o modelo se ajusta ao input, o Mamba lida com diferentes tarefas de modelagem de sequência com facilidade.

Como resultado, o Mamba lida com sequências com uma eficiência inédita, tornando-se uma escolha ideal para tarefas que envolvem dados longos.

A filosofia de design do Mamba parte de um entendimento das capacidades do hardware moderno. Ele é projetado para aproveitar totalmente o poder de computação das GPUs, garantindo:

• Uso de memória otimizado: Tempos menores de transferência de dados e processamento mais rápido são alcançados ao projetar a expansão de estado do Mamba para caber na HBM (High-Bandwidth Memory) das GPUs.

• Paralelismo maximizado: O Mamba atinge um nível de performance que estabelece um novo patamar para modelos de sequência ao alinhar seus cálculos com a natureza paralela da computação em GPU.

Mamba vs Transformers

A chegada dos Transformers, como o GPT-4, transformou o campo de processamento de linguagem natural (NLP) e estabeleceu referências para várias tarefas de linguagem. Sequências longas, porém, sempre foram um calcanhar de aquiles dos transformers, pois prejudicam significativamente sua eficiência.

É justamente aí que o Mamba se destaca. Ele processa sequências extensas mais rapidamente do que transformers e de forma mais simples graças à sua arquitetura única.

Arquitetura Transformer

Transformers são muito habilidosos em trabalhar com sequências de dados, como textos em modelos de linguagem. Eles processam sequências completas de uma só vez, diferente de modelos anteriores que processavam dados de forma sequencial. Essa característica permite capturar relações complexas dentro dos dados.

Eles utilizam um mecanismo de atenção que permite ao modelo se concentrar em diferentes partes da sequência ao gerar previsões. Três conjuntos de pesos são usados para calcular essa atenção: values, keys e queries, obtidos a partir dos dados de entrada.

Cada elemento de uma sequência é ponderado em relação a todos os outros para indicar quanta importância — ou "atenção" — ele deve ter para prever o próximo elemento da série.

Transformers contam com dois blocos principais: o decoder, que gera a saída, e o encoder, que processa os dados de entrada.

O encoder é composto por várias camadas — cada uma com dois sub-blocos: uma rede feed-forward totalmente conectada e position-wise, e um mecanismo de self-attention multi-head. Para facilitar o treino de redes profundas, conexões residuais e normalização são usadas em cada sub-bloco.

Assim como o encoder, o decoder tem camadas com dois sub-blocos, mas adiciona um terceiro sub-bloco que realiza atenção multi-head sobre a saída do encoder. A propriedade autoregressiva do decoder é preservada por sua natureza sequencial, que limita as previsões de uma posição a considerar apenas posições anteriores.

Portanto, enquanto os Transformers tentam contornar o problema de sequências longas com mecanismos de atenção mais complexos, o Mamba segue um caminho diferente.

Arquitetura Mamba

O Mamba aproveita espaços de estados seletivos. Esse método resolve as ineficiências computacionais dos Transformers com sequências extensas.

A arquitetura do Mamba viabiliza inferência mais rápida e escalonamento linear com o tamanho da sequência, criando um novo paradigma para modelagem de sequências que tende a ser mais eficaz à medida que as sequências crescem.

Como já detalhamos a arquitetura do Mamba acima, não vamos repeti-la aqui.

Aqui está um gráfico da Wikipedia para ajudar a visualizar a comparação entre Mamba e Transformers:

Vale notar que, apesar de muitas vantagens dos SSMs em relação aos Transformers, estes conseguem lidar com sequências muito mais longas do que os SSMs conseguem armazenar em memória, precisam de bem menos dados para aprender tarefas similares e superam os SSMs em tarefas que exigem recuperar ou copiar o contexto de entrada, mesmo com menos parâmetros. 

Primeiros passos com Mamba

Se você quer experimentar o Mamba ou usá-lo em um projeto, vai precisar de:

• Linux
• GPU NVIDIA
• PyTorch 1.12+
• CUDA 11.6+

Para instalar os pacotes necessários do repositório do Mamba, use alguns comandos simples do pip:

• [Opcional] pip install causal-conv1d>=1.2.0: uma implementação eficiente de uma camada Conv1d causal simples usada dentro do bloco Mamba.
• pip install mamba-ssm: o pacote principal do Mamba.

Também é possível compilar a partir do código-fonte com pip install . a partir deste repositório.

Se versões do PyTorch causarem problemas de compatibilidade, o pip pode ser usado com a opção --no-build-isolation para ajudar. Esses modelos foram treinados em grandes conjuntos de dados como o Pile e o SlimPajama e foram construídos para atender a diferentes requisitos de computação e referências de desempenho.

O modelo Mamba possui vários níveis de interface, mas o módulo principal é o bloco de arquitetura Mamba que encapsula o SSM seletivo.

Depois de tudo instalado, você pode usar assim:

# Source: Mamba Repository
import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(
    # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape

Aplicações do Mamba

A introdução do Mamba LLM representa uma potencial virada de chave no universo de arquiteturas de LLM. Mais rápido, eficiente e escalável, o Mamba lida com sequências longas mantendo alto desempenho — por isso deve ter um papel decisivo na evolução de sistemas de IA mais sofisticados.

Em especial, a eficiência e a performance do Mamba podem impulsionar a próxima onda de inovações em IA, abrindo caminho para modelos e aplicações cada vez mais complexos. Seu impacto potencial é enorme, abrangendo processamento de áudio e fala, análise de textos longos, criação de conteúdo, tradução automática em tempo real e muito mais.

Setores que podem ser transformados incluem:

• Saúde: o Mamba pode acelerar o desenvolvimento de tratamentos personalizados ao analisar rapidamente dados genéticos.

• Finanças: o Mamba pode ser usado para analisar tendências de mercado de longo prazo, gerando previsões de ações mais precisas.

• Atendimento ao cliente: o Mamba pode alimentar chatbots que acompanham conversas longas, melhorando a comunicação com clientes.

O caminho à frente para o Mamba

O Mamba se destaca como um farol de inovação, apontando para novas oportunidades no enfrentamento de problemas complexos de modelagem de sequências. Sua estreia é um passo rumo a sistemas mais inteligentes, escaláveis e eficientes, capazes de entender e tratar sequências com uma profundidade inédita.

Embora o Mamba seja um marco técnico relevante, seu sucesso não depende apenas das capacidades tecnológicas — pesquisa colaborativa, recursos compartilhados e contribuições open source também têm papel essencial:

• Contribuições open source: incentivar pesquisadores e desenvolvedores a contribuir com o código do Mamba pode resultar em modelos mais resilientes e adaptáveis.

• Recursos compartilhados: a comunidade pode acelerar o progresso ao reunir conhecimento e recursos e compartilhar modelos pré-treinados.

• Pesquisa colaborativa: parcerias entre instituições acadêmicas e empresas podem expandir as capacidades do Mamba.

Conclusão

O Mamba não oferece apenas um avanço incremental nos modelos de sequência atuais; ele redefine o que é possível. Com sua chegada, a história da inteligência artificial ganha um novo capítulo, em que ineficiências computacionais e limitações de comprimento de sequência começam, enfim, a ficar para trás.

Nos últimos anos, vimos a evolução da IA de RNNs para Transformers e agora para o Mamba — cada passo nos aproxima de uma IA capaz de raciocínio e processamento de informações mais próximos dos humanos. Com sua abordagem de espaço de estados seletivo e escalonamento em tempo linear, o Mamba incorpora o espírito revolucionário que impulsiona o campo de IA.

O Mamba marca o início de uma evolução promissora na inteligência artificial. É um paradigma pensado para o futuro, com potencial para impactar profundamente a IA.

Para continuar aprendendo, confira:

• Como criar aplicações com LLMs usando LangChain
• Como treinar um LLM com PyTorch