Pesquisadores do MIT-IBM Watson AI Lab desenvolveram uma arquitetura expressiva que proporciona melhor rastreamento de estado e raciocínio sequencial em modelos de linguagem latente (LLMs) em textos longos.

Grandes modelos de linguagem têm dificuldades com mudanças de estado comuns em textos longos, como a interação de um gato com uma caixa ao longo do tempo e a possibilidade de a caixa se deteriorar. Agora, um trabalho de pesquisadores do MIT-IBM Watson AI Lab consegue superar o método mais avançado atualmente, o RoPE. Créditos: Imagem: AdobeStock
A maioria das línguas usa a posição das palavras e a estrutura da frase para extrair significado. Por exemplo, "O gato sentou-se na caixa" não é o mesmo que "A caixa estava sobre o gato". Ao longo de um texto extenso, como um documento financeiro ou um romance, a sintaxe dessas palavras provavelmente evolui.
Da mesma forma, uma pessoa pode estar rastreando variáveis em um trecho de código ou seguindo instruções que possuem ações condicionais. Esses são exemplos de mudanças de estado e raciocínio sequencial nos quais esperamos que os sistemas de inteligência artificial de última geração se destaquem; no entanto, o mecanismo de atenção de ponta existente nos Transformers — a arquitetura usada principalmente em grandes modelos de linguagem (LLMs) para determinar a importância das palavras — possui limitações teóricas e empíricas quando se trata dessas capacidades.
Um mecanismo de atenção permite que um LLM (Linguagem de Aprendizado de Máquina) revise partes anteriores de uma consulta ou documento e, com base em seu treinamento, determine quais detalhes e palavras são mais importantes; no entanto, esse mecanismo sozinho não entende a ordem das palavras. Ele "enxerga" todas as palavras de entrada, também conhecidas como tokens, simultaneamente e as processa na ordem em que são apresentadas. Por isso, pesquisadores desenvolveram técnicas para codificar informações de posição. Isso é fundamental para domínios altamente estruturados, como a linguagem. Mas o método predominante de codificação de posição, chamado codificação de posição rotativa (RoPE), leva em consideração apenas a distância relativa entre os tokens em uma sequência e é independente dos dados de entrada. Isso significa que, por exemplo, palavras que estão a quatro posições de distância, como "gato" e "caixa" no exemplo acima, receberão a mesma rotação matemática fixa específica para essa distância relativa.
Agora, uma pesquisa liderada pelo MIT e pelo MIT-IBM Watson AI Lab produziu uma técnica de codificação conhecida como "PaTH Attention" que torna a informação posicional adaptativa e sensível ao contexto, em vez de estática, como acontece com o RoPE.
“Os transformadores permitem a modelagem precisa e escalável de muitos domínios, mas apresentam limitações em relação ao rastreamento de estado, uma classe de fenômenos que se acredita ser fundamental para importantes capacidades desejadas em nossos sistemas de IA. Portanto, a questão crucial é: como podemos manter a escalabilidade e a eficiência dos transformadores, ao mesmo tempo que possibilitamos o rastreamento de estado?”, afirma Yoon Kim, autor sênior do artigo, professor associado do Departamento de Engenharia Elétrica e Ciência da Computação (EECS), membro do Laboratório de Ciência da Computação e Inteligência Artificial (CSAIL) e pesquisador do Laboratório de IA Watson do MIT-IBM.
Um novo artigo sobre este trabalho foi apresentado no início deste mês na Conferência sobre Sistemas de Processamento de Informação Neural (NeurIPS). Os coautores de Kim incluem o autor principal Songlin Yang, estudante de pós-graduação em Engenharia Elétrica e Ciência da Computação (EECS) e ex-estagiário do Programa de Verão do Laboratório de IA Watson do MIT-IBM; Kaiyue Wen, da Universidade Stanford; Liliang Ren, da Microsoft; e Yikang Shen, Shawn Tan, Mayank Mishra e Rameswar Panda, da IBM Research e do Laboratório de IA Watson do MIT-IBM.
Caminho para a compreensão
Em vez de atribuir a cada palavra uma rotação fixa com base na distância relativa entre os tokens, como faz o RoPE, o PaTH Attention é flexível, tratando as palavras intermediárias como um caminho composto por pequenas transformações dependentes dos dados. Cada transformação, baseada em uma operação matemática chamada reflexão de Householder, age como um pequeno espelho que se ajusta dependendo do conteúdo de cada token que atravessa. Cada passo em uma sequência pode influenciar como o modelo interpreta as informações posteriormente. O efeito cumulativo permite que o sistema modele como o significado muda ao longo do caminho entre as palavras, e não apenas a distância entre elas. Essa abordagem permite que os transformadores acompanhem como as entidades e os relacionamentos mudam ao longo do tempo, dando-lhes uma sensação de "memória posicional". Imagine isso como caminhar por um caminho enquanto experimenta o ambiente e como ele o afeta. Além disso, a equipe também desenvolveu um algoritmo eficiente em termos de hardware para calcular com mais eficiência as pontuações de atenção entre cada par de tokens, de modo que a transformação matemática cumulativa do PaTH Attention seja comprimida e dividida em cálculos menores, tornando-a compatível com o processamento rápido em GPUs.
Os pesquisadores do MIT-IBM exploraram então o desempenho do PaTH Attention em tarefas sintéticas e do mundo real, incluindo raciocínio, benchmarks de contexto longo e treinamento completo de LLM (Modelo de Aprendizagem Lógica) para verificar se ele melhorava a capacidade de um modelo de rastrear informações ao longo do tempo. A equipe testou sua capacidade de seguir o comando "escrever" mais recente, apesar de várias etapas de distração e testes de recordação em várias etapas, tarefas difíceis para métodos de codificação posicional padrão, como o RoPE. Os pesquisadores também treinaram LLMs de tamanho médio e os compararam com outros métodos. O PaTH Attention melhorou a perplexidade e superou outros métodos em benchmarks de raciocínio nos quais não foi treinado. Eles também avaliaram a recuperação, o raciocínio e a estabilidade com entradas de dezenas de milhares de tokens. O PaTH Attention demonstrou consistentemente capacidade de reconhecimento de conteúdo.
“Descobrimos que, tanto em tarefas de diagnóstico projetadas para testar as limitações dos Transformers quanto em tarefas de modelagem de linguagem do mundo real, nossa nova abordagem foi capaz de superar os mecanismos de atenção existentes, mantendo sua eficiência”, afirma Kim. Além disso, “ficaria entusiasmado em ver se esses tipos de codificações de posição dependentes de dados, como o PATH, melhoram o desempenho dos Transformers em domínios estruturados como a biologia, na análise de proteínas ou DNA”.
Pensar de forma mais abrangente e eficiente
Em seguida, os pesquisadores investigaram o desempenho do mecanismo de Atenção PaTH se ele imitasse mais de perto a cognição humana, na qual ignoramos informações antigas ou menos relevantes ao tomar decisões. Para isso, combinaram a Atenção PaTH com outro esquema de codificação de posição conhecido como Forgetting Transformer (FoX), que permite que os modelos "esqueçam" seletivamente. O sistema PaTH-FoX resultante adiciona uma maneira de reduzir a importância das informações de forma dependente dos dados, alcançando resultados expressivos em benchmarks de raciocínio, compreensão de contextos longos e modelagem de linguagem. Dessa forma, a Atenção PaTH amplia o poder expressivo das arquiteturas Transformer.
Kim afirma que pesquisas como essa fazem parte de um esforço mais amplo para desenvolver a “próxima grande novidade” em IA. Ele explica que um dos principais impulsionadores das revoluções de aprendizado profundo e IA generativa tem sido a criação de “blocos de construção de propósito geral que podem ser aplicados a diversos domínios”, como “camadas de convolução, camadas de RNN [rede neural recorrente]” e, mais recentemente, transformadores. Olhando para o futuro, Kim observa que considerações como precisão, expressividade, flexibilidade e escalabilidade de hardware têm sido e continuarão sendo essenciais. Como ele mesmo diz, “o objetivo principal da pesquisa em arquitetura moderna é criar esses novos elementos básicos que mantenham ou aprimorem a expressividade, ao mesmo tempo que sejam escaláveis”.
Este trabalho foi financiado, em parte, pelo MIT-IBM Watson AI Lab e pelo programa AI2050 da Schmidt Sciences.