💬 NLP

Learning how to Forget: Fine-tuning for Long-Context Sparse Attention

Este artigo apresenta um método de ajuste fino para modelos transformer que possibilita a inferência eficiente de contexto longo com atenção esparsa em hardware moderado, ao permitir que os modelos se coadaptem com várias políticas de cache KV, frequentemente superando abordagens de atenção exata.

Autores originais: Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Publicado 2026-08-21

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Matthias Seeger, Zeyu Zhang, Vihang Patil, Konstantinos Benidis, Sebastian Schelter

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Os sistemas modernos de inteligência artificial que geram texto semelhante ao humano dependem de um mecanismo que atua como uma memória de curto prazo, permitindo que eles se lembrem do que foi dito anteriormente em uma conversa ou em um documento longo. Essa memória é armazenada em um buffer digital que cresce a cada nova palavra que o sistema processa. Para que esses sistemas funcionem bem ao longo de extensos trechos de texto, essa memória precisa ser vasta, mas o hardware de computador necessário para contê-la é caro e limitado. Quando o buffer de memória enche, o sistema deve decidir quais partes antigas de informação descartar para abrir espaço para a nova. Se ele jogar fora a informação errada, o sistema perde sua capacidade de raciocinar ou responder perguntas com precisão. Isso cria um difícil equilíbrio: manter a memória pequena economiza dinheiro e permite que o sistema rode em equipamentos padrão, mas corre o risco de perder o contexto necessário para ser inteligente.

Pesquisadores há muito tempo tentam resolver isso ensinando o sistema a ser seletivo sobre o que manter, um processo conhecido como atenção esparsa. No entanto, um novo estudo revela uma falha crítica na forma como esses sistemas têm sido treinados até agora. A maioria dos métodos existentes treina a IA usando uma memória perfeita e ilimitada e depois tenta forçá-la a operar com uma memória limitada mais tarde. Os pesquisadores descobriram que essa abordagem falha porque a IA nunca aprendeu como funcionar sob as restrições específicas que realmente enfrentaria. Ao treinar o modelo para esquecer intencionalmente e se adaptar a um tamanho de memória fixo desde o início, a equipe demonstrou que o sistema poderia performar significativamente melhor do que aqueles treinados com recursos ilimitados, mesmo rodando em um único chip de computador moderadamente potente.

A equipe, liderada por cientistas da Amazon Web Services e da Universidade de Amsterdã, desenvolveu uma nova maneira de ajustar esses grandes modelos de linguagem. Em vez de usar supercomputadores massivos para simular uma memória perfeita, eles ensinaram os modelos a coadaptar-se com uma política específica de gerenciamento de memória. Imagine um bibliotecário que é treinado para organizar livros em uma biblioteca com prateleiras infinitas, apenas para ser informado mais tarde que deve trabalhar em uma sala minúscula com uma única prateleira. O bibliotecário treinado na biblioteca grande provavelmente teria dificuldade em priorizar o que manter na sala pequena. Em contraste, o método proposto neste artigo treina o bibliotecário diretamente na sala pequena, ensinando-o exatamente quais livros manter e quais descartar com base nas regras daquele espaço específico. Isso permite que o modelo aprenda o ritmo de suas próprias limitações, em vez de tentar desaprender os hábitos de ter muito espaço.

Os pesquisadores testaram essa abordagem em um modelo com quatro bilhões de parâmetros, um tamanho substancial, mas gerenciável. Eles realizaram seus experimentos em uma única placa gráfica com 40 gigabytes de memória, uma configuração acessível para muitas organizações comparada aos clusters de dezenas de placas exigidos pelos métodos anteriores. Eles compararam essa nova técnica de treinamento contra o método padrão, que utiliza uma técnica chamada paralelismo de sequência para dividir a carga de memória entre múltiplos dispositivos caros. Os resultados mostraram que os modelos treinados com o novo método frequentemente superaram os modelos padrão, particularmente quando a tarefa exigia que o sistema gerasse respostas específicas e concisas, em vez de textos longos e prolixos. Em vários testes envolvendo perguntas complexas e extração de dados, o método padrão produziu saídas que eram excessivamente longas e cheias de números aleatórios e sem sentido, enquanto o novo método aprendeu a parar no momento certo e fornecer o valor único correto.

