Renovando redes de memória de longo e curto prazo: XLSTM para IA de última geração

LSTMs, ou Long Short-Term Memory Networks, já existem há muito tempo. Eles têm sido aplicados para diversas tarefas relacionadas a sequências, como geração e tradução de texto ou até mesmo geração de legendas de imagens.

A desvantagem é que eles não podiam ser paralelizados para fazer uso das poderosas GPUs modernas. Essa limitação abriu caminho para o surgimento de transformadores que utilizam GPUs para paralelização massiva de treinamento e inferência.

Se tentarmos agora renovar e paralelizar os LSTMs, poderão eles tornar-se ferramentas para a construção de LLMs de próxima geração?

Esta é exatamente a pergunta respondida pelo artigo “ XLSM – Extended Long Short-term Memory Networks”, que significa memória “ extendida ” de longo prazo. Eles fazem isso propondo dois novos blocos na arquitetura, a saber, sLSTM e mLSTM.

Então, vamos nos aprofundar nos blocos sLSTM e mLSTM propostos neste artigo e ver como podemos empilhá-los para desenvolver a arquitetura XLSTM.

Explicação visual

Se você é alguém como eu e gostaria que os XLSTMs fossem explicados visualmente, verifique o vídeo do YouTube neste artigo:

Atualização LSTM

Uma das primeiras redes projetadas para lidar com dados sequenciais é a Rede Neural Recorrente.

Rede Neural Recorrente

Ele usa conexões recorrentes em sua arquitetura com x como entrada e o como saída. Se desdobrarmos, podemos visualizá-lo como uma sequência de operações acontecendo nos carimbos de data/hora, t-1, t e t+1 . Uma grande desvantagem dos RNNs era o problema do gradiente evanescente, em que o gradiente chega a zero à medida que empilhamos muitos blocos.

LSTMs, ou redes de memória de longo prazo, foram propostas para superar os gradientes de desaparecimento, introduzindo estados de células e mecanismos de controle na rede.

Uma ilustração simplificada de LSTMs

Os estados da célula, c , são memórias de longo prazo que vivem em vários carimbos de data/hora. Os estados ocultos, h , são memórias de curto prazo que são transmitidas de um intervalo de tempo para outro. E, claro, temos as entradas z da sequência de entradas.

Três portões têm funções em forma de S. A porta do esquecimento usa uma função sigmóide para decidir quais informações esquecer na memória de longo prazo. A porta de entrada também usa uma função sigmóide para processar a entrada e adicioná-la à saída da porta de esquecimento. Esta operação de adição tem um termo sofisticado chamado carrossel de erros constantes no artigo XLSTM e na literatura acadêmica. Esta operação de adição é o que resolve o problema de gradientes de fuga encontrado em RNNs. A saída c_t é então processada pela porta de saída , que geralmente é uma função tanh que leva à saída de estado oculto h_t que é passada para a próxima etapa.

Com essas operações, dissecamos as duas equações principais dos LSTMs que são c_t e h_t .

Desvantagem 1 – Revisão das decisões de armazenamento

Uma das principais desvantagens dos LSTMs é a incapacidade de revisar as decisões de armazenamento. O que isto significa é que à medida que o comprimento da sequência aumenta, o modelo deve ser capaz de decidir se mantém ou não as informações passadas na memória.

Por exemplo, se pegarmos esta frase: “Tom foi à loja. Ele comprou algumas bebidas”, e compare com “Tom foi à loja comprar alguns mantimentos que incluíam cenouras, cebolas, bananas, maçãs, laranjas, café e pão. Ele também comprou algumas bebidas. Para cada palavra nova, como bananas ou maçãs, o modelo tem que revisar constantemente se deve manter a palavra anterior “Tom” em sua memória. Este é um grande desafio para os LSTMs e decorre da função sigmóide de seu portão de esquecimento.

A função sigmóide versus função exponencial. O sigmóide se achata nas extremidades, mas o exponencial só continua aumentando.

Portanto, se considerarmos uma porta de esquecimento, ela é composta por uma função sigmóide, que possui uma curva em forma de S que se achata no final. Isto indica que à medida que avançamos em direção a valores mais elevados de entrada, a decisão sobre o que esquecer e o que manter na memória torna-se bastante desafiadora. Mas se usarmos uma função exponencial em seu lugar, o jogo muda e, à medida que avançamos para valores mais altos de entrada, obtemos um intervalo mais amplo para as saídas. Isso, por sua vez, indica que os LSTMs podem melhorar na revisão das decisões de armazenamento.

Solução 1 – sLSTM

Assim, a solução proposta neste artigo são os blocos do sistema. Se voltarmos à equação LSTM clássica que representa o estado da célula, como vimos antes, é uma função da porta de esquecimento e das portas de entrada.

Essas portas, por sua vez, são compostas por funções sigmóides. Então, e se substituirmos essas funções sigmóides por funções exponenciais? As novas portas f_t e i_t agora se tornam exp(f_t) e exp(i_t), e essa é praticamente a principal modificação para criar o bloco sLSTM.

Ao contrário da função sigmóide, que comprime as entradas para ficarem em um intervalo fixo, as funções exponenciais tendem a aumentar de valor à medida que a entrada aumenta e não normaliza naturalmente a saída para ficar entre, digamos, 0 e 1 como a função sigmóide. função.

Portanto, precisamos introduzir um novo estado normalizador , que é uma função das portas de esquecimento e de entrada. Podemos pensar nisso como uma média móvel de um valor de normalização.

