Natural gradient descent with momentum

Resumo Técnico: Descida de Gradiente Natural com Momento

Autores: Anthony Nouy e Agustin Somacal
Contexto: Otimização em variedades não lineares, Aprendizado de Máquina e Aprendizado Informado por Física (Physics-Informed Learning).

1. O Problema

O artigo aborda a otimização de funções de perda para aproximar uma função alvo $u$ pertencente a um espaço de Banach $V$, utilizando um elemento de uma classe de modelos parametrizada não linear $\mathcal{M} = \{D(\theta) : \theta \in \mathbb{R}^d\} \subset V$. Exemplos típicos incluem redes neurais e redes de tensores.

O problema central é que os métodos de otimização tradicionais, como a Descida de Gradiente (GD), operam apenas no espaço de parâmetros, ignorando a geometria da variedade de aproximação. A Descida de Gradiente Natural (NGD) corrige isso ao considerar a geometria do espaço de funções, utilizando a matriz de Gram do sistema gerador do espaço tangente como pré-condicionador.

No entanto, a NGD possui limitações significativas:

1. Estagnação em Mínimos Locais: Assim como o GD, a NGD carece de inércia e pode ficar presa no primeiro mínimo local encontrado.
2. Desvios Não Lineares: Em variedades não lineares ou com funções de perda não quadráticas (ex.: divergência KL, resíduos de EDPs), o gradiente natural fornece a direção de descida mais íngreme apenas localmente. Um passo finito pode desviar-se significativamente da trajetória ótima devido às não linearidades do modelo.
3. Erro de Estimação: Em cenários estocásticos (amostragem aleatória), os erros de estimação do gradiente e da matriz de Gram podem amplificar esses desvios.

O objetivo do trabalho é introduzir uma versão "natural" de métodos dinâmicos inerciais clássicos (como Heavy-Ball e Nesterov) para a NGD, visando acelerar a convergência e melhorar a exploração do espaço de busca.

2. Metodologia

Os autores propõem uma formulação de momentum diretamente no espaço de funções, em vez de apenas no espaço de parâmetros, e depois discretizam essa dinâmica para obter algoritmos no espaço de parâmetros.

Fundamentos Teóricos:

• Gradiente Natural no Espaço de Funções: O gradiente natural é interpretado como a projeção ortogonal do gradiente funcional no espaço tangente da variedade, equipado com uma métrica adequada (induzida pelo produto interno do espaço ou pela Hessiana da função de perda).
• Fluxo de Gradiente Funcional: A NGD é vista como a discretização temporal de um fluxo de gradiente contínuo na variedade: $\frac{\partial v}{\partial t} = -\text{grad}_{\mathcal{M}} L(v)$.

Algoritmos Propostos:
Os autores derivam versões naturais de dois métodos clássicos de momentum:

1. Natural Heavy-Ball (NHB):

   • Derivado da discretização de um sistema dinâmico de segunda ordem no espaço de funções.
   • Introduz uma variável de momento funcional $P(t)$ que é projetada no espaço tangente atual.
   • Desafio Computacional: A formulação exata requer o cálculo de uma matriz de Gram cruzada ($G^{(k,k-1)}_X$) entre os sistemas geradores das tangentes em iterados consecutivos ($\theta^{(k)}$ e $\theta^{(k-1)}$), o que pode ser custoso.

2. Aproximações Eficientes (Para evitar a matriz cruzada):

   • Quasi-Natural Heavy-Ball (QNHB): Substitui o momento anterior projetado no espaço tangente atual por uma aproximação direta no espaço de parâmetros, assumindo que a mudança na base tangente é pequena.
   • Natural Heavy-Ball com Diferença Funcional (NHB-FD): Aproxima o momento anterior usando a diferença direta entre as funções aproximadas ($v^{(k)} - v^{(k-1)}$), evitando a necessidade de reavaliar gradientes antigos ou calcular matrizes cruzadas complexas.

