Resumo Técnico: Melhorando o Treinamento de Redes Neurais através do Desacoplamento da Magnitude e Direção de Vetores de Peso

1. Declaração do Problema

O treinamento de redes neurais modernas depende de otimizadores como Adam e Muon que tratam matrizes de pesos como objetos únicos e monolíticos. No entanto, toda matriz de pesos $W$ carrega inerentemente duas quantidades distintas: uma magnitude ($\|W\|$) e uma direção ($\hat{W} = W/\|W\|$). Na otimização padrão, essas duas quantidades são acopladas, levando a dinâmicas de treinamento instáveis e imprevisíveis:

• Dependência de Mudança Direcional: A mudança angular induzida por uma atualização depende inversamente da magnitude atual. Consequentemente, a taxa de aprendizado (LR) não controla diretamente a taxa de mudança direcional, que pode variar entre camadas e ao longo do tempo.
• Deriva de Magnitude Indesejada: As atualizações tendem a ser perpendiculares aos pesos atuais (devido à escala-invariância ou ruído), fazendo com que a magnitude cresça mesmo na ausência de um gradiente radial. Esse "crescimento desenfreado" força os praticantes a depender de mecanismos indiretos como o decaimento de peso (weight decay) para estabilizar o treinamento e o warmup para prevenir instabilidade precoce.
• Limitações de Escalonamento: Métodos atuais para escalonar o treinamento (ex: $\mu$P) frequentemente exigem o ajuste de hiperparâmetros ao alterar a largura ou profundidade do modelo. Além disso, abordagens recentes que restringem os pesos a uma esfera de norma fixa carecem de parâmetros de magnitude aprendíveis, deixando o controle de escala exclusivamente para camadas de normalização.

2. Metodologia: Desacoplamento Magnitude–Direção (MD)

Os autores propõem o Desacoplamento Magnitude–Direção (MD), uma modificação de otimizador que fatoriza cada matriz de peso em uma direção de norma fixa e ganhos de magnitude aprendíveis, atualizando-os com taxas de aprendizado separadas.

Fatoração Central

Cada matriz de pesos $W$ é reparametrizada como:
$$W = \text{diag}(\gamma_{\text{row}}) \, \hat{W} \, \text{diag}(\gamma_{\text{col}})$$
Onde:

• $\hat{W}$ é a direção, restrita a uma hiperesfera fixa (ex: norma de Frobenius).
• $\gamma_{\text{row}}$ e $\gamma_{\text{col}}$ são ganhos de magnitude aprendíveis (vetores atuando em linhas e colunas, respectivamente).

Procedimento de Otimização

O método opera inteiramente dentro do passo do otimizador, mantendo a passagem direta (forward pass) do modelo inalterada (o modelo vê um único tensor fundido $W$):

1. Recuperação: Recupera $\hat{W}$ e os ganhos $\gamma$ a partir do peso fundido $W$.
2. Divisão de Gradiente: Divide o gradiente $G = \partial L / \partial W$ em um gradiente de direção ($\partial L / \partial \hat{W}$) e gradientes de ganho (somados sobre os respectivos eixos).
3. Atualização de Direção: Atualiza $\hat{W}$ usando o otimizador base (ex: Adam, Muon) com taxa de aprendizado $\eta_W$, e então o projeta de volta na esfera para manter a norma fixa.
4. Atualização de Ganho: Atualiza os ganhos $\gamma$ (opcionalmente via um mapa suave como softplus para garantir positividade) usando uma taxa de aprendizado separada $\eta_\gamma$.
5. Reassemblagem: Reconstrói $W$ para a próxima passagem direta.

Propriedades Principais

• Agnóstico ao Otimizador: O método atua como um invólucro (wrapper) em torno de qualquer otimizador base (Adam, Muon, Shampoo, etc.).
• Controle Direto de LR: Ao fixar a direção em uma esfera, a taxa de aprendizado define diretamente a atualização angular relativa, removendo a dependência da magnitude atual do peso.
• Eliminação de Heurísticas: O método remove a necessidade de decaimento de peso (pois os pesos já são restritos em norma) e de warmup (pois atualizações iniciais grandes e desestabilizadoras são reguladas desde o primeiro passo).
• Controle de Escala Granular: Diferente de métodos que dependem apenas de ganhos de camadas de normalização, o Desacoplamento MD fornece controle explícito e aprendido sobre a escala de linhas e colunas individuais, permitindo que o modelo amplifique ou atenue características específicas.

