Uma análise aprofundada dos Titãs: o caminho para a convergência da memória de longo prazo e a modelagem eficiente de sequências

Titãs: Aprendendo a Memorizar na Hora da Prova Original: https://arxiv.org/pdf/2501.00663v1

Implementação não oficial da Arquitetura Titans: https://github.com/lucidrains/titans-pytorch

Um, Histórico e motivação: Limitações do Transformer e a inspiração da memória humana

1. Limitações do Transformer: um gargalo no processamento de sequências longas

Transformador Desde sua introdução, o modelo tem sido capaz de fazer uma contribuição significativa para o desenvolvimento do modelo em virtude de seu poderosoMecanismo de autoatençãoque fez um progresso revolucionário em áreas como processamento de linguagem natural e visão computacional. No entanto, à medida que a complexidade da tarefa aumenta, o Transformer expõe gradualmente alguns problemas críticos ao lidar com sequências longas:

• A alta complexidade computacional limita a escalabilidade do modelo :.
  • O mecanismo de autoatenção requer o cálculo de cada token similaridade com todos os outros tokens da sequência, com complexidade de tempo e espaço de O(N²)que N é o comprimento da sequência.
  • Isso significa que o consumo de computação e de memória cresce em etapas quadradas quando o comprimento da sequência aumenta, limitando severamente a capacidade do modelo de lidar com sequências longas. Por exemplo, o Transformer geralmente não tem capacidade para lidar com tarefas como texto longo, compreensão de vídeo ou previsão de séries temporais de longo prazo.
  

  Figura 1: Processo computacional do mecanismo de autoatenção.

• Janela de contexto limitada para capturar dependências de longa distância :.
  • Para aliviar a carga computacional, os Transformers normalmente usam uma janela de contexto de comprimento fixo (por exemplo, 512 ou 1024), o que significa que o modelo só pode se concentrar nas informações dentro da janela atual.
  • No entanto, muitas tarefas do mundo real exigem modelos que possam capturar dependências em escalas de tempo mais longas, como a compreensão de informações contextuais em textos ou diálogos longos, a integração de informações de diferentes pontos no tempo em vídeos e a realização de previsões usando tendências e padrões de longo prazo em dados históricos.

2. O compromisso do transformador linear: eficiência versus desempenho

Para resolver o gargalo computacional do Transformer, os pesquisadores propuseram oTransformador linearOs principais aprimoramentos são:

• Substituição do softmax por uma função de kernel. A substituição do cálculo do softmax no mecanismo de autoatenção por uma função de kernel reduz a complexidade computacional para O(N).
• Raciocínio paralelo. O processo computacional do Linear Transformer pode ser representado como uma forma cíclica, o que permite um raciocínio mais eficiente.

No entanto, o Transformador Linear também tem algumas limitações:

• Degradação do desempenho.
  • O truque do kernel degrada o modelo em uma rede recorrente linear em que os dados são compactados em um estado com valor de matriz, resultando em um desempenho inferior ao do Transformer padrão.
  • Esse tipo de compactação é difícil de capturar efetivamente dependências não lineares complexas.
• Problemas de gerenciamento de memória.
  • O Linear Transformer comprime dados históricos em uma matriz de tamanho fixo, mas ao lidar com contextos muito longos, essa compactação resulta em umestouro de memóriaque afeta o desempenho do modelo.
  

  Figura 2: Processo de atualização da memória para o Linear Transformer.

3. Inspiração do sistema de memória humana: criação de mecanismos mais fortes para a memória de longo prazo

Para superar os desafios acima, os autores começaramsistema de memória humanaInspirado no

• A relação entre memória e aprendizado: o
  • A tese se baseia nas definições de memória e aprendizado da literatura neuropsicológica, considerando a memória como uma atualização neural induzida por informações, e definindo o aprendizado como o processo de aquisição de memórias eficazes e úteis com base em um objetivo.
  • Isso significa queNão é possível obter um aprendizado eficaz sem um forte mecanismo de memória.
• A natureza multinível da memória humana: o
  • O sistema de memória humana não é uma estrutura única, mas consiste em vários subsistemas, como a memória de curto prazo, a memória de trabalho e a memória de longo prazo, cada uma delas com funções e estruturas organizacionais diferentes e capaz de operar de forma independente.
  • Essa natureza de várias camadas permite que os seres humanos armazenem, recuperem e gerenciem informações de forma eficiente.
• Deficiências dos modelos existentes.
  • As arquiteturas de redes neurais existentes (de redes Hopfield a LSTMs e Transformers) enfrentam desafios para lidar com a generalização, a extrapolação de comprimento e a inferência, recursos que são essenciais para muitas tarefas complexas do mundo real.
  • Essas arquiteturas, embora inspiradas no cérebro humano, não tinham a capacidade dememória efêmeraA modelagem eficaz doNatureza multinível do sistema de memóriaA simulação.

