Resumo Técnico: Transferência de Hiperparâmetros para Memórias Associativas Densas

Declaração do Problema

A aprendizagem profunda moderna depende fortemente da transferência de hiperparâmetros (HP), onde configurações ótimas (por exemplo, escalas de inicialização, taxas de aprendizagem) são identificadas em modelos e conjuntos de dados pequenos e depois transferidas para escalas maiores, a fim de evitar um ajuste proibitivamente caro. Embora existam estruturas teóricas para a transferência de HP em redes feed-forward (por exemplo, baseadas na heurística $\mu$P), elas não foram desenvolvidas para Memórias Associativas Densas (DenseAMs).

As DenseAMs apresentam desafios únicos para os métodos de transferência existentes devido a três características arquitetônicas específicas:

1. Partilha de Pesos: A função de energia exige que os pesos sejam partilhados entre camadas e dentro das operações, ao contrário das redes feed-forward padrão.
2. Dinâmica Residual: As equações dinâmicas que descrevem a descida de energia contêm nativamente conexões residuais.
3. Ativações Íngremes: Para alcançar leis de escalonamento superiores de armazenamento de informação, as DenseAMs utilizam funções de ativação com picos rápidos (por exemplo, ReLUs de alta potência, softmax) que são raras nas arquiteturas feed-forward padrão.

Estes fatores criam um cenário complexo onde a análise dimensional padrão falha em prever como os hiperparâmetros devem escalar com o tamanho do modelo ($N, K$), o tamanho do conjunto de dados ($P$) e o tamanho do lote ($B$).

Metodologia

Os autores iniciam o estudo da transferência de HP do tipo $\mu$P para DenseAMs, derivando prescrições explícitas para o escalonamento de hiperparâmetros no regime de escalonamento proporcional, definido como:
$$N \to \infty, \quad K = \kappa N, \quad P = \rho N, \quad B = \beta P$$
onde $\kappa, \rho, \beta$ são constantes fixas.

Definição do Modelo

O estudo foca numa regra de atualização de DenseAM rasa que executa um único passo de descida de gradiente numa paisagem de energia:
$$f(x) = s_2 W^\top \sigma_C [s_1 W g(x) + b] + c$$
onde $W \in \mathbb{R}^{K \times N}$ são pesos acoplados, $g$ é uma não-linearidade fixa e $\sigma_C$ é uma função de ativação potencialmente centrada. O objetivo de treino é uma perda de remoção de ruído:
$$L(W) = \frac{1}{2B} \sum_{\mu=1}^B \| f_W(x_\mu + \epsilon_\mu) - x_\mu \|^2$$

Modificações Técnicas Chave

Os autores identificam que as parametrizações padrão falham para DenseAMs com ativações não-lineares devido a "picos induzidos pela média" na matriz de Gram das ativações. Para abordar isto, introduzem a centragem:

• Operador de Centragem: Uma matriz de projeção $C = I_K - \frac{1}{K}\mathbf{1}_K\mathbf{1}_K^\top$ é aplicada às pré-ativações e pós-ativações.
• Necessidade: Para ativações não-lineares (ReLU, Softmax) com médias não nulas, a falha em centrar leva a um pico de posto único na matriz de Gram de ativação ($S^\top S$) que escala de forma diferente dos autovalores do volume. Esta discrepância causa instabilidade numérica e impede a transferência de HP. A centragem remove este pico, garantindo que todas as estatísticas internas escalem consistentemente.

Abordagem de Derivação

Os autores adaptam a heurística $\mu$P realizando uma análise dimensional para satisfazer três desideratos:

1. Estabilidade: As estatísticas de inicialização (média/variação) das pré-ativações ($Z$) e saídas ($F$) devem permanecer $O(1)$ à medida que $N, K \to \infty$.
2. Maximalidade: As atualizações de parâmetros ($\Delta Z, \Delta F$) após $T$ passos devem permanecer $O(1)$.
3. Equilíbrio: Os vários termos que contribuem para as atualizações (resultantes das duas aparições dos pesos acoplados $W$) devem estar equilibrados em magnitude.

