Arquitetura Titan do Google Rompe o Gargalo de Memória do Transformer

Introdução à Titan: Uma Nova Arquitetura do Google

O mundo da tecnologia está agitado com a Titan, uma arquitetura inovadora que está emergindo do Google. Ela foi projetada para desafiar as limitações dos modelos Transformer, particularmente em como eles lidam com a memória. Esta nova arquitetura está ganhando atenção significativa como um potencial sucessor do Transformer, especialmente considerando seu desenvolvimento por uma equipe dentro do Google.

O Desafio da Memória em Modelos Existentes

Modelos tradicionais como LSTM e Transformer, embora inovadores, enfrentam desafios na simulação de memória semelhante à humana. Esses desafios incluem:

Capacidade Limitada: Os dados são frequentemente comprimidos em um estado oculto de tamanho fixo, restringindo a quantidade de informação que pode ser retida.
Sobrecarga Computacional: Embora capaz de capturar dependências de longo alcance, o custo computacional aumenta quadraticamente com o comprimento da sequência, tornando-o ineficiente para sequências muito longas.
Superdependência de Dados de Treinamento: Simplesmente memorizar dados de treinamento nem sempre ajuda na aplicação no mundo real, onde os dados de teste podem ficar fora da distribuição de treinamento.

Abordagem da Titan: Um Módulo de Memória Neuro-Inspirado

A equipe da Titan adotou uma abordagem diferente, buscando codificar informações nos parâmetros de uma rede neural. Eles desenvolveram um metamodelo online projetado para aprender como lembrar e esquecer dados específicos durante o teste. Este modelo é inspirado em princípios neuropsicológicos, incorporando os seguintes elementos-chave:

Surpresa como Gatilho: Eventos inesperados são mais facilmente lembrados. A "surpresa" é medida pelo gradiente da entrada para o módulo de memória. Quanto maior o gradiente, mais inesperada a entrada.
Mecanismos de Momentum e Esquecimento: Um mecanismo de momentum acumula surpresas de curto prazo na memória de longo prazo, enquanto um mecanismo de esquecimento apaga memórias antigas, evitando o estouro da memória.
Memória Baseada em Perceptron Multicamadas (MLP): O módulo de memória é composto por várias camadas MLP, permitindo que ele armazene abstrações profundas de dados, tornando-o mais poderoso do que as memórias tradicionais baseadas em matrizes.

Essa abordagem de meta-aprendizagem online ajuda o modelo a se concentrar em aprender a se adaptar a novos dados, em vez de simplesmente memorizar dados de treinamento. O módulo também foi projetado para computação paralela, melhorando sua eficiência.

Integração do Módulo de Memória em Arquiteturas de Deep Learning

A equipe de pesquisa da Titan propôs três variações para incorporar seu módulo de memória em arquiteturas de deep learning:

MAC (Memória como Contexto): Este método combina memória de longo prazo e persistente (que codifica o conhecimento da tarefa) como contexto que é inserido no mecanismo de atenção.
MAG (Memória como Gate): Essa abordagem usa fusão controlada do módulo de memória com um mecanismo de atenção de janela deslizante em dois ramos.
MAL (Memória como Camada): Aqui, o módulo de memória é implementado como uma camada independente que comprime informações históricas antes de alimentá-las ao mecanismo de atenção.

A equipe descobriu que cada variação tem seus pontos fortes e fracos.

Desempenho e Vantagens da Titan

A Titan demonstrou desempenho superior em diversas tarefas, incluindo modelagem de linguagem, raciocínio de senso comum e previsão de séries temporais. Ela superou modelos de ponta como Transformer e Mamba. Notavelmente, o módulo de memória de longo prazo (LMM) sozinho superou modelos de linha de base em várias tarefas, mostrando suas capacidades de aprendizado independente sem memória de curto prazo (atenção).

Em um teste de "agulha no palheiro" projetado para encontrar pistas detalhadas em textos longos, a Titan manteve cerca de 90% de precisão mesmo quando os comprimentos da sequência aumentaram de 2k para 16k. A equipe indica que os testes padrão não exibem totalmente as vantagens da Titan no tratamento de textos longos. A Titan também superou modelos como GPT4, Mamba e até mesmo Llama3.1 com RAG em uma tarefa que exigia inferência de fatos espalhados por documentos extremamente longos.

A Titan também mostrou um desempenho impressionante em áreas específicas, como previsão de séries temporais e modelagem de sequência de DNA.

A Equipe por Trás da Titan

A pesquisa foi conduzida por uma equipe do grupo de algoritmos e otimização do Google Research NYC, não fazendo parte atualmente do Google DeepMind.

Ali Behrouz, um estagiário da Cornell University, é o primeiro autor do artigo.
Zhong Peilin, um ex-aluno da Tsinghua University e doutor pela Columbia University, é um cientista pesquisador no Google desde 2021. Ele é notável por ter publicado um artigo como primeiro autor no STOC 2016 como estudante de graduação.
Vahab Mirrokni, um Google Fellow e VP, lidera a equipe.

A equipe desenvolveu a Titan usando Pytorch e Jax e planeja liberar o código para treinamento e avaliação em breve.