Dois, Inovação principal: módulo de memória neural de longo prazo e arquitetura de titãs

Com base nessas reflexões, os autores propõem as seguintes inovações:

1. Módulo de memória neural de longo prazo

(1) Conceito de design.

• Mecanismos de aprendizagem metacontextual.
  • O módulo foi projetado como ummetamodeloO sistema de gerenciamento de dados é um sistema de gerenciamento de dados que permite que o usuário memorize/armazene dados em seus parâmetros durante o teste.
  • esse tipo deAprendizagem on-linepermite que o modelo adapte dinamicamente sua memória à entrada atual, em vez de depender de memórias que foram fixadas durante o pré-treinamento.
• Atualizações de memória baseadas em surpresas: o
  • Os autores se baseiam no mecanismo de memória humana segundo o qual "eventos surpreendentes têm maior probabilidade de serem lembrados" e propõem um método baseado emgrau de surpresaO mecanismo de atualização da memória.
  • grau de surpresaAo calcular a rede neural em relação à entradagradientepara medir, quanto maior o gradiente, mais os dados de entrada diferem dos dados históricos e são mais dignos de serem lembrados.
    • Esse método é eficaz para capturar informações importantes nos dados e armazená-las na memória de longo prazo.
    • Por outro lado, o Linear Transformer só pode realizar transformações lineares com base nos dados de entrada atuais, o que dificulta a captura efetiva de dependências de longo alcance.
  

  Figura 3: Mecanismo de atualização de memória baseado em surpresa.

(2) Tecnologias-chave.

• Mecanismos de impulso.
  • Para evitar que o modelo seja excessivamente influenciado por um único evento surpreendente, os autores introduziram omecanismo de impulsolevando em conta a surpresa de momentos passados.
  • Isso significa que o modelo leva em conta a surpresa dos inputs atuais e históricos, o que resulta em atualizações de memória mais suaves.
• Mecanismos de atenuação.
  • Para evitar o estouro de memória, os autores também introduziram a funçãomecanismo de decaimentoporqueda de pesomaneira de esquecer gradualmente informações sem importância.
  • O mecanismo pode ser visto como umMecanismo de bloqueioque pode apagar seletivamente as memórias conforme necessário.
  • Os autores destacam que esse mecanismo de decaimento é uma generalização do mecanismo de esquecimento em modelos recorrentes modernos e é equivalente à otimização de redes meta-neurais sob descida de gradiente de lote pequeno, momento e decaimento de peso.

(3) Estruturas de memória.

• Em contraste com o modelo de memória linear tradicional, os autores usaram umPerceptron de múltiplas camadas (MLP) como um módulo de memória.
  • O MLP tem uma representação não linear mais forte e pode armazenar e recuperar informações complexas com mais eficiência.
  • Por outro lado, os transformadores lineares só podem usar estados com valor de matriz para armazenar informações, o que dificulta a captura de relações não lineares complexas.

2. Arquitetura dos Titãs: Integrando a memória de longo e curto prazo

Depois de projetar o módulo de memória neural de longo prazo, os autores refletiram sobre como integrá-lo efetivamente em uma arquitetura de aprendizagem profunda e propuseram a arquitetura Titans com os seguintes recursos principais:

(1) Três supercabeçotes trabalhando em conjunto.

• Núcleo.
  • Consiste na memória de curto prazo e é responsável pelo processo principal de manipulação de dados.
  • fazer uso deMecanismos de atenção para tamanhos de janela finitoscomo o Sliding Window Attention (SWA) ou o Fully Connected Attention (FCA).
  • A memória de curto prazo pode ser vista comomemória de curto prazoque é usado para capturar dependências no contexto atual.
• Memória de longo prazo.
  • Responsável por armazenar/lembrar informações de um passado longínquo.
  • É usado o módulo de memória neural de longo prazo descrito acima.
  • A memória de longo prazo pode ser vista comomemória efêmeraque é usado para armazenar e recuperar informações em um período de tempo mais longo.
• Memória persistente.
  • é um conjunto de parâmetros aprendíveis, mas independentes de dados, que codificam um conhecimento prévio sobre a tarefa.
  • Semelhante aos parâmetros da camada totalmente conectada no Transformer, mas com funcionalidade diferente.
  • A memória persistente pode ser vista comometamemóriaque é usado para armazenar conhecimento relacionado à tarefa, como regras gramaticais, conhecimento de senso comum, etc.
  

  Figura 4: Diagrama esquemático da arquitetura do Titans (variante MAC).

(2) Três tipos diferentes de integração.

