Resumo Técnico: Leis de Escala Espectral Induzidas pelo Otimizador

Declaração do Problema

As leis de escala clássicas para modelos de linguagem de grande porte (LLMs) estabeleceram relações previsíveis de lei de potência entre a perda de validação e recursos como tamanho do modelo, volume de dados e capacidade de computação. No entanto, essas estruturas geralmente tratam o otimizador como um detalhe fixo de treinamento, assumindo que ele influencia apenas a velocidade de convergência ou atua como um fator de eficiência multiplicativo. Este trabalho desafia essa premissa, postulando que o otimizador é um eixo fundamental da escala de representação. Especificamente, os autores investigam como diferentes otimizadores convertem efetivamente o aumento da largura da Rede de Alimentação Direta (FFN) em capacidade espectral utilizada — a dimensionalidade efetiva do espaço latente. Eles questionam se uma perda de validação equivalente implica uma geometria de representação interna equivalente, ou se diferentes otimizadores realizam a mesma capacidade arquitetônica através de estruturas espectrais fundamentalmente distintas.

Metodologia

O estudo emprega uma estrutura de análise espectral centrada nos espectros de autovalores das representações da FFN dentro de arquiteturas Transformer.

• Métricas Espectrais: Os autores utilizam a família de rank efetivo de Rényi para quantificar a dispersão espectral. Eles ancoram sua análise em duas métricas primárias:
  • Rank Suave ($R_1$): Corresponde a $\alpha=1$ (entropia de Shannon), medindo a dispersão espectral ponderada pela entropia e capturando a capacidade difusa.
  • Rank Rígido ($R_2$): Corresponde a $\alpha=2$ (razão de participação), medindo a dimensionalidade efetiva sensível à concentração e capturando a capacidade de modo dominante.
  • Assimetria Rígida-Suave ($\Delta_{1,2}$): Definida como $\log R_1 - \log R_2$, esta métrica quantifica a lacuna entre a capacidade difusa e a dominante.
• Configuração Experimental: Os experimentos foram conduzidos em Transformers apenas com decodificador no estilo GPT (160M e 350M parâmetros) treinados no conjunto de dados FineWeb-Edu. A dimensão oculta da FFN ($D$) foi variada para estabelecer leis de escala ($R(D) \propto D^\beta$).
• Comparação de Otimizadores: O estudo compara o AdamW padrão contra otimizadores sensíveis a matrizes (Muon, NorMuon) e variantes com restrição de rank (Dion com frações de rank $r$ variáveis).
• Estratificação: Para abordar a distribuição de Zipf dos dados linguísticos, as representações da FFN foram estratificadas pela frequência de tokens em regimes HEAD (frequente), MID e TAIL (raro).
• Controles Arquiteturais: Os autores compararam deslocamentos induzidos pelo otimizador contra intervenções arquitetônicas, especificamente o aumento do rank de atenção por cabeça (reduzindo a contagem de cabeças) e a remoção de Embeddings Posicionais Rotacionais (RoPE).

Principais Resultados

1. Leis de Escala Espectral Dependentes do Otimizador

A mesma arquitetura Transformer exibe leis de escala espectral marcadamente diferentes dependendo do otimizador.

• Escala de Rank Rígido: O AdamW exibe uma escala de rank rígido fraca ($\beta_{hard} \approx 0,29$ a $0,44$), particularmente para tokens raros (TAIL). Em contraste, o Muon alcança uma escala quase linear ($\beta_{hard} \approx 1,02$) nos mesmos regimes. Isso representa um aumento de 2,3$\times$ no expoente de escala.
• Assimetria Rígida-Suave: O AdamW exibe uma grande assimetria positiva ($\Delta_{1,2} \approx 0,37$), indicando que a largura adicional da FFN expande a massa espectral difusa em vez de direções de representação dominantes. O Muon e o NorMuon reduzem essa assimetria para próximo de zero ($\Delta_{1,2} \approx 0,14$), significando que a largura adicionada é convertida eficientemente em capacidade de modo dominante.

