Unlocking Feature Learning in Gated Delta Networks at Scale

Resumo Técnico: Desbloqueando o Aprendizado de Características em Redes Delta com Portão (Gated Delta Networks) em Escala

Definição do Problema
O escalonamento de Grandes Modelos de Linguagem (LLMs) enfrenta dois gargalos principais: o custo computacional acentuado de treinar modelos maiores e a complexidade quadrática da autoatenção padrão de Transformers. Arquiteturas recorrentes lineares, como as Redes Delta com Portão (GDN - Gated Delta Networks), surgiram como alternativas eficientes que mantêm a inferência em tempo linear enquanto capturam dependências de longo alcance. No entanto, existe uma lacuna crítica no escalonamento eficaz desses modelos. Embora a Parametrização de Atualização Máxima (µP - Maximal Update Parametrization) tenha possibilitado com sucesso a transferência de hiperparâmetros zero-shot (especificamente taxas de aprendizado) para Transformers padrão e alguns modelos lineares, sua aplicação às Redes Delta com Portão permanece inexplorada. O desafio central reside na dinâmica do estado recorrente da arquitetura, que é atualizada via mecanismos de portão dependentes de dados e regras delta de matriz completa, em vez das recorrências escalares ou estruturas feedforward padrão assumidas pelas derivações de µP existentes. Trabalhos anteriores sobre Modelos de Espaço de Estado Estruturados (SSMs) como o Mamba mostraram que o µP convencional falha para SSMs diagonais, mas o uso de atualizações de estado por matriz completa nas GDNs torna essas correções específicas inaplicáveis.

Metodologia
Os autores abordam essa lacuna derivando rigorosamente a formulação µP para Redes Delta com Portão. A metodologia envolve a propagação de estimativas de tamanho de coordenada através de todo o passo de propagação direta (forward pass), incluindo os mecanismos de portão únicos e as transições de estado recorrente, para determinar as regras de escalonamento necessárias para inicialização, multiplicadores de propagação direta e taxas de aprendizado.

Principais etapas analíticas incluem:

1. Análise do Passo de Propagação Direta: Os autores derivam os tamanhos de coordenada das características projetadas, o estado latente $S_t$ e a leitura (readout). Eles identificam que a normalização L2 de consultas (queries) e chaves (keys), combinada com as atualizações de escrita de rank-um, resulta em escalonamentos de variância específicos. Crucialmente, eles propõem a inserção de um multiplicador $\sqrt{d}$ antes da camada RMSNorm para garantir que a entrada da camada de normalização mantenha um tamanho de coordenada $\Theta(1)$, evitando a interrupção do escalonamento de gradiente necessário para o aprendizado de características.
2. Análise do Mecanismo de Portão: O artigo analisa os parâmetros de portão não lineares $\alpha_t$ (decaimento) e $\beta_t$ (força de escrita). Diferente dos pesos padrão, estes são derivados de pré-ativações escalares envolvendo vetores de pesos treináveis ($W_\alpha, W_\beta$) e parâmetros escalares ($a_{log}, b$). A análise revela que esses componentes de portão não seguem as leis de escalonamento padrão de µP devido aos seus papéis específicos na equação de atualização de estado.
3. Análise do Passo de Retropropagação e Derivação da Taxa de Aprendizado:
   • AdamW: Devido à normalização dos gradientes pelo segundo momento do AdamW, a magnitude efetiva da atualização é consistente entre diferentes classes de pesos, simplificando os requisitos de taxa de aprendizado.
   • SGD: Sob o SGD puro, as magnitudes brutas dos gradientes variam significativamente entre diferentes classes de pesos. Os autores derivam que as matrizes de pesos de portão ($W_\alpha, W_\beta$) requerem um escalonamento de taxa de aprendizado de $\Theta(1/\sqrt{d})$, e os parâmetros de portão escalares ($a_{log}, b$) requerem um escalonamento de $\Theta(\sqrt{d})$. Esses escalonamentos divergem do cenário µP padrão, onde as taxas de aprendizado são tipicamente uniformes ou seguem uma regra simples de $1/n$.

Principais Contribuições

1. Derivação Teórica: O artigo fornece a primeira formulação completa de µP para Redes Delta com Portão. Ele deriva estimativas de tamanho de coordenada através de todo o passo de propagação direta e estabelece variâncias de inicialização e multiplicadores de propagação direta fundamentados para todas as classes de pesos.
2. Novas Regras de Escalonamento: Os autores descobrem que os componentes de portão requerem escalonamentos de taxa de aprendizado não convencionais sob SGD:
   • Matrizes de pesos de portão ($W_\alpha, W_\beta$) requerem escalonamento de $\Theta(1/\sqrt{d})$.
   • Parâmetros de portão escalares ($a_{log}, b$) requerem escalonamento de $\Theta(\sqrt{d})$.
3. Validação Empírica: Os autores pré-treinam modelos de linguagem GDN através de múltiplas larguras (variando de 256 a 1536 dimensões ocultas) usando ambos os otimizadores AdamW e SGD.
   • Resultos: Sob a configuração µP proposta, a taxa de aprendizado ideal transfere zero-shot entre todas as larguras de modelo para ambos os otimizadores. Em contraste, a Parametrização Padrão (SP) e a configuração µP original falham na transferência, com as taxas de aprendizado ideais mudando significativamente conforme a largura do modelo aumenta.
   • Estabilidade: As configurações propostas permitem o aprendizado estável de taxa de aprendizado, confirmando as derivações teóricas.

Significância e Alegações
O artigo afirma que sua análise resolve o problema aberto de parametrizar Redes Delta com Portão para o aprendizado de características em escala. Ao demonstrar que a formulação µP derivada permite a transferência de hiperparâmetros zero-shot, o trabalho valida a correção da derivação teórica e sua utilidade prática. Os autores enfatizam que este é o primeiro trabalho a derivar e validar uma parametrização consistente com µP especificamente para Redes Delta com Portão, distinguindo-o de trabalhos anteriores em SSMs diagonais. Os achados sugerem que a parametrização adequada é essencial para o escalonamento de modelos recorrentes lineares e que os mecanismos de portão específicos em GDN exigem regras de escalonamento únicas que diferem tanto de Transformers padrão quanto de outras arquiteturas lineares. Os autores esperam que essas derivações guiem pesquisas futuras sobre as leis de escala de outras arquiteturas lineares ou híbridas.