Usamos os valores de normalização calculados para normalizar a saída ou o novo estado oculto.

Embora a normalização cuide dos estados ocultos, para evitar que o exponencial exploda as portas de esquecimento e entrada, precisamos introduzir um estabilizador . Ele vem na forma de funções logarítmicas para neutralizar o efeito das exponenciais e introduzir estabilidade. Portanto, o estado do estabilizador é o máximo do log da porta de esquecimento e da saída da porta de entrada. Subtraímos esses valores do estabilizador da entrada e esquecemos as portas para estabilizá-los.

Desvantagem 2 – Memória e Paralelização

A segunda desvantagem dos LSTMs é a falta de paralelização. Os LSTMs foram projetados para lidar com dados sequenciais, o que significa que precisam da saída do processamento da entrada anterior na sequência para processar a entrada atual na sequência. Essa desvantagem específica impede a paralelização e foi a culpada que levou ao início da era dos Transformers.

A solução proposta neste artigo são os novos blocos mLSTM. Então, vamos dar uma olhada neles a seguir.

Solução - mLSTM

O próximo bloco de construção dos XLSTMs é o bloco mLSTM, onde m significa memória. Vamos voltar à equação LSTM clássica novamente para ver qual é a desvantagem dela. Podemos ver que o estado da célula c_t é um escalar. Isso significa que lidamos apenas com 1 número por vez, quando temos o luxo das GPUs modernas com pelo menos 12 GB de memória.

O bloco mLSTM introduz matrizes no lugar de escalares para os estados das células. Voltando à nossa equação clássica de LSTM, e se substituirmos c_t por uma matriz C *_t* para que o estado da célula agora se torne C *_t* maiúsculo para indicar matrizes e os estados da célula possam ser recuperados não apenas por uma porta i_t, mas armazenando pares de valores-chave que são vetores. cujos valores podem ser recuperados por consultas que são vetores da mesma dimensão.

Para parecer familiar à terminologia do transformador, eles introduziram chave e valor aqui para formar esta matriz.

XLSTM

Com essas informações sobre sLSTM e mLSTM, vamos mergulhar na arquitetura detalhada do XLSTM.

sLSTM

Uma visão detalhada do bloco sLSTM

Quando se trata de sLSTM, usamos projeções post-up. Portanto, a entrada passa primeiro por camadas de convolução causal com uma função de ativação swish. A saída dessas camadas é então alimentada através de uma camada linear diagonal de bloco com quatro blocos diagonais ou “cabeças”. A saída destes é então alimentada através do bloco sLSTM com quatro cabeçotes. Finalmente, a saída é projetada para cima usando uma camada MLP fechada com ativação GeLU e projetada para baixo usando uma função MLP fechada.

mLSTM

Passando para os detalhes do bloco mLSTM, usamos projeções pré-up. O que significa que a entrada é primeiro projetada com um fator de projeção de 2. Uma das saídas de projeção vai para o mLSTM e a outra vai para a porta de saída. A entrada para o bloco mLSTM passa por convolução causal e, em seguida, por matrizes de projeção diagonal de bloco de tamanho de bloco 4, que geram a consulta, a chave e o valor que são prontamente usados pelo bloco mLSTM.

Arquitetura XLSTM

Finalmente, podemos empilhar os dois tipos de blocos para formar a arquitetura LSTM estendida. Portanto, os blocos cinza escuro são os blocos mLSTM e os cinza claro são os blocos sLSTM.

Em termos de vantagens, o artigo menciona que as redes XLSTM possuem uma complexidade computacional linear e uma complexidade de memória constante em relação ao comprimento da sequência.

Avaliações

Os autores treinaram no conjunto de dados SlimPajama para compará-lo com outros métodos baseados em transformadores, como LLAMA, e métodos baseados em espaço de estados, como MAMBA. Eles usaram esta notação de xLSTM a:b, onde a é o número de blocos mLSTM eb é o número de blocos sLSTM na pilha.

Em termos de precisões, eles relatam precisões relativas escalonando as precisões entre 0 e 1, onde 0 é aleatório e 1 é perfeito.

As avaliações indicam que o XLSTM tem melhor desempenho em tarefas como Paridade, enquanto Llama e Mamba têm desempenho ruim.

A partir dos resultados, algo que é de particular interesse é a tarefa de paridade, onde os transformadores ou os modelos de espaço de estados tendem a lutar sem mistura de memória ou rastreamento de estado. Podemos ver que neste tipo de tarefa, o xLSTM atinge uma precisão de 1 quando usamos os blocos sLSTM e mLSTM juntos.

Eles também fizeram alguns estudos de ablação para mostrar a robustez dos XLSTMs. Eles são fáceis de entender no papel. Além disso, este artigo é mais sobre as novidades arquitetônicas dos XLSTMs, por isso não vou entrar nos resultados experimentais aqui.

Gritar

Se você gostou deste artigo, por que não me seguir no Twitter onde compartilho atualizações de pesquisas dos principais laboratórios de IA todos os dias da semana?

Também por favor assine meu Canal do Youtube onde explico conceitos e artigos de IA visualmente.

Conclusão

Espero que este artigo tenha simplificado e facilitado a compreensão da arquitetura XLSTM, por que precisamos deles e como eles poderiam potencialmente ultrapassar os transformadores em um futuro próximo.

Vamos esperar e ver o que eles têm reservado. Vejo você na minha próxima…