2. Perda Equivalente $\neq$ Geometria Equivalente

O trabalho demonstra que a perda de validação é um proxy insuficiente para a estrutura de representação.

• Treinamento Estendido: Treinar o AdamW por mais passos (12K vs. 6K) permitiu que ele correspondesse à perplexidade de uma variante do Dion (1/16). No entanto, a geometria espectral permaneceu distinta: a escala de rank rígido do AdamW quase desapareceu ($\beta_{hard} \to 0,03$), enquanto o Dion manteve uma escala de lei de potência robusta.
• Taxa de Aprendizado: Varreduras de taxas de aprendizado confirmaram que a lacuna espectral AdamW-Muon não é um artefato de hiperparâmetro; o menor expoente de rank rígido do Muon consistentemente superou o maior do AdamW em toda a varredura.

3. Dependências de Frequência e Rank

• Frequência de Token: O efeito do otimizador é mais pronunciado nos regimes de tokens MID e TAIL. O AdamW luta para escalar a capacidade rígida para tokens raros, enquanto o Muon escala linearmente.
• Rank de Atualização: Usando o Dion, os autores isolaram o efeito do rank de atualização. Reduzir a fração de rank ($r$) de $1/2$ para $1/16$ causou uma queda na escala de rank rígido de $\beta \approx 0,88$ para $\beta \approx 0,40$ (aproximando-se dos níveis do AdamW), enquanto a escala de rank suave degradou-se mais gradualmente. Isso indica que o rank de atualização atua como um botão de controle para converter largura em capacidade de modo dominante.

4. Dominância sobre Intervenções Arquiteturais

Deslocamentos espectrais induzidos pelo otimizador frequentemente superam ou remodelam os efeitos de mudanças arquitetônicas:

• Rank de Atenção: Aumentar o rank de atenção por cabeça (reduzindo a contagem de cabeças) produziu deslocamentos de escala espectral menores do que trocar de otimizador em 28 de 30 regimes de frequência/métrica.
• Codificação Posicional: A remoção do RoPE induziu uma redistribuição dependente do otimizador. Para o Muon, o NoPE aumentou a escala de rank rígido HEAD, mas diminuiu a escala MID/TAIL; para o NorMuon, diminuiu a escala em todos os regimes. Isso mostra que sinais arquitetônicos são expressos através da geometria do otimizador.

Significado e Afirmações

O trabalho argumenta que a otimização deve ser tratada como um eixo de primeira classe da escala de representação, ao lado do tamanho do modelo, dados e computação.

• Co-Design Arquitetura-Otimizador: A capacidade espectral realizada é uma propriedade do par arquitetura-otimizador, não da arquitetura isoladamente. Intervenções arquitetônicas (como mudar o rank de atenção ou codificações posicionais) não induzem deslocamentos fixos e independentes do otimizador; seus efeitos são mediados pela geometria do otimizador.
• Estrutura de Representação: A escolha do otimizador determina como a capacidade do modelo é estruturada através dos modos próprios e alocada através dos regimes de frequência de token. Não é meramente um botão de eficiência para a velocidade de convergência.
• Limitações: Os autores notam que essas descobertas foram estabelecidas dentro de um regime empírico controlado (modelos estilo GPT até 350M parâmetros). Embora os resultados de 350M repliquem os padrões de 160M, validação adicional em escalas de fronteira (1B+) é necessária. Além disso, embora os ranks espectrais quantifiquem a utilização da capacidade, eles não caracterizam totalmente o comportamento em tarefas downstream.

Em conclusão, o estudo motiva uma mudança em direção ao co-design otimizador-arquitetura, sugerindo que futuras leis de escala devem levar em conta os vieses indutivos implícitos do otimizador como um determinante primário de como a capacidade do modelo é realizada.