Retrofitting Linear Attention into Diffusion Language Models

Imagine o mundo da inteligência artificial como uma biblioteca enorme e movimentada onde um robô bibliotecário superinteligente está tentando escrever uma história. Durante anos, a maneira padrão como esse bibliotecário trabalhava era "autorregressiva": ele escrevia uma palavra, parava, pensava sobre tudo o que tinha acabado de escrever, escrevia a próxima palavra, parava e pensava novamente. Era como construir uma torre tijolo por tijolo, esperando a cola secar antes de adicionar o próximo. Embora fosse um método confiável, essa forma era lenta, especialmente para histórias longas.

Recentemente, um novo estilo chamado "Difusão" chegou. Em vez de construir tijolo por tijolo, o bibliotecário de Difusão começa com uma página em branco cheia de pontos de interrogação e tenta adivinhar a história inteira de uma vez, refinando os palpites repetidamente até que as palavras façam sentido. Isso é como ter uma equipe de artistas pintando um mural simultaneamente; eles podem trabalhar muito mais rápido porque não estão esperando o tijolo anterior secar. No entanto, mesmo este método rápido tem um problema. À medida que o mural fica maior, os artistas precisam constantemente olhar para cada uma das palavras que já decidiram para garantir que as novas palavras se encaixem. É como tentar lembrar de um livro de 10.000 páginas enquanto se escreve a página 101, o que acaba atrasando a todos e preenchendo a memória da biblioteca.

Este é o quebra-cabeça que uma equipe de pesquisadores abordou. Eles perguntaram: "Podemos tornar este método de Difusão rápido ainda mais rápido mudando a forma como o bibliotecário lembra o passado?". A resposta deles foi um truque inteligente chamado "Atenção Híbrida por Blocos" (Block-Hybrid Attention). Eles perceberam que, embora o bibliotecário precise se lembrar perfeitamente do parágrafo atual (para manter a fluidez das frases), ele não precisa reler cada palavra dos capítulos anteriores em alta definição. Em vez disso, eles podem resumir os capítulos antigos em uma "folha de resumo" minúscula e de tamanho fixo. Isso permite que o bibliotecário escreva as novas partes da história em uma velocidade relâmpago sem ficar sobrecarregado pelo peso de todo o histórico.

A Grande Ideia: Um Sistema de Memória de Duas Velocidades

O artigo apresenta um método para atualizar (ou fazer um upgrade) de um modelo de linguagem de Difusão existente e poderoso (especificamente um modelo de 16 bilhões de parâmetros chamado LLaDA 2.1) para que ele possa usar esta estratégia de "folha de resumo" sem precisar ser treinado do zero.

Pense no cérebro do modelo como tendo 20 camadas diferentes de pensamento. No modelo original, cada uma das camadas atua como um bibliotecário superatento que lê cada palavra anterior para entender a atual. Isso é preciso, mas pesado. A inovação dos pesquisadores, a Atenção Híbrida por Blocos, divide o trabalho em dois modos:

1. O Modo "Agora" (Atenção Exata): Quando o modelo está trabalhando no bloco atual de palavras que está gerando ativamente, ele mantém seu modo "superatento". Ele usa a memória padrão, de alta capacidade, para observar cada palavra no parágrafo atual para garantir que as frases façam sentido perfeito.
2. O Modo "Passado" (Atenção Linear): Quando o modelo precisa se lembrar do que aconteceu nos blocos anteriores (os capítulos finalizados), ele muda para a "Atenção Linear". Em vez de reler o livro inteiro, ele consulta um estado de resumo de tamanho fixo. É como consultar o índice de um livro ou um resumo de uma página, em vez de reler todo o texto. Este resumo permanece do mesmo tamanho, não importa o quão longo o livro se torne.

Como Eles Fizeram: O Upgrade em Duas Etapas

Os pesquisadores não apenas trocaram as peças e esperaram pelo melhor; eles usaram um processo cuidadoso de duas etapas para garantir que o modelo não peresse sua inteligência durante a atualização.

