Forget, Then Recall: Learnable Compression and Selective Unfolding via Gist Sparse Attention

1. O Problema

O escalonamento de Modelos de Linguagem Grandes (LLMs) para contextos longos enfrenta um gargalo fundamental: o custo computacional quadrático ($O(n^2)$) do mecanismo de atenção padrão. À medida que o comprimento do contexto cresce (de milhares para milhões de tokens), o cálculo de atenção torna-se proibitivo para treinamento e inferência.

As abordagens existentes para mitigar isso geralmente caem em duas categorias com limitações:

• Seleção/Compressão de KV-cache (Inferência): Métodos como H2O ou StreamingLLM evitam tokens menos importantes durante a decodificação, mas operam em modelos pré-treinados com atenção completa, limitando a capacidade de aprender padrões esparsos ótimos do zero.
• Atenção Esparsa (Treinamento): Métodos recentes (NSA, DSA, MoBA) introduzem esparsidade durante o pré-treinamento, mas frequentemente exigem modificações na arquitetura do modelo, módulos externos não diferenciáveis (indexadores) ou operações de agrupamento não diferenciáveis (mean-pooling), o que impede a otimização end-to-end.

Além disso, em cenários de Geração Aumentada por Recuperação (RAG), a atenção padrão trata todos os documentos recuperados como igualmente relevantes, dispersando a capacidade do modelo em informações irrelevantes (ruído), em vez de focar no documento correto.

2. Metodologia: Gist Sparse Attention (GSA)

O artigo propõe o Gist Sparse Attention (GSA), uma estrutura unificada que conecta compressão baseada em gist (resumo) e atenção esparsa através de um mecanismo de desdobramento seletivo. A ideia central é que tokens de gist intercalados podem servir como sinais de roteamento aprendíveis para selecionar quais partes do contexto devem ser recuperadas em detalhe.

Componentes Principais:

1. Tokens de Gist Intercalados:

   • O contexto é dividido em blocos (chunks) de tamanho $L$.
   • Um token de gist aprendível ($g_m$) é inserido após cada bloco.
   • Uma máscara causal especial força o token de gist a aprender uma representação comprimida e densa de todo o seu bloco correspondente, atuando como um "resumo" daquela seção.

2. Roteamento e Pontuação de Relevância:

   • Em vez de usar pooling médio não diferenciável (como no MoBA) ou indexadores externos (como no DSA), o GSA usa os próprios tokens de gist para calcular a relevância.
   • Para uma consulta ($q_t$), calcula-se a pontuação de atenção entre a consulta e as chaves dos tokens de gist ($k_{g_m}$). Isso fornece uma estimativa direta e aprendível da relevância de cada bloco de contexto.

3. Desdobramento Seletivo (Selective Unfolding):

   • Com base nas pontuações, o modelo seleciona os $k$ blocos mais relevantes (Top-k).
   • Mecanismo Híbrido: Para os blocos selecionados, o modelo "desdobra" a compressão: reintroduz os tokens brutos (raw tokens) daquele bloco no contexto de atenção, juntamente com o token de gist.
   • Para os blocos não selecionados, apenas o token de gist (ou nenhum token, dependendo da configuração) permanece, mantendo o custo computacional baixo.
   • Isso cria um mecanismo de finais a grossos (coarse-to-fine): o modelo navega por resumos globais e recupera detalhes finos apenas onde é necessário.

4. Extensão Hierárquica (H-GSA):

   • Para atingir complexidade logarítmica ($O(\log n)$), o GSA pode ser estendido recursivamente.
   • Tokens de meta-gist são criados para resumir grupos de tokens de gist.
   • O modelo realiza uma seleção hierárquica: primeiro seleciona os segmentos mais relevantes no nível meta-gist, depois os sub-blocos no nível gist, e finalmente os tokens brutos. Isso permite acesso a contextos ilimitados com custo de decodificação logarítmico.

