Resumo Técnico: Neural Attention Search Linear (NAtS-L)

Declaração do Problema

A complexidade computacional quadrática ($O(L^2)$) da autoatenção baseada em softmax padrão em Transformers tornou-se um gargalo significativo para cenários de contexto longo. Embora modelos de atenção linear ofereçam uma alternativa promissora ao comprimir valores passados de Chave-Valor (KV) em um estado oculto de tamanho fixo para alcançar complexidade linear ($O(L)$), eles frequentemente sofrem de expressividade limitada devido às restrições desse tamanho de estado fixo. Por outroى, os Transformers padrão preservam toda a informação de contexto KV, permitindo uma melhor recuperação de longo prazo, mas incorrendo em altos custos de memória e computação.

Abordagens híbridas existentes tentam mitigar isso intercalando camadas de atenção softmax e linear ou usando cronogramas fixos para alternar entre elas. No entanto, esses métodos frequentemente mantêm estruturas fixas que podem não capturar de forma flexível toda a informação necessária através de contextos variados, e sua eficiência permanece limitada pela presença de camadas softmax completas. Além disso, variações de atenção esparsa frequentemente dependem de heurísticas humanas pré-definidas, potencialmente ignorando correlações não cobertas por seus mapas esparsos.

Metodologia

O artigo propõe o Neural Attention Search Linear (NAtS-L), um framework que aplica operações de atenção linear e softmax dentro da mesma camada, porém em diferentes tokens. Em vez de uma arquitetura fixa, o NAtS-L determina automaticamente a operação de atenção ideal para cada bloco de tokens com base no contexto de entrada.

Mecanismos Principais

1. Hibridização ao Nível de Token: A sequência de entrada é dividida em blocos (chunks). Uma Camada de Pontuação de Atenção (consistindo em pooling de média seguido por uma camada linear) mapeia cada bloco para um conjunto de pontuações que representam diferentes tipos de operação (Atenção Softmax vs. Gated DeltaNet, uma variante de atenção linear). A operação com a maior pontuação é atribuída a esse bloco.
2. Espaço de Busca Unificado: O framework unifica operações de atenção softmax não lineares e de atenção linear, permitindo que sejam processadas dentro de uma única camada. Ele aprende a importância de cada bloco de tokens, aplicando atenção softmax a tokens com alto impacto de longo prazo (que requerem recuperação de longo prazo) e atenção linear a tokens com impacto de curto prazo.
3. Computação por Blocos (Chunk-wise): Para manter a eficiência de hardware, o NAtS-L utiliza uma abordagem paralela por blocos. Dentro de um bloco, tanto as correlações lineares quanto as de softmax são computadas. O modelo mantém estados ocultos para ambos os tipos de operação. Quando um novo bloco chega, o modelo realiza um pseudo-update de ambos os estados; no entanto, para blocos que não foram selecionados como a operação alvo, os estados ocultos são resetados ou revertidos ao início do bloco, garantindo que apenas os valores KV relevantes sejam retidos para o tipo de operação selecionado.
4. Fusão de Saída (Output Merging): As saídas das operações softmax e linear selecionadas são normalizadas individualmente (usando normalização RMS) e então combinadas via uma soma ponderada. Os pesos para essa combinação são gerados dinamicamente a partir das projeções de consulta ($Q$), garantindo que o modelo aprenda a equilibrar as contribuições de ambos os tipos de atenção.
5. Otimização Baseada em Gradiente: A seleção dos tipos de atenção é aprendida de ponta a ponta. Gradientes são computados para a camada de pontuação de atenção, permitindo que o modelo otimize conjuntamente a seleção da operação e os pesos do modelo sem a necessidade de perdas auxiliares ou modelos pré-treinados.

Contribuições Principais

1. Framework Unificado: Os autores fornecem uma descrição unificada para atenção softmax não linear e atenção linear, permitindo que sejam processadas e selecionadas dentro de uma única camada.
2. Arquitetura Adaptativa: O NAtS-L é proposto como um framework que determina automaticamente a operação de atenção ideal para a informação de contexto de entrada em uma base de nível de token (nível de bloco), em vez de depender de razões fixas por camada.
3. Eficiência e Expressividade: O artigo demonstra que o NAtS-L alcança um forte equilíbrio entre desempenho e eficiência, oferecendo melhores capacidades de modelagem de contexto longo do que modelos puramente lineares, enquanto mantém menor latência computacional do que Transformers de softmax total.

Resultados Experimentais

Os autores avaliaram o NAtS-L em tarefas de pré-treinamento de escala acadêmica usando o dataset Fineweb-Edu com tamanhos de modelo de 380M e 800M de parâmetros.

• Modelagem de Linguagem e Raciocínio: Em benchmarks de raciocínio de senso comum zero-shot (ex: LAMBADA, PIQA, Hellaswag), o NAtS-L e sua variante híbrida tiveram desempenho comparável ou ligeiramente superior ao Gated DeltaNet (GDN) puro, Mamba2 e baselines de Transformer.
• Tarefas de Recuperação (Retrieval): Em tarefas de recuperação do mundo real (ex: SWDE, SQD, DROP), o NAtS-L Hybrid alcançou as melhores pontuações médias, superando modelos puramente lineares (que tiveram dificuldades com tarefas específicas como FDA) e Transformers puros (que tiveram dificuldades com SQD). Isso destaca a robustez do modelo em preservar a informação de contexto necessária.
• Extrapolação de Comprimento: Avaliado em contextos de até 65.536 tokens (PG19, NarrativeQA, CodeParrot), o NAtS-L Hybrid manteve níveis de perplexidade mesmo quando o contexto de entrada excedeu o comprimento de treinamento (4096 tokens), enquanto os híbridos puramente lineares (GDN Hybrid) falharam em extrapolar efetivamente.
• Benchmarks de Contexto Longo: No RULER e LongBench, o NAtS-L Hybrid alcançou desempenho de estado da arte entre os modelos comparados, particularmente em comprimentos de contexto de 8k e 16k.
• Latência de Inferência: O NAtS-L Hybrid mostrou uma aceleração de 5.4x no pre-filling em comparação aos Transformers e uma aceleração de 2.3x no decoding em um contexto de 128k, sendo apenas 1.66x mais lento que os modelos GDN puros.

Significância e Alegações

O artigo afirma que o NAtS-L fornece uma arquitetura híbrida flexível que atribui adaptativamente os tokens para atenção linear ou softmax com base em suas necessidades informacionais específicas. Ao aprender essa distribuição de ponta a end-to-end, o modelo evita a rigidez de estruturas híbridas fixas e a ineficiência da atenção softmax total.

Os autores enfatizam que esta abordagem permite ao modelo:

• Preservar informações relacionadas ao contexto de longo prazo via atenção softmax onde necessário.
• Reduzir os custos computacionais totais via atenção linear para tokens com impacto de curto prazo.
• Alcançar um melhor compromisso entre precisão e eficiência sem exigir modelos pré-treinados ou restrições auxiliares sobre a proporção dos tipos de atenção.

O trabalho sugere que uma mistura de nível de token entre atenção linear e softmax é um caminho viável para modelos de linguagem de contexto longo escaláveis que não comprometam suas capacidades de recuperação ou velocidade de inferência.