Linear Attention Architectures: Mechanisms, Trade-offs, and Cross-Layer Routing

Resumo Técnico: Arquiteturas de Atenção Linear: Mecanismos, Trade-offs e Roteamento entre Camadas

Definição do Problema
A autoatenção softmax padrão permite que cada token recupere informações de todo o contexto, um mecanismo central para a expressividade dos modelos de linguagem Transformer. No entanto, essa troca de token para token incorre em um custo quadrático em relação ao comprimento da sequência ($O(T^2)$) tanto para o treinamento quanto para a inferência, tornando-se um gargalo dominante à medida que as janelas de contexto se expandem. Enquanto as arquiteturas de atenção linear oferecem uma solução ao reformular a atenção como uma atualização recorrente sobre uma matriz de memória de tamanho constante (resultando em tempo de treinamento linear e tempo de inferência constante por token), as primeiras variantes sofreram perdas significativas de precisão devido à interferência entre associações armazenadas e à falta de mecanismos para esquecimento ou edição seletiva. Desenvolvimentos recentes, incluindo o DeltaNet e seus sucessores, reduziram essa lacuna ao introduzir regras delta de correção de erro e portões de controle granulares, contudo, os trade-offs específicos entre esses mecanismos, seus comportamentos de escala e o potencial para compartilhamento de informação leve entre camadas ainda precisam ser sistematicamente mapeados.

Metodologia e Estrutura
Os autores apresentam um estudo comparativo da atenção softmax e de quatro arquiteturas recentes de atenção linear recorrente: DeltaNet, Gated DeltaNet, Kimi Delta Attention e Gated DeltaNet-2.

1. Notação Unificada: O artigo estabelece uma notação recorrente de memória comum para expressar esses mecanismos. Este framework isola os papéis da matriz de memória recorrente $W$, o resíduo da regra delta $r$ (o erro entre o valor atual e a previsão da memória) e os vários mecanismos de controle de decaimento, apagamento e escrita.

   • DeltaNet: Substitui as atualizações aditivas ingênuas por uma regra delta de correção de erro, escrevendo apenas o resíduo $r = v - \bar{v}$ na memória.
   • Gated DeltaNet: Adiciona um portão de decaimento escalar aprendido $\alpha$ ao estado da memória antes de computar o resíduo, introduzindo um mecanismo de esquecimento grosseiro.
   • Kimi Delta Attention: Refina o esquecimento ao usar um vetor de decaimento por canal $\alpha$, permitindo que diferentes dimensões de características decaiam em taxas distintas.
   • Gated DeltaNet-2: Desacopla ainda mais a atualização ao introduzir portões de apagamento e de escrita separados por canal, permitindo o controle independente sobre quais canais de chave são removidos e quais canais de valor são armazenados.

2. Configuração Experimental: O estudo utiliza modelos de 350 milhões de parâmetros tipo decoder-only treinados em 15 bilhões de tokens (FineWeb-Edu). Os experimentos comparam a perda de validação, o throughput de treinamento, a sensibilidade do otimizador (AdamW vs. Muon) e o escalonamento de comprimento de sequência. Os autores também realizam execuções de maior escala (1,3B e 3B de parâmetros) e avaliam tarefas de downstream (HellaSwag, PIQA, WinoGrande).

3. Roteamento entre Camadas (Cross-Layer Routing): Para abordar a diluição de informação em pilhas profundas, o artigo introduz e avalia mecanismos de roteamento leves. Ele explora dois eixos: o que rotear (erro de escrita vs. valor de escrita) e onde injetá-lo (o alvo de valor da próxima camada vs. o fluxo oculto compartilhado).

   • CLER (Cross-Layer Error Residuals): Encaminha o erro de escrita da regra delta para o alvo de valor da próxima camada.
   • CLER-H: Roteia o erro de escrita para o fluxo oculto compartilhado.
   • CLVR (Cross-Layer Value Routing): Roteia o valor de escrita da camada (o alvo antes da subtração) para o fluxo oculto compartilhado via uma projeção inicializada com zero.