5. Treinamento End-to-End:

   • Pré-treinamento Contínuo: O modelo aprende a comprimir o contexto em tokens de gist e a prever o próximo token usando apenas esses resumos.
   • Ajuste Fino Seletivo (Opcional): O mecanismo de desdobramento seletivo é incorporado diretamente no treinamento. O modelo aprende a ativar o desdobramento de tokens brutos apenas para os blocos relevantes, otimizando a seleção via gradiente.

3. Contribuições Chave

• Pivôs Aprendíveis: A observação de que tokens de gist intercalados podem atuar como sinais de roteamento eficientes para atenção esparsa, eliminando a necessidade de módulos externos ou modificações arquiteturais.
• Framework Unificado: Uma abordagem que integra compressão e atenção esparsa de forma end-to-end, sem alterar a arquitetura do Transformer padrão e sem usar operações não diferenciáveis.
• Desdobramento Seletivo: Um mecanismo que supera a perda de informação típica da compressão, recuperando detalhes finos apenas quando necessário, melhorando a precisão sob orçamentos de tokens rigorosos.
• Complexidade Logarítmica: A extensão hierárquica (H-GSA) permite escalar para contextos extremamente longos com complexidade de decodificação $O(\log n)$.

4. Resultados Experimentais

Os autores avaliaram o GSA em dois modelos (Qwen2-7B e Llama3.2-1B) em benchmarks de LongBench e RAG (Geração Aumentada por Recuperação).

• Desempenho em LongBench:

  • O GSA superou consistentemente métodos de compressão (como ActivationBeacon, UniGist) e métodos de atenção esparsa apenas na inferência.
  • Em taxas de compressão de 8× a 32×, o GSA obteve ganhos significativos. Por exemplo, no pré-treinamento contínuo com 8× de compressão, o GSA alcançou 46.20 de média, superando o ActivationBeacon (42.52) e quase igualando o modelo de atenção completa (47.78).
  • O ajuste fino seletivo melhorou ainda mais os resultados, com o GSA superando o baseline de atenção completa em algumas tarefas específicas (ex: MF-en), sugerindo que o mecanismo atua como um viés indutivo benéfico.

• Desempenho em RAG:

  • O GSA mostrou vantagens dramáticas em cenários de múltiplos documentos. No ajuste fino com 8× de compressão, alcançou 53.39, superando o KVLink (41.45) e o UniGist (45.27) em mais de 8-11 pontos.
  • A capacidade de filtrar documentos irrelevantes e focar apenas nos blocos relevantes do documento correto foi crucial para o sucesso.

• Reutilização de KV-cache:

  • Em tarefas onde documentos longos são pré-processados e reutilizados para várias consultas, o GSA demonstrou eficiência superior, permitindo roteamento eficiente entre documentos sem precisar acessar todos os tokens brutos.

• Análise de Ablação:

  • A seleção Top-k adaptativa superou consistentemente a seleção baseada em top-p (nucleus), oferecendo maior estabilidade e controle sobre o orçamento de tokens.
  • Manter os tokens de gist selecionados junto com os tokens brutos (SG+SR) foi a configuração mais eficaz, fornecendo informações complementares.

5. Significado e Conclusão

O GSA representa um avanço significativo na modelagem de contextos longos ao demonstrar que a compressão não precisa ser um processo irreversível de perda de informação. Ao transformar a compressão em um mecanismo de roteamento aprendível, o GSA permite que os modelos "esqueçam" o contexto irrelevante (comprimindo-o em gists) e "lembrem" (desdobrem) apenas os detalhes críticos quando necessário.

Impacto Principal:

• Eficiência: Reduz o custo computacional e de memória para contextos longos, permitindo escalabilidade para milhões de tokens.
• Simplicidade: Não requer mudanças na arquitetura do modelo ou módulos externos complexos, facilitando a integração com modelos existentes.
• Qualidade: Mantém ou melhora a qualidade da geração em comparação com a atenção completa, especialmente em tarefas que exigem recuperação precisa de informações em grandes volumes de dados (como RAG).

O código do projeto está disponível publicamente, promovendo a adoção e o desenvolvimento futuro de métodos de atenção esparsa aprendíveis.