Resumo Técnico: Seleção com Portão de Incerteza para Atenção de Esparsidade em Blocos

Declaração do Problema

Os mecanismos de atenção de esparsidade em blocos (ex: SSA, Quest, H2O, SnapKV) escalam modelos de linguagem de contexto longo substituindo o softmax $O(N^2)$ por uma seleção top-$k$ por consulta sobre blocos de chaves. Embora eficientes, esse processo de seleção é míope. Quando os escores do $k$-ésimo e do $(k+1)$-ésimo blocos estão quase empatados, o seletor assume uma decisão sem orçamento adicional. Se o bloco descartado contiver evidências críticas (ex: em raciocínio de múltiplos saltos ou recuperação needle-in-a-haystack), a informação é perdida irremediavelmente, pois as camadas subsequentes não podem recuperá-la. Esta falha é estrutural, não meramente ruidosa, e degrada o desempenho particularmente em cenários onde a relevância depende do contexto latente em vez de correspondências de superfície entre consulta e chave.

Métodos existentes tratam o corte como uma regra de decisão fixa. Eles não consideram a incerteza dessa decisão ou o valor da informação (VoI) que poderia ser ganho ao expandir o conjunto de atenção quando o corte é ambíguo.

Metodologia

O artigo propõe um roteador com portão de incerteza (uncertainty-gated router) que se posiciona sobre as espinhas dorsais (backbones) de pontuação de blocos existentes (como o mean-pooling do estilo SSA ou os limites superiores K-min/K-max do Quest) sem exigir retreinamento ou parâmetros adicionais. O método opera em cinco etapas:

1. Pontuação de Blocos (Independente de Backbone): As chaves são agrupadas em blocos. Os escores são computados usando um backbone existente (ex: produto interno de média agrupada ou limites superiores por coordenada).
2. Seleção por Bloco de Consulta (Per-Tile): Para reduzir o custo computacional, a seleção é realizada por Bloco de Consulta (Query Tile — um grupo de linhas de consulta consecutivas) em vez de por linha. O escore do bloco é o valor máximo entre as linhas do bloco. Uma seleção top-$k$ determina o conjunto inicial de blocos ($kv\_idx$).
3. Medição de Incerteza (Sinal de VoI): O método computa uma margem de corte normalizada ($\sigma$) para cada bloco de consulta e cabeça:
   $$\sigma = \frac{s_{(k-1)} - s_{(k)}}{s_{(0)} - s_{(k)}}$$
   onde $s_{(i)}$ são os escores dos blocos ordenados.
   • $\sigma \to 1$: O corte é decisivo (o conjunto mantido está muito acima da cauda).
   • $\sigma \to 0$: O corte é ambíguo (o $k$-ésimo e o $(k+1)$-ésimo blocos estão empatados).
     Este $\sigma$ serve como um proxy para a perda esperada da decisão top-$k$; um $\sigma$ baixo indica alto valor na aquisição de mais informação.
4. Agregação e Gatilho: Os valores de $\sigma$ por cabeça são calculados pela média para um sinal por bloco de consulta ($\bar{\sigma}$). Um gatilho é ativado se $\bar{\sigma}$ cair abaixo de um limiar de quantil empírico por camada ($q$). Isso identifica os blocos de consulta "arriscados" onde a decisão foi ambígua.
5. Expansão Controlada: Os blocos de consulta disparados expandem seu conjunto de atenção por um fator $\rho$ (ex: $\rho=2$), atendendo aos top-$\rho k$ blocos em vez de top-$k$. Blocos não disparados mantêm o orçamento base. A implementação utiliza um kernel fundido que preenche os blocos não disparados com sentinelas para manter formas de despacho uniformes, garantindo que o custo escale de forma sublinear com a expansão.

Principais Contribuições