3. Natural Nesterov (NN):

   • Adaptação do método de Nesterov para variedades não lineares.
   • Envolve um "passo de arrasto" (look-ahead) no espaço de funções, seguido por uma retração para a variedade e um passo de gradiente.
   • Variantes:
     • NN-I: Utiliza a projeção do momento anterior baseada na matriz cruzada.
     • NN-II: Projeta o momento no espaço tangente atual usando a diferença funcional, evitando a matriz cruzada.
   • O artigo discute também versões "Gauss-Newton" (GN), onde a métrica no espaço tangente é induzida pela Hessiana da função de perda, o que é particularmente útil para problemas de estimação de densidade e EDPs.

3. Contribuições Chave

1. Formulação Funcional de Momentum: Derivação rigorosa de algoritmos de momentum (Heavy-Ball e Nesterov) diretamente no espaço de funções, respeitando a geometria da variedade de aproximação, em vez de aplicar momentum apenas no espaço de parâmetros.
2. Algoritmos Práticos e Eficientes: Proposta de variantes (QNHB, NHB-FD, NN-II) que eliminam a necessidade de calcular matrizes de Gram cruzadas ou reavaliar gradientes passados, mantendo a complexidade computacional comparável à NGD padrão.
3. Análise de Convergência e Estabilidade: Demonstração teórica de que, em variedades quase planas, as versões aproximadas (QNHB, NHB-FD) seguem dinâmicas similares às versões exatas.
4. Validação Empírica: Testes extensivos mostrando que a adição de momentum à NGD supera a NGD pura em termos de número de iterações e tempo de computação.

4. Resultados Experimentais

Os autores testaram os métodos em quatro cenários:

1. Série Temporal de Mackey-Glass (Regressão):

   • Os métodos de momentum (NHB, NN-II) alcançaram o erro mínimo em menos da metade das iterações e do tempo comparado à NGD.
   • As versões aproximadas (QNHB, NHB-FD) foram ligeiramente mais lentas que as versões exatas, mas ainda muito superiores à NGD.
   • O método NN-II mostrou a aceleração mais rápida, embora tenha apresentado oscilações devido ao termo de amortecimento decrescente do Nesterov.

2. Classificação XOR Estendida:

   • Uso de perda de entropia cruzada.
   • Os métodos de momentum aceleraram a convergência em ~50% das iterações e ~20% do tempo.
   • A versão Gauss-Newton (usando a Hessiana como métrica) mostrou-se superior à NGD padrão.
   • Ajustes no coeficiente de momentum ($\beta_k$) foram necessários para estabilizar a versão NN-II-FD.

3. Equação de Advecção-Difusão Linear (Physics-Informed Learning):

   • Resolução de EDP via minimização do resíduo.
   • A adição de momentum reduziu o tempo e iterações pela metade para atingir a convergência.

4. Equação de Reação-Difusão Não Linear:

   • Cenário mais complexo com não linearidades fortes.
   • Momentum acelerou a convergência, mas as versões Nesterov puras divergiram ou se comportaram como NGD sem um ajuste cuidadoso do coeficiente de momentum (redução para $0.75\beta_k$).

5. Significado e Conclusão

O trabalho demonstra que a incorporação de inércia (momentum) em métodos de gradiente natural é crucial para otimização em variedades não lineares complexas.

• Superação de Limitações: Os métodos propostos ajudam a escapar de mínimos locais e corrigem desvios causados pela não linearidade da variedade, algo que a NGD pura não consegue fazer.
• Eficiência Computacional: Ao oferecer variantes que evitam o cálculo de matrizes cruzadas caras, os autores tornam a NGD acelerada viável para problemas de grande escala e aprendizado ativo.
• Aplicabilidade: Os resultados são particularmente relevantes para Aprendizado Informado por Física (PINNs), onde a função de perda (resíduo de EDP) frequentemente não é bem condicionada e a geometria do espaço de soluções é complexa.

Em suma, o artigo fornece uma ponte teórica e prática entre a otimização geométrica (NGD) e técnicas de aceleração clássicas (momentum), resultando em algoritmos mais robustos e rápidos para treinamento de modelos não lineares avançados.