• Memória como contexto (MAC).
  • Conectar memórias longas e persistentes a sequências de entrada como informações adicionais ao contexto atual.
  • Os mecanismos de atenção determinam quais informações precisam ser armazenadas na memória de longo prazo.
  • No momento do teste, os parâmetros de memória persistente foram mantidos fixos e o módulo de atenção foi ponderado para o aprendizado contextual, enquanto o módulo de memória de longo prazo continuou a aprender/lembrar informações.
  • Esse design permite que o modelo utilize de forma flexível as informações da memória de longo prazo com base na entrada atual.
• Memória como Gating (MAG).
  • Usando a atenção de janela deslizante como memória de curto prazo e módulos de memória neural como memória de longo prazo.
  • Combinando os dois por meio de mecanismos de gating, por exemplo, normalizando os dois usando pesos de valores vetoriais aprendidos e, em seguida, aplicando uma função de ativação não linear.
  • Esse design pode ser visto como uma arquitetura de vários cabeçotes, na qual os diferentes cabeçotes são estruturados de forma diferente.

  
  Figura 5: Diferentes variantes da arquitetura Titans (MAC e MAG).

• Memória como uma camada (MAL).
  • O uso do módulo de memória neural como uma camada de uma rede neural profunda comprime as informações contextuais passadas e atuais antes do módulo de atenção.
  • Esse design é mais comum na literatura, por exemplo, o modelo H3.

(3) Pontos fortes.

• Gerenciamento de memória mais flexível.
  • Com o uso de módulos de memória como ramificações contextuais ou controladas, a arquitetura Titans é capaz de utilizar dinamicamente as informações da memória de longo prazo com base na entrada atual.
  • Isso proporciona mais flexibilidade do que a abordagem tradicional de usar módulos de memória como camadas.
• Expressão mais forte.
  • A sinergia dos três supercabeçotes permite que a arquitetura Titans processe dados sequenciais longos com mais eficiência e integre os benefícios da memória de curto prazo, da memória de longo prazo e da memória persistente.
• Escalabilidade.
  • Em comparação com o Transformer, a arquitetura Titans tem melhor escalabilidade ao lidar com sequências longas e é capaz de manter o alto desempenho em uma janela de contexto maior.

Três, Resultados experimentais e análise: validação da arquitetura Titans

Os autores realizaram experimentos abrangentes em várias tarefas para avaliar o desempenho da arquitetura Titans e suas variantes:

1. Modelagem linguística e raciocínio de senso comum.

• Configuração experimental.
  • Três modelos Titans de tamanhos diferentes (340M, 400M, 760M parâmetros) foram usados, bem como vários modelos de linha de base, incluindo Transformer++, RetNet, GLA, Mamba, Mamba2, DeltaNet, TTT e Gated DeltaNet.
  • O conjunto de dados FineWeb-Edu foi usado como dados de treinamento.
• Principais conclusões.
  • No modelo não misto, o módulo Neural Long-Term Memory obteve o melhor desempenho nas métricas de perplexidade e precisão.
  • Todas as três variantes do Titans (MAC, MAG, MAL) superaram o Samba (Mamba + atenção) e o Gated DeltaNet-H2 (Gated DeltaNet + atenção).
  • O MAC tem melhor desempenho ao lidar com dependências de longo alcance, enquanto o MAG e o MAC superam a variante MAL.
  

  Figura 6: Comparação do desempenho do Titans com o modelo de linha de base em tarefas de modelagem de linguagem e raciocínio de senso comum.

2. Missão "Needle in the Haystack" (agulha no palheiro).

• Configuração experimental.
  • Usando a tarefa Single NIAH (S-NIAH) no benchmarking RULER, a capacidade de recuperação do modelo foi avaliada em sequências de 2K, 4K, 8K e 16K de comprimento.
• Principais conclusões.
  • O módulo Neural Long-Term Memory obteve os melhores resultados em todas as três tarefas.
  • As variantes Titans também tiveram um bom desempenho, sendo que a variante MAC teve o melhor desempenho.

3. Benchmarking do BABILong.

• Configuração experimental.
  • A tarefa exige que o modelo raciocine sobre informações factuais distribuídas em documentos extremamente longos.
  • Há configurações de menos amostras e configurações de ajuste fino.
• Principais conclusões.
  • Na configuração de menos amostras, o Titans supera todas as linhas de base, inclusive os modelos com maior número de parâmetros, como o GPT-4 e o GPT4o-mini.
  • O Titans também supera todos os modelos em configurações de ajuste fino, até mesmo modelos muito grandes como o GPT-4.
  • Em comparação com os modelos de memória baseados em transformador (RMT), os Titãs apresentam melhor desempenho, principalmente devido a seus fortes recursos de memória.
  

  Figura 7: Comparação do desempenho do Titans com o modelo de linha de base no benchmark BABILong.

4. Previsão de séries temporais.

