Os aplicativos empresariais de IA que lidam com documentos grandes ou tarefas de longo prazo enfrentam um grave gargalo de memória. À medida que o contexto aumenta, o mesmo acontece com o cache KV, a área onde a memória de trabalho do modelo é armazenada.
Uma nova técnica desenvolvida por pesquisadores do MIT aborda esse desafio com um método de compactação rápida para o cache KV. A técnica, chamada Correspondência de atençãoconsegue compactar o contexto em até 50x com pouquíssima perda de qualidade.
Embora não seja a única técnica de compactação de memória disponível, o Attention Matching se destaca por sua velocidade de execução e impressionantes capacidades de preservação de informações.
O gargalo de memória do cache KV
Grandes modelos de linguagem geram suas respostas sequencialmente, um token por vez. Para evitar o recálculo de todo o histórico de conversas do zero para cada palavra prevista, o modelo armazena uma representação matemática de cada token anterior processado, também conhecido como pares de chave e valor. Essa memória de trabalho crítica é conhecida como cache KV.
O cache KV é dimensionado com a duração da conversa porque o modelo é forçado a reter essas chaves e valores para todos os tokens anteriores em uma determinada interação. Isso consome recursos de hardware caros. “Na prática, a memória cache KV é o maior gargalo para servir modelos em contexto ultralongo”, disse Adam Zweiger, coautor do artigo, ao VentureBeat. “Ele limita a simultaneidade, força lotes menores e/ou requer descarregamento mais agressivo.”
Em casos de uso corporativos modernos, como análise de contratos legais massivos, manutenção de diálogos multissessões com clientes ou execução de agentes de codificação autônomos, o cache KV pode atingir muitos gigabytes de memória para uma única solicitação do usuário.
Para resolver esse enorme gargalo, a indústria de IA tentou diversas estratégias, mas esses métodos são insuficientes quando implantados em ambientes empresariais onde a compactação extrema é necessária. Uma classe de correções técnicas inclui a otimização do cache KV, removendo tokens que o modelo considera menos importantes ou mesclando tokens semelhantes em uma única representação. Essas técnicas funcionam para compressão leve, mas “degradam-se rapidamente em altas taxas de redução”, segundo os autores.
Os aplicativos do mundo real geralmente dependem de técnicas mais simples, sendo a abordagem mais comum simplesmente descartar o contexto mais antigo quando o limite de memória for atingido. Mas esta abordagem faz com que o modelo perca informações mais antigas à medida que o contexto aumenta. Outra alternativa é o resumo do contexto, onde o sistema faz uma pausa, escreve um breve resumo em texto do contexto mais antigo e substitui a memória original por esse resumo. Embora este seja um padrão do setor, o resumo apresenta muitas perdas e prejudica gravemente o desempenho downstream porque pode remover informações pertinentes do contexto.
Pesquisas recentes provaram que é tecnicamente possível comprimir altamente esta memória usando um método chamado Cartuchos. No entanto, esta abordagem requer o treinamento de modelos de cache KV latentes por meio de otimização matemática lenta e completa. Esse treinamento baseado em gradiente pode levar várias horas em GPUs caras apenas para compactar um único contexto, tornando-o completamente inviável para aplicativos empresariais em tempo real.
Como a correspondência de atenção é compactada sem custo
O Attention Matching atinge taxas de compactação e qualidade de alto nível, ao mesmo tempo que é muito mais rápido do que a otimização baseada em gradiente. Ele contorna o lento processo de treinamento por meio de truques matemáticos inteligentes.
Os pesquisadores perceberam que, para imitar perfeitamente como uma IA interage com sua memória, eles precisam preservar duas propriedades matemáticas ao compactar os vetores originais de chave e valor em um espaço menor. A primeira é a “saída de atenção”, que é a informação real que a IA extrai quando consulta sua memória. A segunda é a “massa de atenção”, que atua como o peso matemático que um token tem em relação a tudo o mais na memória de trabalho do modelo. Se a memória compactada puder corresponder a essas duas propriedades, ela se comportará exatamente como a enorme memória original, mesmo quando novos prompts imprevisíveis do usuário forem adicionados posteriormente.
“A correspondência de atenção é, de certa forma, o objetivo ‘correto’ para fazer a compactação de contexto latente, pois visa diretamente a preservação do comportamento de cada cabeça de atenção após a compactação”, disse Zweiger. Embora a eliminação de tokens e heurísticas relacionadas possam funcionar, combinar explicitamente o comportamento de atenção simplesmente leva a melhores resultados.
Antes de compactar a memória, o sistema gera um pequeno conjunto de “consultas de referência” que atuam como um proxy para os tipos de pesquisas internas que o modelo provavelmente realizará ao raciocinar sobre o contexto específico. Se a memória compactada puder responder com precisão a essas consultas de referência, muito provavelmente conseguirá responder às perguntas reais do usuário posteriormente. Os autores sugerem vários métodos para gerar essas consultas de referência, incluindo anexar um prompt oculto ao documento informando ao modelo para repetir o contexto anterior, conhecido como técnica “repeat-prefill”. Eles também sugerem uma abordagem de “autoestudo”, em que o modelo é solicitado a realizar algumas tarefas sintéticas rápidas no documento, como agregar todos os fatos importantes ou estruturar datas e números em um formato JSON.
Com essas consultas em mãos, o sistema escolhe um conjunto de chaves para preservar no cache KV compactado com base em sinais como o valor de atenção mais alto. Em seguida, ele usa as chaves e consultas de referência para calcular os valores correspondentes junto com um termo de polarização escalar. Essa tendência garante que as informações pertinentes sejam preservadas, permitindo que cada chave retida represente a massa de muitas chaves removidas.
Esta formulação torna possível ajustar os valores com técnicas algébricas simples, como mínimos quadrados ordinários e mínimos quadrados não negativos, evitando totalmente a otimização baseada em gradientes com muitos cálculos. Isso é o que torna o Attention Matching super rápido em comparação com métodos de compactação com muita otimização. Os pesquisadores também aplicam compactação em partes, processando partes contíguas da entrada de forma independente e concatenando-as, para melhorar ainda mais o desempenho em contextos longos.
Correspondência de atenção em ação
Para entender o desempenho desse método no mundo real, os pesquisadores realizaram uma série de testes de estresse usando modelos populares de código aberto, como Llama 3.1 e Qwen-3, em dois tipos distintos de conjuntos de dados corporativos. O primeiro foi o QuALITY, um benchmark padrão de compreensão de leitura que utiliza documentos de 5.000 a 8.000 palavras. O segundo, representando um verdadeiro desafio empresarial, foi o LongHealth, um conjunto de dados altamente denso de 60.000 tokens que contém registros médicos complexos de vários pacientes.
A principal descoberta foi a capacidade do Attention Matching de compactar o cache KV do modelo em 50x sem reduzir a precisão, enquanto leva apenas alguns segundos para processar os documentos. Para atingir o mesmo nível de qualidade anterior, os cartuchos exigiam horas de computação intensiva de GPU por contexto.
Ao lidar com os densos registros médicos, as soluções alternativas padrão da indústria entraram em colapso completamente. Os pesquisadores observaram que quando tentaram usar o resumo de texto padrão nesses registros de pacientes, a precisão do modelo caiu tão baixo que correspondeu à linha de base “sem contexto”, o que significa que a IA funcionou como se não tivesse lido o documento.
A correspondência de atenção supera drasticamente o resumo, mas os arquitetos corporativos precisarão diminuir a taxa de compactação para tarefas densas em comparação com testes de compreensão de leitura mais simples. Como explica Zweiger: “A principal desvantagem prática é que, se você estiver tentando preservar quase tudo no contexto em tarefas com alta densidade de informações, geralmente precisará de uma taxa de compactação mais suave para manter uma precisão forte”.
Os pesquisadores também exploraram o que acontece nos casos em que a precisão absoluta não é necessária, mas sim a economia extrema de memória. Eles executaram o Attention Matching sobre um resumo de texto padrão. Essa abordagem combinada alcançou compactação de 200x. Ele correspondeu com sucesso apenas à precisão do resumo padrão, mas com um consumo de memória muito pequeno.
Uma das experiências interessantes para fluxos de trabalho empresariais foi testar a compactação online, embora observem que esta é uma prova de conceito e não foi testada rigorosamente em ambientes de produção. Os pesquisadores testaram o modelo no teste avançado de raciocínio matemático AIME. Eles forçaram a IA a resolver um problema com um limite de memória física estritamente limitado. Sempre que a memória do modelo ficava cheia, o sistema fazia uma pausa, comprimia instantaneamente sua memória de trabalho em 50% usando o Attention Matching e deixava-o continuar pensando. Mesmo depois de atingir a barreira da memória e ter seu cache KV reduzido até seis vezes consecutivas no meio do pensamento, o modelo resolveu com sucesso os problemas matemáticos. Seu desempenho correspondia a um modelo que recebeu memória enorme e ilimitada.
Há advertências a serem consideradas. Com uma taxa de compressão de 50x, Attention Matching é o vencedor claro no equilíbrio entre velocidade e qualidade. No entanto, se uma empresa tentar levar a compactação a limites extremos de 100x em dados altamente complexos, o método de cartuchos mais lento e baseado em gradiente na verdade o superará.
Os pesquisadores divulgaram o código para correspondência de atenção. No entanto, eles observam que atualmente não se trata de uma simples atualização de software plug-and-play. “Acho que a compactação latente é melhor considerada uma técnica de camada modelo”, observa Zweiger. “Embora possa ser aplicado sobre qualquer modelo existente, requer acesso aos pesos do modelo.” Isso significa que as empresas que dependem inteiramente de APIs fechadas não podem implementar isso sozinhas; eles precisam de modelos de peso aberto.
Os autores observam que a integração desta compactação KV do espaço latente em mecanismos de inferência comerciais existentes e altamente otimizados ainda requer um esforço significativo. A infraestrutura moderna de IA usa truques complexos, como cache de prefixo e empacotamento de memória de comprimento variável, para manter os servidores funcionando de maneira eficiente, e incorporar perfeitamente essa nova técnica de compactação aos sistemas existentes exigirá um trabalho de engenharia dedicado. No entanto, existem aplicações empresariais imediatas. “Acreditamos que a compactação após a ingestão é um caso de uso promissor, onde grandes saídas de chamadas de ferramentas ou documentos longos são compactados logo após serem processados”, disse Zweiger.
Em última análise, a mudança em direção à compactação mecânica do espaço latente se alinha com os futuros roteiros de produtos dos principais players de IA, argumenta Zweiger. “Estamos vendo a compactação passar de algo que as próprias empresas implementam para algo que os fornecedores modelo fornecem”, disse Zweiger. “Isso é ainda mais verdadeiro para compactação latente, onde é necessário acesso aos pesos do modelo. Por exemplo, o OpenAI agora expõe um endpoint de compactação de caixa preta que retorna um objeto opaco em vez de um resumo de texto simples.”
