FG$^2$-GDN: Enhancing Long-Context Gated Delta Networks with Doubly Fine-Grained Control

Resumo Técnico: FG2-GDN

1. Problema e Contexto

As arquiteturas baseadas em Transformers com mecanismos de atenção softmax sofrem de complexidade quadrática ($O(L^2)$) em relação ao comprimento da sequência, o que limita a eficiência no treinamento e a escalabilidade na inferência para contextos longos. Mecanismos de atenção linear surgiram como alternativas promissoras, reduzindo a complexidade para $O(L)$.

Recentes avanços, como a Gated DeltaNet (GDN) e a Kimi Delta Attention (KDA), demonstraram que a regra delta (uma atualização baseada em descida de gradiente online) supera as atualizações aditivas simples no reconhecimento associativo.

• A GDN introduziu um fator de decaimento escalar ($\alpha_t$) para controle de esquecimento.
• A KDA refinou isso para um decaimento channel-wise (por canal), permitindo que cada dimensão de característica tenha sua própria taxa de esquecimento.

A Lacuna Identificada: Embora a KDA tenha refinado o controle de esquecimento (decaimento) para ser granular por canal, a taxa de aprendizado ($\beta_t$) na atualização delta permanece um escalar compartilhado por todas as dimensões. Isso limita a capacidade do modelo de realizar adaptação específica por dimensão durante o processo de "escrita" (atualização) da memória, impedindo que o modelo ajuste a força de escrita independentemente para cada canal de características.

2. Metodologia Proposta

Os autores propõem o FG2-GDN (Fine-Grained Gated Delta Network), que estende a KDA introduzindo taxas de aprendizado adaptativas finas e, posteriormente, o FG2-GDN+, que desacopla ainda mais o controle.

A. FG2-GDN (Controle Fino Duplo)

• Mudança Fundamental: Substitui o escalar $\beta_t \in \mathbb{R}$ por um vetor $\beta_t \in \mathbb{R}^{d_k}$, análogo à transição de SGD para otimizadores adaptativos por coordenada (como AdaGrad ou Adam).
• Regra de Atualização: Para manter a estrutura de rank-1 necessária para a paralelização eficiente em chunks (baseada na representação WY), o vetor $\beta_t$ é absorvido simetricamente nos vetores de chave ($k_t$) e valor ($v_t$):
  $$\tilde{k}_t = \sqrt{\beta_t} \odot k_t, \quad \tilde{v}_t = \sqrt{\beta_t} \odot v_t$$
  A atualização do estado oculto $S_t$ torna-se:
  $$S_t = (I - \tilde{k}_t \tilde{k}_t^\top) \text{Diag}(\alpha_t) S_{t-1} + \tilde{k}_t \tilde{v}_t^\top$$
• Interpretação: Isso permite que cada dimensão de característica mantenha dinâmicas de adaptação independentes, concentrando a "energia de escrita" nas dimensões mais discriminativas para o contexto atual.

B. FG2-GDN+ (Desacoplamento de Forças)

• Inovação: Introduz vetores separados $\beta^k_t$ e $\beta^v_t$ para chaves e valores, respectivamente.
• Objetivo: Desacoplar a "força de apagamento" (erasure strength, controlada por $\beta^k$) da "força de escrita" (write strength, controlada por $\beta^v$). Isso permite que o modelo apague informações antigas de forma agressiva em certas dimensões enquanto escreve novos valores com intensidade diferente em outras.

C. Arquitetura Híbrida
O modelo utiliza uma arquitetura híbrida que intercala camadas de atenção linear (FG2-GDN) com camadas de Multi-head Latent Attention (MLA) em uma proporção de 3:1. Isso combina a eficiência de longo alcance da atenção linear com a capacidade de interação completa do MLA para tarefas que exigem atenção total.

3. Contribuições Principais

1. Generalização da Regra Delta: Estender o controle fino do decaimento (já presente na KDA) para a taxa de aprendizado/escrita, permitindo adaptação por dimensão.
2. Eficiência Computacional: A nova regra de atualização preserva a estrutura Diagonal-Plus-Low-Rank (DPLR). Isso garante que o FG2-GDN possa reutilizar algoritmos de treinamento paralelo em chunks existentes (como o da KDA), com um overhead de pré-preenchimento (prefill) inferior a 5% em relação à KDA.
3. FG2-GDN+: Uma variante que oferece controle independente sobre a intensidade de escrita e apagamento, melhorando a gestão de múltiplas associações concorrentes em contextos longos.
4. Validação Empírica: Demonstração de que o controle fino duplo melhora o reconhecimento associativo e a compreensão de contexto longo sem sacrificar a qualidade da modelagem de linguagem.

4. Resultados Experimentais

Os experimentos foram conduzidos em escalas de 340M e 1.3B parâmetros, comparando FG2-GDN e FG2-GDN+ contra GDN e KDA.

• Modelagem de Linguagem (LM):
  • O FG2-GDN obteve consistentemente a melhor precisão média zero-shot em benchmarks de raciocínio de senso comum (LAMBADA, ARC, etc.) em ambas as escalas.
  • Redução de perplexidade no LAMBADA de 13.80 (KDA) para 13.09 (FG2-GDN) na escala de 1.3B.
• Contexto Longo (RULER e LongBench):
  • RULER: O FG2-GDN superou significativamente os baselines em tarefas de recuperação de chaves únicas (MQ, CWE) e múltiplas. O FG2-GDN+ liderou em tarefas de recuperação de múltiplas chaves (MK2), especialmente em sequências longas (16k), sugerindo que o desacoplamento de escrita e apagamento ajuda a gerenciar interferências.
  • LongBench: O FG2-GDN+ alcançou a melhor pontuação média geral (18.3%) e em tarefas few-shot (32.1%), superando todos os outros modelos.
• Extrapolação: O modelo manteve alta precisão ao ser testado em sequências de 16k, embora tenha sido treinado apenas com 8k, indicando uma representação de estado mais estruturada e robusta a mudanças de distribuição posicional.
• Eficiência: A taxa de transferência (throughput) de pré-preenchimento manteve-se comparável à KDA, com menos de 5% de sobrecarga, preservando a escalabilidade linear.

5. Significado e Conclusão

O trabalho FG2-GDN estabelece que o controle granular não deve se limitar apenas ao esquecimento (decaimento), mas também deve ser aplicado à taxa de aprendizado (escrita) em redes de memória recorrente linear.

Ao introduzir taxas de aprendizado por coordenada, o modelo consegue formar ligações chave-valor mais precisas e "apertadas", sobrevivendo melhor a perturbações na formulação da consulta. A proposta demonstra que é possível obter ganhos significativos na capacidade de memória associativa e compreensão de contexto longo com um custo computacional marginal, oferecendo um novo ponto de dados crucial para o desenvolvimento de modelos de linguagem eficientes e escaláveis para contextos extensos.