3. Contribuições Principais e Descobertas

Resultados Empíricos

Os autores validam o Desacoplamento MD em modelos densos (181M a 1.29B de parâmetros) e grandes modelos de Mistura de Especialistas (MoE) (até 6.7B de parâmetros totais).

• Ganhos de Desempenho: O Desacoplamento MD melhora sobre as linhas de base bem ajustadas (AdamW e Muon). Em modelos MoE grandes, o MuonMD (Muon com Desacrupamento MD) atinge a perda do AdamW com aproximadamente 2$\times$ menos computação.
• Transferência de Taxa de Aprendizado: O método permite a transferência de taxa de aprendizado zero-shot através da largura do modelo. A LR de matriz ótima encontrada em um modelo pequeno permanece ótima para modelos significativamente maiores sem necessidade de novo ajuste. Isso se aplica tanto ao escalonamento de largura quanto de profundidade.
• Treinamento Sem Warmup: Experimentos mostram que descartar o warmup não apenas mantém a estabilidade, mas melhora a perda final, pois os passos iniciais de treinamento não são mais desperdiçados em taxas de aprendizado reduzidas.
• Insights de Ablação:
  • Eixo de Normalização: Restringir a direção via norma de Frobenius (matriz inteira) mostrou-se o mais flexível e performa comparavelmente às restrições de linha/coluna.
  • Parametrização de Ganho: Um ganho combinado por linha e coluna ($\gamma_{\text{row}} + \gamma_{\text{col}}$) supera ganhos escalares ou de eixo único. A parametrização softplus para os ganhos fornece uma vantagem leve e consistente sobre atualizações diretas ou exponenciais.
  • Embeddings: Restringir vetores de embedding à norma L2 unitária performa ligeiramente melhor do que deixá-los sem restrição.

Comparação com Trabalhos Relacionados

O artigo situa o Desacrupamento MD dentro de um corpo crescente de trabalhos sobre treinamento esférico e reparametrização:

• Diferente do nGPT, que agrupa restrições esféricas com mudanças arquiteturais específicas (ex: normalização L2, resíduos interpolados), o Desacrupamento MD é uma modificação do lado do otimizador que pode ser aplicada a arquiteturas padrão. Quando aplicado à arquitetura do nGPT, o Desacrupamento MD melhora ainda mais o desempenho.
• Diferente de Muown ou HyperP, que podem fixar normas sem magnitudes aprendíveis ou vincular o método a otimizadores específicos, o Desacrupamento MD aprende explicitamente as magnitudes e é agnóstico ao otimizador.
• Ele aborda a "interferência de magnitude-direção" identificada em trabalhos anteriores (ex: Kosson et al., 2024a) ao separar a dinâmica das duas quantidades.

4. Significância e Alegações

O artigo afirma que tratar magnitude e direção como quantidades controladas separadamente produz dinâmicas de treinamento mais previsíveis e é uma melhoria simples e amplamente aplicável aos otimizadores modernos.

• Simplicidade: O método não requer mudanças arquiteturais e adiciona um overhead computacional mínimo (operações elemento a elemento durante o passo do otimizador).
• Robustez: Simplifica a receita de treinamento ao remover a necessidade de decaimento de peso e warmup, reduzindo o espaço de busca de hiperparâmetros.
• Escalabilidade: Fornece um mecanismo fundamentado para a transferência de taxa de aprendizado através da largura do modelo, abordando um gargalo fundamental no escalonamento de grandes modelos de linguagem.

Os autores concluem que, embora a lógica adicional no passo do otimizador introduza um overhead modesto, os benefícios na estabilidade do treinamento, eficiência computacional e transferência de hiperparâmetros tornam o Desacrupamento MD um avanço significativo para o treinamento de redes neurais de larga escala.