1. Formulação de VoI: Uma formulação inovadora e agnóstica a backbone do corte do seletor como uma decisão sob incerteza. Adiciona apenas um escalar ($\sigma$) por bloco de consulta e um limiar de quantil por camada, não exigindo retreinamento ou parâmetros extras.
2. Componibilidade: Demonstração empírica de que o roteador se compõe com distintos backbones de pontuação (K-mean do estilo SSA e K-max do Quest). O roteador consistentemente eleva o desempenho de qualquer backbone que seja superior para um modelo específico, provando que "melhor pontuação" e "melhor alocação de orçamento" são melhorias ortogonais.
3. Implementação Fundida: Um kernel Triton customizado que lida com seleção e expansão em um único caminho de código, mantendo o despacho de forma uniforme. Isso permite que o sistema rode a $0,62\times$ a $0,80\times$ o tempo de execução da atenção densa em contextos de 128K.
4. Benchmarking Diagnóstico: Introdução do benchmark Pointer-Chase Haystack (PCH) para isolar a qualidade do seletor da capacidade do modelo, e o uso do LongBench-v1 como um controle negativo para demonstrar que os benefícios do roteador são condicionais à compressão do orçamento do seletor em relação à densidade de evidências.

Resultados

O método foi avaliado em quatro modelos (Qwen2.5-14B, Qwen2.5-7B-1M, Mistral-Nemo, Qwen3.6-35B) em contextos de 32K a 128K.

• Ganhos de Desempenho:
  • No LongBench-v2 medium (contexto de 32K), a configuração roteador-sobre-Quest alcançou um recall pareado de 0,75 comparado a 0,47 para o baseline top-$k$ (+28 pontos percentuais), com um McNemar $p < 0,01$.
  • No RULER NIAH multi-key, o roteador-sobre-Quest atingiu 0,98 de precisão (vs. 0,93 para top-$k$) no Qwen2.5-14B e 1,00 no Qwen2.5-7B-1M.
  • O roteador consistentemente melhorou o backbone vencedor para cada modelo. Notavelmente, o backbone ideal mudou conforme a arquitetura: o K-max do Quest venceu no Qwen2.5/Mistral, enquanto o K-mean do estilo SSA venceu no Qwen3.6 (devido ao QK-Norm regularizar a variância de magnitude). O roteador elevou ambos.
• Eficiência:
  • Em 128K de contexto, o roteador preservou 0,81 da precisão densa no Qwen2.5-7B-1M e 0,89 no Qwen3.6, rodando a $0,62\times$ e $0,80\times$ o tempo de execução da atenção densa, respectivamente.
  • A "vitória de Pareto" (melhor velocidade com qualidade comparável) emerge em contextos mais longos (64K–128K) à medida que a parcela de atenção do computo total aumenta, consistente com a lei de Amdahl.
• Orçamento vs. Seletividade: Estudos de ablação mostram que o ganho do roteador vem tanto do aumento do orçamento quanto da alocação seletiva. Em modelos próximos ao teto de desempenho (ex: Qwen-14B), o ganho é majoritariamente impulsionado pelo orçamento; em modelos com mais margem (ex: Mistral-Nemo), a seletividade (gastar orçamento apenas em blocos de consulta ambíguos) fornece a maior parte do ganho.

Significância e Alegações

O artigo afirma que o roteador com portão de incerteza fornece uma camada genérica e plug-and-play para a atenção de esparsidade em blocos que mitiga o modo de falha estrutural da seleção míope de top-$k$. Sua significância reside em:

• Robustez: Ele se adapta às características específicas de pontuação de diferentes arquiteturas de modelos (ex: lidando com o efeito QK-Norm no Qwen3.6) sem ajuste manual.
• Eficiência: Alcança precisão próxima à densa em 128K de contexto com tempo de execução significativamente reduzido, tornando a inferência de contexto longo mais prática.
• Fundamentação Teórica: Enquadra o problema de seleção através de uma lente de Valor da Informação (VoI), fornecendo um limite de arrependimento de amostra finita sob suposições de ruído Gaussiano (embora o artigo note que este limite é atualmente fraco nos regimes dominados por ruído onde o método opera).

Os autores afirmam explicitamente que o ganho do roteador é condicional: ele ativa apenas quando o orçamento por consulta é pequeno em relação à localização da evidência relevante para a resposta. Em cenários onde o orçamento é generoso (ex: LongBench-v1 com prompts curtos nativos), o roteador fornece benefício negligenciável, servindo como um controle negativo que valida a especificidade do método para cenários de orçamento apertado e contexto longo.