Uma parte fundamental desse sucesso foi a melhoria do "oráculo de heavy-hitter", uma estratégia popular para decidir qual informação manter. Essa estratégia funciona rastreando em quais partes de informação o modelo presta mais atenção ao longo do tempo. Os pesquisadores descobriram que a versão original dessa estratégia era lenta e ineficiente. Eles reescreveram o código subjacente para funcionar muito mais rápido, permitindo que o sistema calculasse essas pontuações de importância sem atrasar todo o processo. Essa otimização significou que o sistema poderia tomar decisões inteligentes sobre o que esquecer em tempo real, sem precisar do enorme poder computacional que geralmente acompanha tais tarefas. A equipe também lançou uma nova biblioteca de software de código aberto para disponibilizar essas técnicas a outros, visando baixar a barreira para qualquer pessoa que deseje construir sistemas de IA de contexto longo sem precisar de uma fortuna em hardware.

O estudo destaca que a maneira como um modelo é treinado é tão importante quanto o hardware em que ele roda. Quando os pesquisadores forçaram o modelo a treinar com as exatas mesmas restrições de memória que ele enfrentaria durante o uso, ele aprendeu a navegar por essas restrições de forma eficaz. Em um teste específico envolvendo dados JSON, o método padrão falhou completamente, sendo incapaz de encontrar os pontos de dados corretos, enquanto o novo método teve sucesso em identificá-los cerca de metade das vezes. Isso sugere que a capacidade de lidar com contextos longos não é apenas uma questão de ter mais memória, mas de ensinar o sistema a gerenciar a memória que ele possui. As descobertas indicam que, para muitas aplicações, o caminho mais eficaz à frente não é construir computadores maiores, mas ensinar o software a ser mais eficiente com os recursos que ele já possui.

Resumo Técnico: Aprendendo a Esquecer: Fine-tuning para Atenção Esparsa de Longo Contexto

1. Definição do Problema

Os modelos de linguagem de grande escala (LLMs) modernos requerem a capacidade de processar contextos extremamente longos para tarefas como conversas de múltiplos turnos, uso de ferramentas e raciocínio de cadeia de pensamento (chain-of-thought). Embora as implementações padrão de transformer escalem quadraticamente em computação e linearmente em memória com a largura do contexto, métodos de atenção esparsa foram propostos para alcançar escalonamento de tempo linear e memória constante. Esses métodos utilizam um cache de Chave-Valor (KV) de tamanho fixo onde a informação é descartada com base em políticas específicas (ex: Heavy-Hitter Oracle, H2O) uma vez que o cache fica cheio.

No entanto, existe uma lacuna significativa entre a inferência e o treinamento neste domínio:

Restrições de Treinamento: O ajuste fino (fine-tuning) de modelos com políticas de atenção esparsa é notoriamente difícil em orçamentos de hardware moderados. A diferenciação automática padrão (autograd) exige o armazenamento de todo o grafo de computação, incluindo buffers de cache KV para cada bloco da sequência. Para contextos longos, esse requisito de memória excede a memória disponível na GPU, muitas vezes necessitando de configurações multi-GPU caras com paralelismo de sequência ou de contexto (SP/CP).
Lacuna de Coadaptação: Modelos treinados com atenção exata (usando SP/CP) frequentemente apresentam desempenho significativamente inferior quando implantados com atenção esparsa na inferência. Isso ocorre porque o modelo nunca é exposto às restrições de memória específicas e à lógica de descarte que enfrentará durante a inferência.
Ineficiência de Implementação: Políticas de atenção esparsa líderes, como a H2O, são subutilizadas em bibliotecas de inferência de estado da arte (SotA) (ex: vLLM) porque as implementações atuais são lentas, carecendo de kernels otimizados para os requisitos específicos de atenção esparsa (ex: retornar pesos de atenção somados).

