Resumo Técnico: Atenção Linear Kaczmarz (KLA)

Declaração do Problema

A modelagem de linguagem em contexto longo é prejudicada pelos custos computacionais e de memória quadráticos da atenção Softmax padrão em Transformers. Embora modelos recorrentes de tempo linear abordem isso comprimindo o contexto em um estado de tamanho fixo, o projeto da regra de atualização de estado torna-se o gargalo crítico. Abordagens existentes, como o Gated DeltaNet (GDN), utilizam uma "regra delta" para escrever o resíduo entre um valor alvo e a previsão atual do estado, editando assim associações existentes em vez de apenas acumulá-las.

No entanto, o GDN depende de um coeficiente escalar aprendido para equilibrar o esquecimento e a atualização. Os autores argumentam que este coeficiente é uma escolha de projeto empírica, e não um tamanho de passo derivado teoricamente. Um coeficiente inadequado pode levar a magnitudes de atualização subótimas: supercorreção para chaves de grande norma e subcorreção para chaves de pequena norma, resultando em corrupção do estado e desempenho de recuperação subótimo.

Metodologia: Atenção Linear Kaczmarz

O artigo propõe a Atenção Linear Kaczmarz (KLA), uma modificação do GDN que substitui o escalar aprendido por um tamanho de passo derivado teoricamente com base no método de projeção de Kaczmarz.

Derivação Teórica

Os autores enquadram a operação de escrita chave-valor como um problema de otimização com restrições. Para um token $t$, o objetivo é encontrar o estado $S_t$ mais próximo do estado anterior decaído $\tilde{S}_{t-1}$ que satisfaça a restrição linear $S_t^T k_t = v_t$ (onde $k_t$ é a chave e $v_t$ é o valor).

• Objetivo de Otimização: $\min_S \frac{1}{2}\|S - \tilde{S}_{t-1}\|_F^2$ sujeito a $S^T k_t = v_t$.
• Solução: A solução ótima é a projeção ortogonal de $\tilde{S}_{t-1}$ no hiperplano de restrição. Isso produz um tamanho de passo (coeficiente) de $\beta_t = 1/\|k_t\|_2^2$.

A Regra de Atualização KLA

A KLA generaliza essa projeção exata em uma regra de atualização prática que inclui um portão de relaxação aprendível $\eta_t$ e um estabilizador numérico $\epsilon$:
$$\beta_t = \frac{\eta_t}{\|k_t\|_2^2 + \epsilon}$$
A atualização do estado torna-se:
$$S_t = \tilde{S}_{t-1} + \beta_t k_t e_t^T$$
onde $e_t = v_t - \tilde{S}_{t-1}^T k_t$ é o erro residual.

Propriedades Chave

O artigo estabelece que este coeficiente satisfaz três propriedades desejáveis simultaneamente (no caso não relaxado onde $\eta_t=1, \epsilon=0$):

1. Projeção Ortogonal: Realiza a mudança de norma mínima no estado para satisfazer a restrição.
2. Busca Linear Exata: É o minimizador exato da perda por token ao longo da direção de atualização.
3. Passo de Gradiente Normalizado: Atua como um passo de Descida de Gradiente Estocástico (SGD) normalizado.

Implementação

A KLA é projetada como uma substituição direta para o GDN. Ela preserva:

• A forma do estado ($S \in \mathbb{R}^{d_k \times d_v}$).
• Os mecanismos de portão (portões de decaimento e atualização).
• A estrutura de recorrência linear.
• O algoritmo de treinamento paralelo por blocos (chunkwise).

A única mudança estrutural é o cálculo da matriz de coeficientes diagonais $B = \text{Diag}(\beta_1, \dots, \beta_C)$ dentro do solucionador por blocos. Não são necessários novos kernels de hardware ou layouts de estado.

Contribuições Principais

1. Regra Delta Auto-Normalizada: Derivação do coeficiente KLA a partir da projeção de Kaczmarz, provando que é uma projeção ortogonal, minimizador de busca linear exata e passo de gradiente normalizado para a perda por token.
2. Algoritmo Recorrente Direto: Demonstração de que a KLA mantém a recorrência do GDN e o solucionador por blocos, exigindo apenas uma mudança na matriz de coeficientes diagonais.
3. Ganhos Empíricos: Validação de que a KLA alcança menor perplexidade e melhor rastreamento de estado do que o GDN e outras linhas de base de tempo linear, sem expansão arquitetônica.
4. Eficiência Preservada: Confirmação de que a KLA mantém complexidade de tempo/memória $O(n)$ e iguala a latência de preenchimento (prefill) do GDN, enquanto melhora a vazão de decodificação.

Resultados Experimentais

Os experimentos foram conduzidos na escala de 0,4 bilhão de parâmetros com um orçamento de 1 bilhão de tokens no conjunto de dados SlimPajama-672B.

Modelagem de Linguagem

• Perplexidade de Validação (PPL): A KLA alcançou 8,09 PPL, superando o GDN (8,50), DeltaNet (8,26), GLA (8,42), Longhorn (8,79) e Mamba2 (9,41).
• Extrapolação de Contexto Longo: Sem ajuste fino em contexto longo, a KLA manteve a menor perplexidade em comprimentos de contexto de 1K a 65K tokens, mostrando uma vantagem consistente de 3,7–4,1 PPL sobre o GDN.

Tarefas Sintéticas

• Agulha Única no Palheiro (S-NIAH): A KLA atingiu 100% de precisão em todos os comprimentos de contexto testados (1K–8K), enquanto o GDN caiu ligeiramente e o Mamba2 falhou significativamente.
• Recordação Associativa de Múltiplas Consultas (MQAR): A KLA melhorou a precisão de extrapolação de comprimento 8× em 7,03 pontos sobre o GDN (73,84% vs 66,81%).
• Pilha e Palíndromo: A KLA mostrou desempenho competitivo ou superior em tarefas de rastreamento de estado e inversão.

Eficiência

• Latência de Preenchimento (Prefill): A KLA igualou a latência de preenchimento do GDN (por exemplo, ~1220ms em contexto de 131K), confirmando que o cálculo do coeficiente de Kaczmarz não impacta o custo dominante do kernel por blocos.
• Vazão de Decodificação: Em contexto de 32K, a KLA alcançou 3,54 tokens/s, o que é 2,1× maior que o GDN (1,65 tokens/s) e 1,7× maior que o Mamba2.

Estudos de Ablação

Os autores abalaram a estratégia de normalização, o mecanismo de portão e a expansão do estado. Os resultados indicaram que a normalização pela norma da chave é crítica para a estabilidade; removê-la ou usar apenas um escalar aprendido levou a uma degradação significativa na estabilidade do pré-treinamento (PPL) e no desempenho em tarefas sintéticas.

Significado e Alegações

O artigo alega que o coeficiente de Kaczmarz normalizado pela norma da chave é um "eixo de projeto de primeira ordem" para modelos de sequência de regra delta. Ao tratar a operação de escrita como uma projeção com restrição, a KLA resolve o problema das magnitudes de atualização subótimas inerentes a escalares aprendidos empiricamente.

O significado reside no fato de que essas melhorias são alcançadas sem alterar o estado recorrente, o kernel de hardware ou a complexidade arquitetônica. Os autores postulam que o coeficiente em si, e não a capacidade do estado ou a estrutura de portões, é o fator primário que determina se uma escrita de regra delta resolve com sucesso seu problema de regressão local. Isso sugere que modelos futuros de tempo linear podem alcançar ganhos substanciais em precisão, extrapolação e eficiência de decodificação simplesmente adotando este tamanho de passo fundamentado teoricamente.