A Training-Time Diagnostic for Generalization via the Log-Alignment Ratio

Título: Um Diagnóstico em Tempo de Treinamento para Generalização via a Razão de Alinhamento Logarítmico

Declaração do Problema
À medida que modelos de redes neurais e conjuntos de dados aumentam de escala, o custo computacional das execuções de treinamento cresce, tornando o estudo da dinâmica de treinamento e o desenho de ferramentas de diagnóstico baratas cada vez mais valiosos. Os praticantes necessitam de métricas que possam sinalizar problemas como sobreajuste, instabilidade ou saturação usando apenas quantidades disponíveis durante o treinamento, sem depender de dados de validação retidos e custosos. Embora trabalhos anteriores tenham explorado propriedades espectrais de matrizes de pesos (por exemplo, rank efetivo, rank estável) ou teorias de parametrização (por exemplo, parametrização de atualização máxima, $\mu$P) para compreender a generalização, as métricas existentes frequentemente focam exclusivamente em pesos ou ativações, ou exigem operações computacionalmente caras, como a Decomposição em Valores Singulares (SVD) completa ou Análise de Componentes Principais (PCA).

Metodologia
Os autores investigam a Razão de Alinhamento Logarítmico (LAR), uma métrica originalmente introduzida na teoria de parametrização por Everett et al. [2024] para medir o alinhamento entre a matriz de parâmetros de uma camada e suas ativações de entrada. O artigo reformula a LAR teoricamente para revelar um segundo significado, anteriormente inexplorado: ela mede a sobreposição entre um espectro de pesos ($p$) e um espectro de ativações ($q$).

• Definições:
  • Seja $W$ uma matriz com valores singulares $s_i$. O espectro de pesos $p$ é definido como os valores singulares quadrados normalizados: $p_i = s_i^2 / \sum s_j^2$.
  • Seja $X$ um lote de entradas. O espectro de ativações $q$ é definido como as projeções quadradas normalizadas das entradas sobre as direções singulares de $W$: $q_i = \sum_{x \in B} (v_i^\top x)^2 / \sum_{x \in B} \|x\|^2$.
• Reformulação: Os autores derivam que a LAR pode ser expressa como:
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( \sum_{i} p_i q_i \right)$$
  Esta formulação mostra que a LAR é uma medida do sobreposição logarítmica entre as distribuições $p$ e $q$. Ela também pode ser escrita em termos de covariância e correlação de Pearson:
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( n \cdot \text{cov}(p, q) + \frac{1}{n} \right)$$
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( n \sigma_p \sigma_q \text{corr}(p, q) + \frac{1}{n} \right)$$
  onde $\sigma_p$ e $\sigma_q$ representam a "dispersão" ou concentração das distribuições.
• Eficiência Computacional: A LAR é computável a partir de quantidades disponíveis durante a passagem direta (normas de matriz) com sobrecarga negligenciável ($O(|B|m + |B|n + mn)$), não exigindo dados retidos ou SVD completa.

Principais Contribuições

1. Reformulação Distribucional: O artigo demonstra que a LAR não é meramente uma medida de alinhamento, mas um quantificador da sobreposição entre os espectros de pesos e ativações. Ela fornece formas estatísticas equivalentes (covariância, correlação) que tornam o comportamento da métrica interpretável em termos de quantidades elementares.
2. Previsão de Dimensão Efetiva: Os autores mostram que, quando uma matriz de pesos concentra energia em $k$ direções singulares e as ativações se alinham a elas, a LAR relaciona-se com a dimensão efetiva $k$ via $k \approx n^2(1 - \text{LAR})$. Esta fórmula é empiricamente verificada em experimentos de "grokking".
3. Diagnóstico de Sobreajuste em Escala: Em pré-treinamento em grande escala (modelos de 3B de parâmetros), o artigo estabelece que a LAR de desembedding estabiliza durante a generalização, mas declina abruptamente à medida que o sobreajuste se aproxima. O desvio da LAR em relação a uma linha de base sem sobreajuste rastreia a lacuna de generalização.
4. Proxy para Dispersão Espectral: O estudo encontra que a otimização leva $p$ e $q$ a serem altamente correlacionados. Consequentemente, a LAR atua como um proxy para a variância (dispersão) dessas distribuições. No pré-treinamento, a direção singular superior domina, tornando a contribuição do componente superior ($L_1$) um indicador primário da assinatura pré-sobreajuste.

Resultados Experimentais
Os autores validam essas descobertas em dois cenários distintos:

1. Experimentos de Grokking (Tarefas Algorítmicas Pequenas):

   • Treinaram Transformers de 2 camadas em 12 tarefas de operações binárias.
   • Observação: Durante a fase de memorização, a LAR é mais baixa (indicando $p$ e $q$ difusas). À medida que o modelo transita para o "grokking" (generalização), a LAR aumenta à medida que $p$ e $q$ se concentram em um conjunto compartilhado de direções.
   • Dimensionalidade: A dimensão prevista pela LAR, $k_{pred} = n^2(1 - \text{LAR})$, coincide estreitamente com $k_{95}$ (o número de componentes principais que explicam 95% da variância) nos checkpoints finais, com baixo Erro Quadrático Médio (RMSE).
   • Diferenciação: Modelos que realizam "grokking" alcançam valores finais de LAR mais altos ($>0,65$) em comparação com modelos que não realizam "grokking" ($\le 0,65$).

2. Pré-Treinamento em Grande Escala (Modelos de Linguagem de 3B):

   • Treinaram modelos baseados em Gemma-2 (3B de parâmetros) usando otimizadores Adam e Muon em conjuntos de dados variando de 200M a 100B de tokens.
   • Generalização vs. Sobreajuste: Em regimes sem sobreajuste (por exemplo, 100B de tokens), a LAR sobe até um pico e depois declina lentamente, estabilizando-se. Em regimes de sobreajuste (conjuntos de dados menores), a LAR declina abruptamente à medida que o passo de sobreajuste se aproxima.
   • Lacuna de Generalização: A diferença entre a LAR de uma execução específica e a execução de linha de base de 100B (diferencial de LAR) espelha estreitamente a lacuna de generalização (perda de validação menos perda de treinamento).
   • Mecanismo: O declínio na LAR anterior ao sobreajuste é impulsionado principalmente pelo componente singular superior ($L_1$). Com o Adam, isso é impulsionado pela difusão das ativações ($q_1$); com o Muon, envolve tanto pesos ($p_1$) quanto ativações.

Significado e Alegações
O artigo alega que a LAR oferece um diagnóstico prático, computacionalmente barato e livre de dados para a generalização.

• Eficiência: Diferentemente do rank efetivo ou outras medidas espectrais que exigem SVD, a LAR pode ser computada com sobrecarga negligenciável durante a passagem direta.
• Poder Preditivo: Ela rastreia com sucesso a transição da memorização para a generalização no "grokking" e serve como um sinal de alerta precoce para sobreajuste no pré-treinamento em grande escala, potencialmente permitindo métodos de treinamento sem avaliação.
• Interpretabilidade: Ao vincular a LAR à dispersão dos espectros de pesos e ativações, a métrica fornece uma visão unificada de como as dinâmicas de otimização (concentração vs. difusão de espectros) se relacionam com a generalização do modelo.

Os autores reconhecem limitações, notando que suas evidências são correlacionais, focam principalmente na camada de desembedding e ainda não testaram se regularizar diretamente a dispersão espectral impediria o sobreajuste. Eles sugerem que trabalhos futuros devem explorar a LAR para outros tipos de matrizes e seu comportamento sob tamanhos de lote variáveis.