2. Metodologia

O artigo propõe um novo método de fine-tuning que permite que modelos transformer se coadaptem com políticas arbitrárias de cache KV em hardware moderado (demonstrado em uma única Nvidia A100 40GB).

2.1 Checkpointing de Ativação Aninhado com Offloading para CPU

Para abordar o gargalo de memória ao treinar em sequências longas ( $N \gg N_C$ , onde $N_C$ é o tamanho do cache), os autores empregam uma estratégia de checkpointing de ativação aninhada:

Loop Externo (Camadas): A passagem para trás (backward pass) itera sobre as camadas do modelo.
Loop Interno (Células): A sequência é dividida em "células" (grupos de blocos). Para cada célula, uma passagem direta (forward pass) é executada para armazenar os buffers de cache KV necessários na memória da CPU.
Mecanismo de Replay: Em vez de diferenciar através da política de cache KV (que é frequentemente não diferenciável ou dispendiosa), a passagem direta registra todas as decisões de descarte. Durante a passagem para trás, essas decisões são reproduzidas usando "caches de replay" que mimetizam a lógica de descarte do forward sem exigir o cálculo de gradiente através da própria política.

2.2 Codificação Delta para Buffers de Cache KV

A otimização de memória mais crítica envolve como os buffers de cache KV são armazenados dentro do grafo autograd.