• Configuração experimental.
  • Usando a estrutura do Simba, o módulo Mamba foi substituído por um módulo de memória neural de longo prazo.
  • Avaliado nos conjuntos de dados de referência ETT, ECL, Traffic e Weather.
• Principais conclusões.
  • O módulo Neural Long-Term Memory supera todas as linhas de base, incluindo as arquiteturas baseadas em Mamba, Linear Model e Transformer.

5. Modelagem de DNA.

• Configuração experimental.
  • Avalie o desempenho de tarefas downstream de modelos pré-treinados no GenomicsBenchmarks.
• Principais conclusões.
  • Titans (LMM) é competitivo em diferentes tarefas de genômica downstream e está no mesmo nível dos métodos mais avançados.

6. Análise de eficiência.

• Principais conclusões.
  • O módulo Neural Long-Term Memory é um pouco mais lento para treinar em comparação com outros modelos recorrentes, principalmente devido à sua memória mais profunda e ao processo de transformação mais complexo, bem como ao kernel altamente otimizado implementado no Mamba2.
  • Titans (MAL) é mais rápido do que a linha de base e do que o módulo de memória, principalmente devido ao uso do kernel altamente otimizado do FlashAttention.

7. Estudos de ablação.

• Principais conclusões.
  • Todos os componentes do design da memória neural contribuem positivamente para o desempenho, sendo que o decaimento do peso, o momentum, a convolução e a memória persistente são os que mais contribuem.
  • O projeto arquitetônico também tem um impacto significativo no desempenho, com o MAC e o MAG apresentando desempenho próximo um do outro em tarefas de modelagem de linguagem e raciocínio de senso comum, enquanto o MAC apresenta melhor desempenho em tarefas de contexto longo.

Quatro, Pontos inovadores e pontos fortes da tese

1. É proposto um novo módulo de memória neural de longo prazo: o
   • Ele se baseia nos principais elementos dos mecanismos de memória humana, como surpresa, impulso e esquecimento, para obter uma atualização e um armazenamento mais eficientes da memória.
   • Uma rede neural profunda é usada como um módulo de memória para dar ao modelo maior poder de expressão.
2. A arquitetura do Titans foi projetada para combinar a memória de longo e curto prazo.
   • Três abordagens de integração diferentes são propostas, oferecendo opções flexíveis para diferentes cenários de aplicativos.
   • A sinergia dos três supercabeçotes, núcleo, memória de longo prazo e memória persistente, permite que o modelo processe dados sequenciais longos com mais eficiência.
3. É excelente em várias tarefas.
   • Seja para modelagem de linguagem, raciocínio de senso comum, previsão de séries temporais e modelagem de DNA, a arquitetura Titans demonstra um desempenho avançado que supera os modelos existentes de transformador e loop linear.
4. Escalável.
   • A capacidade de manter o alto desempenho em uma janela de contexto maior abre a possibilidade de processar sequências muito longas.

Cinco, perspectivas futuras

Embora a arquitetura Titans tenha produzido resultados impressionantes em várias frentes, ainda há as seguintes direções que merecem ser mais exploradas.

1. Explorando arquiteturas de módulos de memória mais complexas: o
   • Por exemplo, a introdução de estruturas de memória hierárquicas ou a combinação de módulos de memória com outros modelos, como redes neurais gráficas.
2. Desenvolvimento de mecanismos mais eficientes para atualização e armazenamento de memória.
   • Por exemplo, técnicas de esparsificação ou técnicas de quantificação são usadas para reduzir o consumo de memória e os custos computacionais.
3. Aplicar a arquitetura dos Titãs a uma gama maior de áreas.
   • Os exemplos incluem compreensão de vídeo, controle de robôs, sistemas de recomendação, etc.
4. Explorar estratégias de treinamento mais eficazes.
   • Por exemplo, a introdução de algoritmos de otimização mais avançados ou o uso de meta-aprendizagem para acelerar o treinamento de modelos.
5. Investigando a interpretabilidade da arquitetura dos Titãs.
   • Uma compreensão mais profunda de como os Titãs armazenam e utilizam as informações da memória de longo prazo pode fornecer novas ideias para a criação de sistemas de IA mais poderosos.

Seis, resumos

A principal contribuição deste documento é:

• Um novo módulo de memória neural de longo prazo é propostoO design é inspirado no sistema de memória humana e incorpora os principais conceitos da aprendizagem profunda, como descida de gradiente, momentum e decaimento de peso.
• Construiu a arquitetura dos Titãsque combina organicamente a memória de longo e curto prazo e explora três métodos de integração diferentes, oferecendo opções flexíveis para diferentes cenários de aplicativos.
• O desempenho superior do Titans foi verificado por meio de experimentos rigorosos.que apresenta bom desempenho em várias tarefas, especialmente ao lidar com tarefas de contexto longo, demonstrando forte escalabilidade e maior precisão.