Principais Resultados

• Trade-offs de Arquitetura e Otimizador: Dentro do sweep de 350M de parâmetros e 15B de tokens, a Kimi Delta Attention treinada com o otimizador Muon em uma pilha híbrida (intercalada com camadas softmax) alcança a menor perda de validação final (2,273). No entanto, pilhas puras de Gated DeltaNet treinadas com AdamW demonstram o maior throughput de treinamento normalizado (definido como 100% da baseline). Pilhas híbridas geralmente melhoram a perda ao custo de algum throughput, enquanto pilhas puras preservam o escalonamento de tempo de iteração mais forte para contextos longos.
• Sensibilidade do Otimizador: A escolha do otimizador é inseparável do ajuste da taxa de aprendizado. O Muon prefere taxas de aprendizado mais baixas ($3 \times 10^{-4}$) para variantes de atenção linear, enquanto as variantes AdamW preferem taxas mais altas ($10^{-3}$). O Muon consistentemente reduz a perda de validação final em relação ao AdamW em configurações de arquitetura correspondentes.
• Comportamento de Escalonamento: Pilhas recorrentes puras exibem um escalonamento superior com o comprimento da sequência. A 32k tokens, a Gated DeltaNet pura mostra um aumento mínimo no tempo de iteração (~1,1x de 4k), enquanto a atenção softmax aumenta em ~2,9x. Pilhas híbridas ficam entre esses dois extremos, com aumentos no tempo de iteração aproximadamente proporcionais à densidade de suas camadas softmax.
• Descobertas de Roteamento entre Camadas:
  • CLER (Erro para o Alvo de Valor): Encaminhar o erro de escrita da regra delta para o alvo de valor da próxima camada não melhora o desempenho em relação às baselines correspondentes. Os autores atribuem isso a um "desajuste de espaço", onde o resíduo roteado compete com a própria geometria de valor aprendida pelo receptor.
  • CLVR (Valor para o Fluxo Oculto): Rotear o valor de escrita para o fluxo oculto compartilhado produz uma melhoria modesta, mas consistente, na perda de validação final tanto para DeltaNet quanto para Gated DeltaNet. Isso sugere que o sinal útil entre camadas é o alvo de escrita da camada, não a correção da regra delta.
  • Retornos Decrescentes: Os ganhos do CLVR são pequenos (reduzindo a perda em ~0,01 em 350M/1B de tokens) e diminuem à medida que a escala do modelo ou a contagem de tokens de treinamento aumenta, sugerindo que o benefício é mais pronunciado quando a memória base ainda não capturou toda a informação necessária.

Significância e Alegações
O artigo não declara uma única "melhor" arquitetura, mas visa tornar o espaço de design legível. Suas principais contribuições são:

1. Um Framework Unificado: Fornecer uma notação comum que esclarece como as recentes variantes de atenção linear diferem em expressividade, decaimento de memória e granularidade de controle.
2. Mapeamento de Trade-offs Empíricos: Demonstrar que a escolha entre pilhas lineares ou híbridas, e entre otimizadores, envolve um trade-off multiobjetivo entre perda de validação, throughput e escalonamento de comprimento de sequência.
3. Insights de Roteamento: Identificar que, embora o roteamento de erro ingênuo (CLER) falhe devido ao desalinhamento de base, rotear o valor de escrita para o fluxo oculto alinhado (CLVR) oferece um caminho leve, de tempo linear, para o compartilhamento de informação entre camadas que melhora o desempenho sem quebrar a estrutura recorrente.

Os autores enfatizam que suas conclusões baseiam-se em comparações de execução única e orçamentos de tokens específicos, instando cautela ao generalizar pequenas diferenças de perda de validação sem verificação de múltiplas sementes. O trabalho serve como uma fundação para entender os mecanismos específicos que impulsionam o desempenho na atenção linear e destaca questões abertas sobre a generalização de estratégias de roteamento para arquiteturas mais complexas, como a Gated DeltaNet-2.