Don't Stop Me Yet: Sampling Loss Minima via Dissipative Riemannian Mechanics

Resumo Técnico: Don't Stop Me Yet: Amostragem de Mínimos de Perda via Mecânica Riemanniana Dissipativa

Declaração do Problema
As redes neurais modernas são tipicamente superparametrizadas, resultando em paisagens de perda onde os mínimos globais não são pontos isolados, mas formam subvariedades contínuas e conectadas de soluções invariantes à reparametrização. Embora essas soluções performem igualmente bem nos dados de treinamento, elas podem induzir comportamentos funcionais diversos em entradas não vistas. As abordagens existentes para quantificação de incerteza enfrentam uma dicotomia: métodos que amostram regiões de baixa perda frequentemente falham em impor a interpolação exata dos dados de treinamento, enquanto métodos restritos ao conjunto invariante à reparametrização são frequentemente inerentemente locais, limitando sua capacidade de explorar vales de mínimos distintos. O artigo aborda o desafio de navegar no espaço de "boas soluções" para amostrar modelos diversos que minimizam estritamente a perda de treinamento, mantendo alta incerteza preditiva em dados fora da distribuição (OOD).

Metodologia: Amostrador de Mínimos Dissipativo (DIMS)
Os autores propõem o Amostrador de Mínimos Dissipativo (DIMS), um sistema dinâmico guiado por geometria projetado para amostrar exatamente os conjuntos de nível mínimos da função de perda. A abordagem fundamenta-se na mecânica Riemanniana e em sistemas dissipativos clássicos.

1. Formulação Riemanniana: A superfície de perda é tratada como uma variedade Riemanniana $(M, g)$ embutida em $\mathbb{R}^{K+1}$. O tensor métrico $G(\theta)$ é definido como uma métrica de pullback, $G(\theta) = I_K + \nabla_\theta L(\theta)\nabla_\theta L(\theta)^\top$, que incorpora a curvatura da paisagem de perda na geometria do espaço de parâmetros.
2. Sistema Dinâmico: O processo de amostragem modela uma partícula movendo-se sobre esta variedade. A dinâmica é governada por uma equação diferencial ordinária (EDO) de segunda ordem derivada das equações de Euler-Lagrange. O sistema inclui:
   • Energia Cinética: $T(\alpha, \dot{\alpha}) = \frac{1}{2}\dot{\alpha}^\top G(\alpha)\dot{\alpha}$.
   • Energia Potencial: $V(\alpha) = \kappa(L(\alpha) - L(\theta^*_{glob}))$, atuando como uma atração gravitacional em direção aos mínimos.
   • Dissipação: Um termo de atrito $-\eta(t)\dot{\alpha}$ é introduzido para quebrar a conservação de energia.
3. Atrito Dependente da Velocidade: Uma inovação chave é a definição da função de dissipação $\eta(t)$, que escala com a velocidade da partícula (a raiz quadrada da energia cinética) e o ângulo entre a velocidade e o gradiente. Isso garante que a energia seja dissipada mais rapidamente ao mover-se contra o gradiente, enquanto permite que a partícula traverse o conjunto de nível mínimo (onde o gradiente é ortogonal à velocidade) com atrito mínimo.
4. Procedimento de Amostragem: O amostrador inicializa uma partícula em um mínimo treinado $\theta^*$ com uma velocidade aleatória amostrada de uma distribuição (tipicamente uma aproximação de Laplace respeitando a geometria local). A EDO é integrada numericamente até que a partícula convirja para um novo ponto na subvariedade mínima $S^*(D)$.

Principais Contribuições

1. Amostrador Determinístico para Subvariedades Mínimas: O DIMS fornece um método para visitar continuamente pontos distintos nas subvariedades mínimas da perda de treinamento. É formulado como um sistema dinâmico guiado por geometria com hiperparâmetros fisicamente interpretáveis (coeficiente de atrito $\eta_0$) que controlam a exploração.
2. Garantias Teóricas: O artigo estabelece que a dinâmica proposta explora a superfície de perda antes de convergir para a variedade mínima em tempo finito. Especificamente, o Teorema 1 garante a convergência para o conjunto de equilíbrios (mínimos locais) sob a propriedade de sela estrita, enquanto os Teoremas 2 e 3 fornecem limites sobre o tempo de convergência e o comprimento de arco percorrido, garantindo uma exploração não trivial.
3. Validação Empírica: Os autores demonstram que aproximações posteriores concentradas na variedade mínima generalizam bem. O DIMS melhora o desempenho preditivo, a calibração e, especificamente, a quantificação de incerteza OOD em comparação com aproximações baseadas em Laplace existentes (LA Padrão, LA Linearizada e LA Riemanniana).

Resultados Experimentais
O método foi avaliado em tarefas de regressão (conjunto de dados Snelson) e classificação (Banana, conjuntos de dados UCI, MNIST, Fashion MNIST).

• Regressão: No conjunto de dados Snelson, o DIMS convergiu com sucesso para soluções que ajustam perfeitamente os dados de treinamento enquanto exibem alta incerteza OOD, whereas a Aproximação de Laplace Riemanniana (RLA) falhou em garantir interpolação quando inicializada a partir de posições subótimas.
• Classificação: Nos conjuntos de dados UCI, o DIMS alcançou métricas de Negativo Log-Likelihood (NLL) e calibração competitivas ou superiores em comparação com a RLA e outras linhas de base. Notavelmente, o DIMS foi robusto ao hiperparâmetro de atrito $\eta_0$, frequentemente performando bem com um valor padrão fixo em diversos conjuntos de dados.
• Desempenho OOD: No MNIST e Fashion MNIST, o DIMS superou consistentemente outros métodos nas pontuações AUROC para detecção OOD, indicando uma capacidade superior de distinguir amostras dentro da distribuição daquelas fora da distribuição.

Significado e Alegações
O artigo alega que o DIMS oferece uma abordagem principial para quantificação de incerteza em aprendizado profundo ao mirar explicitamente a geometria das subvariedades de perda mínima. Ao contrário de posteriors aproximados tradicionais que podem se espalhar por regiões de baixa perda sem impor interpolação exata dos dados, o DIMS garante a convergência para a variedade mínima. Os autores postulam que essa abordagem permite a geração de amostras diversas no espaço de funções que concordam nos dados de treinamento, mas variam em outros lugares, uma propriedade crucial para inferência bayesiana. O trabalho sugere que dinâmicas inspiradas na física e conscientes da geometria podem navegar efetivamente pelas paisagens de perda complexas de redes superparametrizadas, oferecendo uma alternativa robusta às técnicas existentes de amostragem e aproximação. Os autores reconhecem que o método é computacionalmente exigente devido à necessidade de resolver EDOs em espaços de alta dimensão, mas argumentam que o trade-off produz estimativas de incerteza superiores.