Resumo Técnico: OSDN: Melhorando a Regra Delta com Pré-condicionamento Online Provável em Atenção Linear

1. Declaração do Problema

Atenção linear e modelos de espaço de estados (SSMs) oferecem complexidade de inferência $O(N)$ e alternativas de memória constante à atenção softmax, tornando-as adequadas para processamento de sequências longas. No entanto, elas frequentemente lutam com recordação associativa em contexto, especialmente quando comparadas a transformers baseados em softmax.

A Regra Delta (como implementada no DeltaNet) tenta preencher essa lacuna interpretando a atualização do estado recorrente como um único passo de descida de gradiente online (OGD) em uma perda de regressão por token: $f_t(S) = \frac{1}{2}\|S k_t - v_t\|_F^2$. A regra de atualização é $S_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)k_t^\top$, onde $\beta_t$ é um portão escalar.

A Limitação Central: O portão escalar $\beta_t$ atua como uma taxa de aprendizado uniforme aplicada a todas as dimensões da chave. Isso ignora a curvatura específica de características do objetivo interno. Em tarefas de recordação associativa, diferentes chaves (por exemplo, identificadores frequentes vs. raros, valores estáveis vs. de alta entropia) exibem curvaturas empíricas diferentes. Um único tamanho de passo escalar força um compromisso: aumentá-lo para ajudar a sobrescrever memória obsoleta em uma direção pode supercorrigir outra direção que já está bem calibrada. Otimizadores adaptativos padrão (como Adam) usam pré-condicionamento diagonal para abordar isso, mas aplicar tais mecanismos a camadas recorrentes geralmente requer estados de segunda ordem densos ou quebra os pipelines de paralelização em blocos amigáveis ao hardware (como a transformada WY) que tornam o DeltaNet eficiente.

2. Metodologia: OSDN (DeltaNet Escalado Online)

Os autores propõem o OSDN, que aumenta o portão escalar com um pré-condicionador diagonal $D_t = \text{diag}(d_t)$, aprendido online via feedback de hipergradiente.

2.1. Equivalência Algébrica e Preservação de Hardware

Uma ideia central do OSDN é que o pré-condicionamento à direita do passo de gradiente é algebricamente equivalente a um escalamento por característica da chave do lado de escrita.
A atualização padrão:
$$S_t = S_{t-1} - \beta_t \nabla f_t(S_{t-1}) D_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)(d_t \odot k_t)^\top$$
Ao definir uma chave pré-condicionada $\tilde{k}_t = d_t \odot k_t$, a atualização torna-se:
$$S_t = S_{t-1} + \beta_t (v_t - S_{t-1}k_t)\tilde{k}_t^\top$$
Essa transformação permite que o OSDN preserve estritamente o pipeline paralelo em blocos WY do DeltaNet. A implementação requer apenas substituir a chave de armazenamento $K$ por $\tilde{K}$ na matriz Gram do bloco e nas pontuações intra-bloco, incorrendo em nenhuma sobrecarga de estado de alta dimensão além do vetor pré-condicionador $O(K)$.

2.2. Feedback de Hipergradiente Desacoplado

O pré-condicionador $d_t$ é atualizado via Descida de Gradiente Online (OGD) em um surrogado de hipergradiente derivado do framework do Método de Gradiente Escalado Online (OSGM) [18].
O surrogado de perda para o pré-condicionador é:
$$h_t(d_t) = \frac{f_t(S_t) - f_t(S_{t-1})}{\|\nabla f_t(S_{t-1})\|_F^2}$$
Crucialmente, devido à estrutura exata-quadrática da perda de regressão interna $f_t$, esse surrogado admite uma expressão de forma fechada que depende apenas da chave $k_t$ e do portão escalar $\beta_t$, desacoplando-o do estado de alta dimensão $S_t$, valor $v_t$ ou resíduo $u_t$.
$$h_t(d_t) = \frac{(1 - \beta_t \langle d_t, k_t^2 \rangle)^2 - 1}{2 \|k_t\|_2^2}$$
Esse desacoplamento permite uma implementação em duas fases:

1. Fase 1: Uma varredura leve $O(K)$ sobre o fluxo de chaves para calcular a sequência de chaves pré-condicionadas $\{\tilde{k}_t\}$ e atualizar $d_t$.
2. Fase 2: A passagem padrão do DeltaNet em blocos usando $\tilde{K}$ em vez de $K$.

2.3. Esquecimento Adaptativo do Pré-condicionador (APF)

Para lidar com contextos de linguagem não estacionários onde as distribuições de chaves mudam, os autores introduzem o APF. Isso adiciona um portão de retenção aprendido, token a token e cabeça a cabeça, $r_{t,h}$ aplicado apenas ao estado do pré-condicionador $d_t$ (não ao estado de memória $S_t$):
$$\bar{d}_{t+1} = r_{t,h} d_t + \eta \nabla_d h_t(d_t)$$
Isso mantém a estrutura de recorrência afim necessária para a varredura em duas fases, enquanto permite que o meta-otimizador "esqueça" calibrações obsoletas.