• Etapa 1: O Treinamento de Sombra. Eles pegaram o modelo original e inteligente (o "Professor") e o congelaram. Em seguida, substituíram 6 das 20 camadas de atenção pela versão "Híbrida por Blocos" (o "Aluno"). O Aluno foi treinado para imitar exatamente a saída do Professor, usando uma versão "corrompida" do texto como entrada. Foi como ter um aluno observando um mestre chef, tentando copiar os movimentos exatos do chef e provar os mesmos sabores, mas apenas para os pratos específicos que o aluno foi designado para mudar. Esta etapa levou cerca de 24 horas.
• Etapa 2: O Ajuste Fino (Fine-Tuning). Uma vez que o Aluno aprendeu a imitar o Professor, os pesquisadores deixaram todo o modelo trabalhar junto novamente. Eles usaram uma técnica chamada LoRA (Low-Rank Adaptation) para fazer ajustes minúsculos e precisos nas conexões do modelo. Isso foi como dar ao modelo atualizado um rápido "polimento" para corrigir quaisquer pequenos problemas que surgiram porque as novas partes agora estavam conversando com as partes antigas. Esta etapa levou cerca de 6 horas.

Os Resultados: Mais Rápido, Mais Leve e Ainda Inteligente

A equipe testou seu modelo atualizado, que nomearam de LLaDA-HYBRID, para ver se ele ainda era bom em escrever e resolver problemas, e se era realmente mais rápido.

Ele ainda sabe pensar?
Surpreendentemente, sim. Mesmo que tenham substituído 6 das 20 camadas com este método de "folha de resumo", o desempenho do modelo permaneceu muito próximo do original.

• Em um teste de codificação chamado HumanEval, o modelo original obteve 75,6%, e o modelo híbrido obteve 72,0%.
• Em um teste de matemática chamado CMATH, o original obteve 88,3%, e o híbrido obteve 86,7%.
• Curiosamente, em outro teste de codificação, o MBPP+, o modelo híbrido melhorou, passando de 57,7% para 63,0%.
  O artigo sugere que, embora o modelo seja ligeiramente menos perfeito em tarefas de raciocínio mais difíceis (como um quiz científico difícil chamado GPQA-Diamond, onde caiu de 38,9% para 30,8%), ele preservou amplamente sua capacidade de escrever código e resolver problemas matemáticos.

Ele é realmente mais rápido?
É aqui que acontece a mágica. Como o modelo não precisa mais reler todo o histórico da conversa para cada nova palavra, ele pode lidar com mais solicitações ao mesmo tempo.

• Ao testar quantas palavras o modelo poderia gerar por segundo, o modelo híbrido foi 1,7 vezes mais rápido que o original em seu pico de desempenho (lidando com 128 solicitações simultâneas).
• O uso de memória do modelo original crescia à medida que a história ficava mais longa, eventualmente atingindo um limite onde não conseguia lidar com mais usuários. O modelo híbrido, graças à sua "folha de resumo" de tamanho fixo, pôde lidar com mais solicitações simultâneas antes de ficar sem memória.

A Conclusão

Este artigo sugere que nem sempre precisamos construir uma nova IA mais rápida do zero. Em vez disso, podemos pegar uma IA poderosa já existente e dar a ela um sistema de "memória híbrida". Ao manter seu foco aguçado no momento presente enquanto resume o passado em uma forma compacta, podemos tornar esses modelos significativamente mais rápidos e eficientes. Os pesquisadores descobriram que essa atualização poderia ser feita em cerca de 60 horas de treinamento em dados públicos, oferecendo um caminho promissor para tornar os serviços de IA mais rápidos e baratos sem sacrificar muito de sua inteligência. No entanto, os autores observam que isto é apenas o começo; eles atualizaram apenas 6 camadas, e trabalhos futuros podem encontrar maneiras de atualizar ainda mais camadas ou lidar com histórias ainda mais longas.