Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scale

Resumo Técnico: Modelos de Linguagem Conseguem Realmente Recuperar em Contexto?

Definição do Problema
Embora os modelos de linguagem (LMs) ofereçam uma alternativa potencial à recuperação baseada em vetores ao condicionar-se a um corpus em contexto para gerar diretamente identificadores de documentos relevantes, essa capacidade permanece amplamente inexplorada em escalas realistas. Trabalhos anteriores focaram em sistemas proprietários ou tarefas de reranking de pequena escala, deixando dois desafios críticos não abordados: (C1) escalar para corpora contendo milhões de tokens, e (C2) generalizar para tamanhos de corpus muito além daqueles vistos durante o treinamento. Os LMs atuais lutam tanto com os custos de atenção de um milhão de tokens quanto com a fragilidade da extrapolação para contextos mais longos. A questão central é se um LM pode atuar de forma confiável como um recuperador de escala de corpus sob essas condições sem colapsar.

Metodologia: BLOCKSEARCH e Análise
Os autores introduzem o BLOCKSEARCH, um estudo sistemático e uma arquitetura de modelo projetada para lidar com a recuperação em contexto (ICR) em escala de corpus.

• Arquitetura e Treinamento: Construído sobre um backbone de 0,6B de parâmetros (Qwen3-0.6B), o BLOCKSEARCH emprega atenção esparsa por blocos para lidar com documentos disjuntos de forma eficiente. As principais modificações arquiteturais e de treinamento incluem:

  • Identificadores de Documentos Aleatórios: Em vez de IDs sequenciais, os documentos recebem códigos de quatro dígitos aleatórios por etapa de treinamento para evitar o overfitting à posição absoluta.
  • Negativos no Lote (In-Batch Negatives) e Perda On-Policy: O modelo é treinado com 1 positivo e 15 negativos difíceis por consulta. Para mitigar o viés de exposição (onde o modelo se condiciona aos seus próprios prefixos gerados potencialmente incorretos durante a inferência), uma perda auxiliar on-policy é adicionada. Isso envolve realizar o rollout de um código a partir da própria distribuição do modelo e treinar o modelo para corrigir seu caminho com base nos scores do lote.
  • Configuração de Avaliação: O modelo é avaliado em datasets BEIR (MS MARCO, NQ, HotpotQA) com tamanhos de corpus variando de 500 a 10.000 documentos (até ~1 milhão de tokens), excedendo em muito a janela de contexto nativa de 32k tokens do modelo.

• Análise Mecanística de Falha: O estudo identifica que, embora o BLOCKSEARCH melhore a extrapolação, ele ainda colapsa em grandes escalas (recall próximo de zero em 1M de tokens). Os autores rastreiam isso não como uma falha em identificar o documento relevante (ranking), mas como um efeito de diluição de atenção:

  • Lacuna entre Recuperação e Geração: Mesmo com 10.000 documentos, os cabeças de atenção individuais ainda costumam atribuir o maior score pré-softmax ao documento de ouro ($R_{any} \approx 1.0$).
  • Colapso do Softmax: No entanto, conforme o corpus ($N$) cresce, a contribuição agregada dos documentos irrelevantes para o denominador do softmax aumenta mais rápido do que o score do documento de ouro. Consequentemente, a massa de atenção normalizada (peso pós-softmax) no documento de ouro colapsa, fazendo com que a saída da camada se torne uma média de distratores em vez do conteúdo relevante.

• Mitigações Propostas: Para combater a diluição de atenção, os autores propõem dois ajustes sensíveis ao comprimento:

  1. Ajustes de Softmax Sensíveis ao Comprimento:
     • Sinks Aditivos: Anexar uma constante aprendida ao denominador do softmax para absorver a massa de atenção difusa (inspirado no Streaming-LLM).
     • Redimensionamento Multiplicativo (SSMax): Escalar os scores pré-softmax por $s \cdot \log N$ para ampliar a lacuna entre o score de ouro e os distratores, evitando que o peso de ouro desapareça conforme $N$ aumenta.
  2. Atenção Esparsa ao Nível de Documento: Um mecanismo de roteamento em uma camada intermediária (L16) que pontua todos os documentos e restringe a atenção densa aos top-$B$ (por exemplo, 256) documentos para as camadas subsequentes, efetivamente podando o contexto antes da geração final.

Principais Resultados

• Desempenho da Linha de Base (Baseline): O modelo BLOCKSEARCH inicial iguala o desempenho de recuperação densa em benchmarks padrão (MS MARCO, NQ) em tamanhos de corpus menores (~45k tokens), e generaliza até ~500k tokens (10x o contexto de treinamento), mas colapsa próximo a 1M de tokens.
• Impacto das Modificações:
  • SSMax: Melhora significativamente o desempenho através de todo o sweep de $N$. No MS MARCO com 1M de tokens, atinge 16,5% de Recall@1 (vs. 0,2% para o modelo não modificado), fechando a maior parte da lacuna para a linha de base densa (20,2%).
  • Roteamento (Routing): O roteamento por documento melhora ainda mais o desempenho, alcançando 18,8% de Recall@1 no MS MARCO com 1M de tokens, quase igualando a linha de base densa.
  • Abordagem Combinada: A combinação de SSMax e roteamento (BLOCKSEARCH-SSMax-routing) atinge 20,5% de Recall@1 no MS MARCO com 1M de tokens, superando ligeiramente a linha de base densa (20,2%) e superando o modelo concorrente MSA-4B, 7x maior (16,0%).
• Fora de Distribuição (LIMIT): No benchmark LIMIT, que exige similaridade lexical em vez de similaridade de embedding semântico (onde a recuperação densa tem dificuldades), o BLOCKSEARCH modificado supera significativamente a linha de base densa. Em 850k tokens, o BLOCKSEARCH-SSMax-routing atinge 14,9% de Recall@1 comparado aos 3,5% do recuperador denso (uma melhoria de ~3x).
• Teto de Atenção (Attention Ceiling): Em todos os experimentos, o "teto de atenção" (a fração de consultas onde pelo menos um cabeça classifica o documento de ouro em primeiro lugar) permanece próximo de 100%, confirmando que o sinal de recuperação está presente internamente, mas é perdido durante a fase de leitura/geração devido à diluição da atenção.

Significância e Alegações
O artigo afirma fornecer o primeiro estudo sistemático de recuperação em contexto na escala de um milhão de tokens, estabelecendo-a como uma alternativa viável à recuperação densa clássica sob condições específicas.

• Viabilidade: Com ajustes arquiteturais (SSMax e roteamento), os LMs podem igualar ou exceder a recuperação densa em benchmarks padrão e superar significativamente a densa em tarefas que exigem similaridade não semântica.
• Eficiência: O proposto modelo de 0,6B alcança esses resultados com 7x menos parâmetros do que modelos de contexto longo concorrentes (MSA-4B) e sem a necessidade de retreinamento quando o corpus muda (ao contrário da recuperação generativa).
• Gargalo Fundamental: O trabalho identifica a diluição de atenção (dispersão do softmax) como o principal gargalo para a recuperação em contexto escalável. Embora intervenções simples possam mitigar isso, o artigo enfatiza que controlar a atenção sob crescimento extremo de contexto continua sendo um desafio fundamental para futuras arquiteturas de recuperação.

Os autores mantêm a modéstia, observando que, embora seus métodos estendam a recuperação funcional para corpora maiores, eles não "resolvem" totalmente o problema da recuperação em escalas extremas, já que o desempenho ainda declina conforme $N$ aumenta, e uma lacuna residual em relação à recuperação densa permanece em alguns datasets.