Ao analisar os gradientes da perda MSE e as atualizações resultantes para os otimizadores SGD e Adam, os autores derivam leis de escalonamento específicas para as escalas de inicialização ($s_1, s_2$) e as taxas de aprendizagem ($\eta_W, \eta_b, \eta_c$).

Contribuições e Resultados Chave

1. Prescrições Explícitas de Escalonamento

O artigo fornece uma tabela de regras de escalonamento (Tabela 1) para diferentes funções de ativação e otimizadores. As descobertas principais incluem:

• Ativações Lineares: No regime proporcional, $s_1 \sim N^{-1/2}$, $s_2 \sim K^{-1/2}$ e $\eta_W \sim K$ (para SGD).
• Ativações do Tipo ReLU: A centragem é crítica. Sem centragem, a transferência de HP falha devido à instabilidade. Com centragem, as prescrições lineares mantêm-se para SGD e Adam. O Adam permanece estável mesmo para ReLUs de alta potência ($p=10$), enquanto o SGD torna-se instável em taxas de aprendizagem elevadas.
• Ativações Softmax:
  • Instabilidade do SGD: O SGD sofre de instabilidade numérica em grandes escalas devido ao "feedback de amplificação de linhas", onde linhas com pesos softmax maiores recebem atualizações desproporcionalmente grandes, levando à localização e divergência. O artigo não fornece uma prescrição de SGD estável para softmax.
  • Sucesso do Adam: O Adam é invariante à escala e exibe transferência de HP bem-sucedida para DenseAMs de softmax centradas com $s_1 \sim N^{-1/2}$, $s_2 \sim K^{1/2}$ e $\eta_W \sim 1$.

2. Validação Empírica

Os autores validam estas descobertas teóricas em múltiplos cenários:

• Escalonamento Proporcional: Experiências em dados gaussianos isotrópicos e MNIST (via subamostragem de plaquetas) demonstram que a dinâmica de treino (curvas de perda, covariâncias de características internas) colapsa entre escalas quando as prescrições derivadas são utilizadas.
• Dados Anisotrópicos: As prescrições funcionam surpreendentemente bem mesmo para dados anisotrópicos e realistas (MNIST), apesar das derivações teóricas assumirem entradas isotrópicas.
• Centragem vs. Não Centrada: Os resultados empíricos (Figuras 7, 8, 10, 11, 13) mostram claramente que os modelos não centrados divergem ou falham em transferir hiperparâmetros, enquanto os modelos centrados exibem dinâmica consistente e transferência bem-sucedida.

3. Colapso Dinâmico

Para além da simples transferência de perda, o artigo demonstra que a parametrização proposta leva a um colapso da dinâmica de treino. Estatísticas internas, como o RMS das atualizações de pesos e as covariâncias de características, permanecem constantes entre diferentes tamanhos de modelo ($N, K$) quando as leis de escalonamento são aplicadas.

Significado e Alegações

O artigo afirma ser o primeiro trabalho a iniciar o estudo da transferência de hiperparâmetros para Memórias Associativas Densas. O seu significado reside em:

• Preenchimento de uma Lacuna: Estende a teoria da transferência de hiperparâmetros de redes feed-forward padrão para modelos baseados em energia com pesos acoplados e ativações íngremes.
• Utilidade Prática: Fornece regras acionáveis e fundamentadas teoricamente para escalar DenseAMs, permitindo que profissionais treinem modelos de grande escala (por exemplo, Transformers de Energia, NRGPT) sem um ajuste de hiperparâmetros em grande escala dispendioso.
• Insight Arquitetural: Identifica a centragem como um requisito fundamental para a estabilidade de DenseAMs com ativações não-lineares, uma nuance anteriormente não abordada no contexto da transferência de HP.

Os autores permanecem modestos quanto ao escopo, notando que a sua análise assume dados isotrópicos para as derivações (embora seja empiricamente robusta de outro modo) e que uma descrição analítica completa da dinâmica de treino em escala infinita para DenseAMs não-lineares permanece um assunto para trabalho futuro. Eles não afirmam resolver a transferência de HP para todas as arquiteturas modernas, mas especificamente para a família DenseAM.