Inteligência artificial
Mamba: Redefinindo a modelagem de sequência e superando a arquitetura de transformadores
Neste artigo sobre Mamba, exploraremos como esse inovador modelo de espaço de estados (SSM) revoluciona a modelagem de sequências. Desenvolvido por Albert Gu e Tri Dao, o Mamba se destaca por sua eficiência no processamento de sequências complexas em áreas como processamento de linguagem, genômica e análise de áudio. Sua modelagem de sequência de tempo linear com espaços de estado seletivos garante desempenho excepcional nessas diversas modalidades.
Iremos nos aprofundar na capacidade do Mamba de superar os desafios computacionais enfrentados pelos Transformers tradicionais, especialmente com sequências longas. Sua abordagem seletiva em modelos de espaço de estados permite inferência mais rápida e escalonamento linear com comprimento de sequência, melhorando significativamente o rendimento.
A singularidade do Mamba reside na sua capacidade de processamento rápido, camada SSM seletiva e design amigável ao hardware inspirado no FlashAttention. Esses recursos permitem que o Mamba supere muitos modelos existentes, incluindo aqueles baseados na abordagem do transformador, tornando-o um avanço notável no aprendizado de máquina.
Transformadores vs. Mamba
Os transformadores, como o GPT-4, estabeleceram padrões de referência no processamento de linguagem natural. No entanto, a sua eficiência diminui com sequências mais longas. É aqui que o Mamba dá um salto à frente, com sua capacidade de processar sequências longas com mais eficiência e sua arquitetura única que simplifica todo o processo.
Transformadores adeptos do tratamento de sequências de dados, como texto para modelos de linguagem. Ao contrário dos modelos anteriores que processavam dados sequencialmente, os Transformers processam sequências inteiras simultaneamente, permitindo-lhes capturar relações complexas dentro dos dados.
Eles usam um mecanismo de atenção, que permite ao modelo focar em diferentes partes da sequência ao fazer previsões.
Essa atenção é calculada usando três conjuntos de pesos: consultas, chaves e valores, derivados dos dados de entrada. Cada elemento de uma sequência é comparado a todos os outros elementos, fornecendo um peso que significa a importância, ou ‘atenção’, que cada elemento deve receber ao prever o próximo elemento da sequência.
Os transformadores mantêm dois blocos principais: o codificador, que processa os dados de entrada, e o decodificador, que gera a saída. O codificador consiste em múltiplas camadas, cada uma contendo duas subcamadas: um mecanismo de autoatenção com múltiplas cabeças e uma rede feed-forward simples, posicionada e totalmente conectada. Normalização e conexões residuais são usadas em cada subcamada para ajudar no treinamento de redes profundas.
O decodificador também possui camadas com duas subcamadas semelhantes ao codificador, mas adiciona uma terceira subcamada que realiza atenção de múltiplas cabeças na saída do codificador. A natureza sequencial do decodificador garante que as previsões para uma posição possam considerar apenas posições anteriores, preservando a propriedade autoregressiva.
Em contraste com os Transformers, o modelo Mamba adota uma abordagem diferente. Enquanto os Transformers lidam com a questão das sequências longas usando mecanismos de atenção mais complexos, o Mamba usa espaços de estado seletivos, fornecendo uma abordagem mais computacional.
Aqui está uma visão geral de alto nível de como funciona um transformador:
- Processamento de entrada: os transformadores primeiro codificam os dados de entrada em um formato que o modelo possa entender, geralmente usando embeddings que também incorporam a posição de cada elemento na sequência.
- Mecanismo de Atenção: em sua essência, o mecanismo de atenção calcula uma pontuação que representa quanto foco deve ser colocado em outras partes da sequência de entrada ao compreender um elemento atual.
- Arquitetura do Codificador-Decodificador: O modelo do transformador é composto por um codificador para processar a entrada e um decodificador para gerar a saída. Cada uma consiste em múltiplas camadas que refinam a compreensão do modelo sobre a entrada.
- Atenção Multi-Cabeça: Tanto no codificador quanto no decodificador, a atenção multicabeças permite que o modelo atenda simultaneamente a diferentes partes da sequência a partir de diferentes espaços representacionais, melhorando sua capacidade de aprender em diversos contextos.
- Redes Feed-Forward posicionadas: Após atenção, uma rede neural simples processa a saída de cada posição separadamente e de forma idêntica. Isto é combinado com a entrada através de uma conexão residual e seguido pela normalização da camada.
- Geração de resultados: o decodificador então prevê uma sequência de saída, influenciada pelo contexto do codificador e pelo que ele gerou até o momento.
A capacidade do transformador de lidar com sequências em paralelo e seu robusto mecanismo de atenção o tornam poderoso para tarefas como tradução e geração de texto.
Em contraste, o modelo Mamba opera de forma diferente ao usar espaços de estados seletivos para processar sequências. Esta abordagem aborda a ineficiência computacional em Transformers ao lidar com sequências longas. O design do Mamba permite inferência mais rápida e escala linearmente com o comprimento da sequência, estabelecendo um novo paradigma para modelagem de sequência que poderia ser mais eficiente, especialmente à medida que as sequências se tornam cada vez mais longas.
Mamba
O que torna o Mamba verdadeiramente único é o seu afastamento da atenção tradicional e dos blocos MLP. Essa simplificação leva a um modelo mais leve e rápido, que é dimensionado linearmente com o comprimento da sequência – um feito incomparável aos seus antecessores.
Os principais recursos do Mamba incluem:
- SSMs seletivos: Permitem ao Mamba filtrar informações irrelevantes e focar em dados relevantes, melhorando o tratamento de sequências. Essa seletividade é crucial para um raciocínio eficiente baseado em conteúdo.
- Algoritmo com reconhecimento de hardware: o Mamba usa um algoritmo paralelo otimizado para hardware moderno, especialmente GPUs. Este design permite uma computação mais rápida e reduz os requisitos de memória em comparação com os modelos tradicionais.
- Arquitetura Simplificada: Ao integrar SSMs seletivos e eliminar blocos de atenção e MLP, o Mamba oferece uma estrutura mais simples e homogênea. Isso leva a melhor escalabilidade e desempenho.
O Mamba demonstrou desempenho superior em vários domínios, incluindo linguagem, áudio e genômica, destacando-se tanto em tarefas de pré-treinamento quanto em tarefas específicas de domínio. Por exemplo, na modelagem de linguagem, o Mamba iguala ou excede o desempenho de modelos maiores do Transformer.
O código do Mamba e os modelos pré-treinados estão disponíveis abertamente para uso da comunidade em GitHub.
As tarefas de cópia padrão são simples para modelos lineares. As cabeças seletivas de cópia e indução requerem memória dinâmica e com reconhecimento de conteúdo para LLMs.
Os modelos de espaço de estados estruturados (S4) emergiram recentemente como uma classe promissora de modelos de sequência, abrangendo características de RNNs, CNNs e modelos clássicos de espaço de estados. Os modelos S4 inspiram-se em sistemas contínuos, especificamente um tipo de sistema que mapeia funções ou sequências unidimensionais através de um estado latente implícito. No contexto da aprendizagem profunda, representam uma inovação significativa, proporcionando uma nova metodologia para projetar modelos de sequência que sejam eficientes e altamente adaptáveis.
A dinâmica dos modelos S4
SSM (S4) Este é o modelo básico de espaço de estados estruturado. É preciso uma sequência x e produz uma saída y usando parâmetros aprendidos A, B, C, e um parâmetro de atraso Δ. A transformação envolve discretizar os parâmetros (transformar funções contínuas em funções discretas) e aplicar a operação SSM, que é invariante no tempo – o que significa que não muda em diferentes intervalos de tempo.
O significado da discretização
A discretização é um processo chave que transforma os parâmetros contínuos em discretos através de fórmulas fixas, permitindo que os modelos S4 mantenham uma conexão com sistemas de tempo contínuo. Isso confere aos modelos propriedades adicionais, como invariância de resolução, e garante a normalização adequada, melhorando a estabilidade e o desempenho do modelo. A discretização também traça paralelos com os mecanismos de controle encontrados nas RNNs, que são essenciais para gerenciar o fluxo de informações através da rede.
Invariância Linear de Tempo (LTI)
Uma característica central dos modelos S4 é sua invariância linear no tempo. Esta propriedade implica que a dinâmica do modelo permanece consistente ao longo do tempo, com os parâmetros fixos para todos os intervalos de tempo. LTI é a base da recorrência e convoluções, oferecendo uma estrutura simplificada, porém poderosa, para a construção de modelos de sequência.
Superando Limitações Fundamentais
A estrutura S4 tem sido tradicionalmente limitada pela sua natureza LTI, o que coloca desafios na modelação de dados que requerem dinâmica adaptativa. O recente artigo de pesquisa apresenta uma abordagem que supera essas limitações, introduzindo parâmetros que variam no tempo, eliminando assim a restrição do LTI. Isto permite que os modelos S4 lidem com um conjunto mais diversificado de sequências e tarefas, expandindo significativamente a sua aplicabilidade.
O termo “modelo de espaço de estados” abrange amplamente qualquer processo recorrente que envolva um estado latente e tem sido usado para descrever vários conceitos em múltiplas disciplinas. No contexto da aprendizagem profunda, os modelos S4, ou SSMs estruturados, referem-se a uma classe específica de modelos que foram otimizados para computação eficiente, mantendo a capacidade de modelar sequências complexas.
Os modelos S4 podem ser integrados em arquiteturas de redes neurais ponta a ponta, funcionando como transformações de sequência autônomas. Eles podem ser vistos como análogos às camadas de convolução nas CNNs, fornecendo a espinha dorsal para a modelagem de sequências em uma variedade de arquiteturas de redes neurais.
SSM vs SSM + Seleção
Motivação para Seletividade na Modelagem de Sequência
SSMs estruturados
O artigo argumenta que um aspecto fundamental da modelagem de sequência é a compressão do contexto em um estado gerenciável. Modelos que podem focar seletivamente ou filtrar entradas fornecem um meio mais eficaz de manter esse estado compactado, levando a modelos de sequência mais eficientes e poderosos. Essa seletividade é vital para que os modelos controlem de forma adaptativa como as informações fluem ao longo da dimensão da sequência, uma capacidade essencial para lidar com tarefas complexas na modelagem de linguagem e além.
Os SSMs seletivos aprimoram os SSMs convencionais, permitindo que seus parâmetros sejam dependentes de entrada, o que introduz um grau de adaptabilidade anteriormente inatingível com modelos invariantes no tempo. Isto resulta em SSMs variantes no tempo que não podem mais usar convoluções para computação eficiente, mas, em vez disso, dependem de um mecanismo de recorrência linear, um desvio significativo dos modelos tradicionais.
SSM + Seleção (S6) Esta variante inclui um mecanismo de seleção, adicionando dependência de entrada aos parâmetros B e C, e um parâmetro de atraso Δ. Isso permite que o modelo se concentre seletivamente em certas partes da sequência de entrada x. Os parâmetros são discretizados levando em consideração a seleção, e a operação SSM é aplicada de forma variável no tempo por meio de uma operação de varredura, que processa os elementos sequencialmente, ajustando o foco dinamicamente ao longo do tempo.
Destaques de desempenho do Mamba
Mamba é o melhor da categoria em todos os resultados de avaliação
Em termos de desempenho, o Mamba se destaca tanto na velocidade quanto na precisão de inferência. Seu design permite uma melhor utilização de contextos mais longos, o que é demonstrado tanto na modelagem de DNA quanto de áudio, superando os modelos anteriores em tarefas complexas que exigem dependências de longo alcance. Sua versatilidade também é destacada em avaliações zero-shot em múltiplas tarefas, estabelecendo um novo padrão para tais modelos em termos de eficiência e escalabilidade.
Introdução ao Mamba
Para aqueles interessados em aproveitar o Mamba, os requisitos técnicos incluem um sistema operacional Linux, uma GPU NVIDIA, PyTorch 1.12+ e CUDA 11.6+. A instalação envolve comandos pip simples para instalar os pacotes necessários do Repositório Mamba. Se surgirem problemas de compatibilidade com as versões do PyTorch, usar o sinalizador –no-build-isolation com pip pode ajudar. Esses modelos, treinados em extensos conjuntos de dados como o conjunto de dados Pile e SlimPajama, são projetados para atender a diversas necessidades computacionais e benchmarks de desempenho.
O Mamba oferece diferentes níveis de interfaces, desde a camada SSM seletiva até o bloco Mamba e estruturas completas de modelos de linguagem. O bloco Mamba, que é o módulo principal da arquitetura, utiliza uma camada causal Conv1d e pode ser facilmente integrado em projetos de redes neurais. O exemplo de uso fornecido em Python demonstra a instanciação de um modelo Mamba e o processamento de dados por meio dele, destacando a simplicidade e flexibilidade do sistema.
Modelos Mamba pré-treinados estão disponíveis em Abraçando o rosto, com tamanhos variando de parâmetros de 130M a 2.8B, treinados no extenso conjunto de dados Pile e no conjunto de dados SlimPajama. Esses modelos são projetados para atender diversos requisitos computacionais e de desempenho, aderindo aos padrões dimensionais do GPT-3. Os usuários podem esperar alto rendimento e precisão desses modelos, tornando o Mamba uma escolha competitiva para diversas aplicações, incluindo, entre outras, modelagem de linguagem.
O impacto do Mamba
O Mamba representa um avanço na modelagem sequencial, oferecendo uma alternativa poderosa às arquiteturas Transformer para processamento de dados com muita informação. Seu design se alinha às demandas do hardware moderno, otimizando o uso de memória e os recursos de processamento paralelo. A disponibilidade de código aberto da base de código do Mamba e seus modelos pré-treinados o tornam uma ferramenta acessível e robusta para pesquisadores e desenvolvedores na área de IA e aprendizado profundo.