Recorrência Linear: Os autores observam que o buffer de cache KV para um novo bloco ($keys'$) é uma transformação linear do buffer anterior ($keys$), diferindo apenas pelos valores dos novos tokens e pelos índices específicos de slots sobrescritos pela política de descarte.
Armazenamento Delta: Em vez de armazenar os tensores completos de $keys $e$ values$ para cada bloco no grafo de computação, o método armazena apenas o delta (a diferença) e os índices dos slots sobrescritos.
Hooks de Tensores Salvos de Autograd: Usando os autograd_saved_tensors_hooks do PyTorch, os autores implementam um mecanismo de "empacotamento" (packing). Durante a passagem direta, os tensores completos são substituídos por referências a tuplas (índice, delta). Durante a passagem para trás, os tensores completos são reconstruídos ("desempacotados") a partir dos deltas e dos buffers subsequentes. Isso reduz o footprint de memória do grafo autograd por um fator de $k$ (o número de blocos por célula), trazendo-o para níveis comparáveis à inferência.

2.3 Melhorias na Política H2O

O artigo introduz melhorias específicas na política Heavy-Hitter Oracle (H2O):

Independência de Batch: A implementação original da H2O agregava pontuações através das dimensões de batch. O novo método permite que as decisões de descarte sejam tomadas independentemente para cada posição de batch.
Normalização: Uma pontuação H2O normalizada é introduzida para evitar o viés em direção a tokens que estiveram no cache por mais tempo.
Otimização de Kernel: Para permitir treinamento e inferência eficientes, os autores fornecem código Triton que retorna pesos de atenção somados junto com um kernel de Scaled Dot Product Attention (SDPA) do FlashInfer. Isso permite que a H2O rode eficientemente sem recorrer a implementações SDPA lentas e ingênuas.

3. Principais Contribuições

Novo Método de Fine-Tuning: Um método para realizar o fine-tuning de modelos transformer com atenção esparsa e políticas arbitrárias de cache KV em recursos comparáveis à inferência de atenção esparsa (ex: uma única A100). Combina checkpointing de ativação aninhado, offloading para CPU e codificação delta via hooks de autograd.
Implementação e Otimização da H2O: Melhorias metodológicas na política H2O (independência de batch, normalização) e uma implementação de alto desempenho usando kernels SDPA dedicados que retornam pesos de atenção somados, dando um passo significativo em direção a latências competitivas com bibliotecas de inferência SotA.
Avaliação Abrangente: Experimentos no conjunto de benchmarks Helmet (janelas de contexto de 64k e 128k) demonstrando que, embora o desempenho varie por dataset, o método proposto permite que modelos se coadaptem com políticas de cache KV, frequentemente superando modelos treinados com paralelismo de sequência quando avaliados sob inferência de atenção esparsa.
Biblioteca KeysAndValues: O lançamento de uma biblioteca de código aberto (awslabs/keys_values) que fornece código performante e fácil de usar para inferência de longo contexto e fine-tuning com atenção esparsa.

4. Resultados Experimentais

Os autores avaliaram um modelo Qwen3-4B-Instruct no benchmark Helmet através de comprimentos de contexto de 64k e 128k.

Comparação: Os modelos foram treinados usando:
- SP: Paralelismo de Sequência com atenção exata (baseline SotA).
- US: O método proposto com a política de cache KV específica em vigor.
- NO: O checkpoint base sem fine-tuning.
Descobertas:
- Resultados Mistos nos Datasets da Tabela 1: Para os datasets da Tabela 1, os resultados foram mistos e inconclusivos. O baseline SP (atenção exata) alcançou os melhores resultados para nq e hotpot qa. Para trivia qa, o checkpoint base sem fine-tuning (no) performou melhor, indicando que o fine-tuning não ajudou neste caso específico. O método proposto (us) superou o SP em pop qa.
- Forte Desempenho nos Datasets da Tabela 2: Para os datasets da Tabela 2, o us superou fortemente tanto o sp quanto o no.
- Modos de Falha do SP: Modelos treinados com SP exibiram um modo de falha onde as saídas eram excessivamente longas e continham nonsense aleatório quando submetidos ao descarte de atenção esparsa, provavelmente porque nunca foram treinados para lidar com o mecanismo de "esquecimento". Em contraste, o us aprendeu a parar adequadamente e geralmente produz um único número ou o alvo correto.
- Variância de Política: Os resultados variaram significamente por dataset. Variantes da H2O performaram bem em várias tarefas, mas a política lastrec (lr) superou a H2O em datasets como nq e hotpot qa. Além disso, para nq e hotpot qa, o baseline SP (atenção exata) frequentemente alcançou os melhores resultados, enquanto para pop qa e outros datasets da Tabela 2, o método US com políticas de atenção esparsa foi superior.
- Eficiência: A implementação H2O com o novo suporte a kernel representa uma grande melhoria de eficiência, dando um grande passo em direção a latências competitivas com bibliotecas SotA, embora o artigo reconheça que a latência atual da atenção esparsa permanece significativamente maior do que métodos usando paralelismo de sequência ou de contexto.

5. Significância e Alegações

O artigo afirma que a atenção esparsa tem sido subutilizada na prática principalmente devido a problemas de latência e à dificuldade de realizar o fine-tuning de modelos para se adaptarem a políticas de descarte específicas.

Viabilidade: Os autores demonstram que o fine-tuning de alta qualidade para atenção esparsa de longo contexto é viável em hardware moderado, removendo a necessidade de clusters massivos de treinamento distribuído apenas para adaptar a uma política de cache.
Necessidade de Coadaptação: Os resultados sugerem fortemente que os modelos devem ser treinados sob as mesmas restrições (políticas de descarte) que enfrentarão durante a inferência para evitar a degradação de desempenho.
Direção Futura: O artigo argumenta que, ao simplificar a estrutura do cache KV e esclarecer os requisitos do kernel SDPA (ex: retornar pesos somados), a comunidade pode melhor integrar a atenção esparsa em bibliotecas de inferência SotA.
Limitações: Os autores reconhecem que, embora seu método melhore a latência, a natureza sequencial das decisões de atenção esparsa ainda pode resultar em maior latência em comparação com RingAttention ou Paralelismo de Sequência para orçamentos de hardware muito grandes. Eles planejam abordar isso via fusão de kernels e offloading assíncrono para CPU em trabalhos futuros.

Em resumo, o artigo fornece um caminho prático para "aprender a esquecer", permitindo que modelos se especializem para inferência de longo contexto em hardware acessível. Embora o desempenho relativo ao paralelismo de sequência varie por dataset, o método frequentemente supera o treinamento de atenção exata em cenários que exigem adaptação a políticas específicas de atenção esparsa.