3. Garantias Teóricas

O artigo estabelece duas garantias de convergência baseadas na natureza exata-quadrática da perda interna:

1. Convergência Super-Geométrica no Limite da População (Teorema 5.1): Contra um comparador de Newton à direita ideal ($D^\star = \Sigma_k^\dagger$), se o aprendiz online atingir arrependimento sublinear, a subotimalidade do objetivo global contrai a uma taxa super-geométrica de $O((C/T)^T)$. Isso é significativamente mais rápido que a convergência geométrica (linear) padrão.
2. Contração de Resíduo Local do Algoritmo (Teorema 5.2): Para a atualização diagonal implementada, o produto das razões de resíduo local por token $\prod_{t=1}^T \frac{f_t(S_t)}{f_t(S_{t-1})}$ é limitado. Se um comparador diagonal satisfizer a condição de Newton com portão ($\beta_t \langle d^\star, k_t^2 \rangle = 1$), a taxa de contração também é super-geométrica, especificamente $O((1/\sqrt{T})^T)$ sob limites de arrependimento OGD padrão.

4. Resultados Experimentais

Os experimentos foram conduzidos nos backbones DeltaNet, Gated DeltaNet (GDN) e Kimi Delta Attention (KDA) nas escalas de 340M e 1.3B parâmetros.

4.1. Recordação em Contexto (Escala 340M)

• Desempenho: O OSDN vanilla melhorou a recordação em estilo JRT em 32% sobre o DeltaNet de base (0.150 $\to$ 0.198). O ganho concentrou-se em tarefas de contexto repetido (por exemplo, variantes "JRT-duas vezes"), onde as chaves recorrentes permitem que o pré-condicionador se calibre.
• Impacto do APF: O OSDN-APF manteve um ganho de 17% (0.176) e forneceu a menor perplexidade WikiText/LAMBADA entre as variantes do DeltaNet (30.99 vs. 32.00).
• Generalização: Em benchmarks mais amplos (Commonsense, LongBench), as variantes do OSDN permaneceram em paridade com suas bases, confirmando que o mecanismo é direcionado e não um impulsionador universal de desempenho.

4.2. Escalonamento para 1.3B Parâmetros

• Contração de Resíduo: O sinal ao nível do mecanismo amplificou-se na escala. A razão média geométrica de resíduo ($q_{geo}$) caiu de 0.432 para 0.265 (uma redução de 39%), quase dobrando a melhoria observada em 340M.
• Paridade a Jusante: Apesar da melhoria massiva na contração de resíduo interna, as métricas a jusante (perplexidade WikiText/LAMBADA, Commonsense, LongBench) permaneceram em paridade com a base do DeltaNet de 1.3B, demonstrando que o mecanismo transfere efetivamente sem degradar capacidades gerais.

4.3. Eficiência

• Vazão: As variantes do OSDN incorreram em sobrecarga mínima de inferência. Em 340M, as mudanças de vazão ficaram dentro de $\pm 2.2\%$. Em 1.3B, o OSDN-APF foi 6.8% mais lento que a base, atribuído à varredura sequencial adicional para o pré-condicionador.
• Sobrecarga de Estado: O estado recorrente persistente aumentou apenas $\le 0.05\%$ devido ao vetor pré-condicionador $O(K)$.

5. Significado e Alegações

O artigo posiciona o OSDN como uma extensão mínima e direcionada à Regra Delta que preenche a lacuna entre atualizações escalares de primeira ordem e otimização adaptativa.

• Melhoria ao Nível do Mecanismo: A contribuição principal é uma melhoria provável, ao nível do mecanismo, na recordação associativa via pré-condicionamento online, validada por um diagnóstico direto da razão de contração de resíduo.
• Compatibilidade com Hardware: Ao contrário de outros métodos de pré-condicionamento de segunda ordem ou densos, o OSDN preserva o pipeline paralelo em blocos e a complexidade de estado $O(K)$, tornando-o prático para implantação em grande escala.
• Escalabilidade: Os resultados demonstram que os benefícios do pré-condicionamento online transferem-se e amplificam-se na escala de bilhões de parâmetros, sugerindo que o "desajuste de curvatura" na atenção linear é um gargalo fundamental que pode ser abordado sem uma reestruturação arquitetônica.
• Limitações: Os autores notam explicitamente que o OSDN não é um impulsionador universal de benchmarks; os ganhos concentram-se em tarefas de recuperação envolvendo chaves repetidas. As garantias teóricas dependem da suposição de perda exata-quadrática e descida monótona, que são condições idealizadas. O método não afirma resolver todas as limitações de aprendizado em contexto, mas aborda especificamente a geometria do passo de escrita da Regra Delta.