Contrastive Regularization of Machine Learning Potentials

Resumo Técnico: Regularização Contrastiva de Potenciais de Aprendizado de Máquina

Definição do Problema
Potenciais de Interação Atômica de Aprendizado de Máquina (MLIPs) são tipicamente treinados como regressores para prever energias e forças a partir de dados de Teoria do Funcional da Densidade (DFT), minimizando o Erro Quadrático Médio (MSE). Embora esses modelos frequentemente alcancem precisão química em conjuntos de teste não vistos, os autores demonstram que eles frequentemente falham como geradores quando implantados para conduzir simulações de dinâmica molecular ou Monte Carlo.

A questão central é estatística, e não arquitetural. Em espaços de configuração de alta dimensão, um conjunto de treinamento finito cobre uma fração ínfima do espaço. A minimização padrão do MSE restringe o potencial apenas onde os dados existem, deixando o restante do espaço de configuração não restringido. Consequentemente, redes neurais expressivas inevitavelmente esculpem mínimos de baixa energia espúrios nessas regiões não visitadas. Quando usados como amostradores, as trajetórias derivam para esses mínimos "fantasmagóricos", resultando em configurações não físicas e observáveis termodinâmicos (ex: perfis de energia livre, distribuições estruturais) que se desviam drasticamente da referência, embora as previsões pontuais de energia e força permaneçam precisas.

Metodologia: MSE com Regularização Contrastiva (CRMSE)
Para abordar isso, os autores introduzem o MSE com Regularização Contrastiva (CRMSE), um procedimento de pós-treinamento que aumenta o MSE padrão com um termo contrastivo derivado da divergência de Kullback–Leibler (KL) entre a distribuição de Boltzmann implícita do modelo e a distribuição alvo.

Componentes metodológicos principais incluem:

• Função Objetivo: A perda total é $L_{CRMSE} = L_{MSE} + \lambda L_{CR}$. O termo MSE mantém a precisão nos dados de treinamento, enquanto o termo contrastivo ($L_{CR}$) otimiza a distribuição de probabilidade.
• Mecanismo Contrastivo: O termo contrastivo é derivado da minimização da divergência KL. Ele atua como um objetivo de energia livre variacional: ele reduz a energia das amostras "positivas" (os dados de treinamento) e aumenta a energia das amostras "negativas" (configurações extraídas da distribuição atual do modelo).
• Amostragem Autoconsistente: Como a distribuição do modelo é intratável, as amostras negativas são geradas usando o Algoritmo de Langevin Não Ajustado (ULA). Os autores empregam a Divergência Contrastiva Persistente (PCD), onde um buffer de replay de configurações negativas é mantido e atualizado com alguns passos de ULA por iteração.
• Natureza de Pós-Treinamento: O CRMSE é aplicado como um passo de pós-treinamento a um potencial já otimizado por MSE. Ele não requer novos dados ab initio nem mudanças arquiteturais. O modelo efetivamente "rotula seus próprios erros" ao identificar configurações nas quais ele deriva durante a simulação e aumenta a energia delas para criar uma barreira protetora ao redor do poço físico.

Principais Resultados
O método foi validado nas moléculas de etanol e aspirina do conjunto de dados MD17 usando uma arquitetura de Rede Neural de Grafo com Referencial Local (GNN-LF).

1. Falha do MSE Padrão: Modelos treinados com MSE padrão alcançaram alta precisão em energias e forças de conjuntos de teste. No entanto, ao serem usados em simulações de algoritmo de Langevin com ajuste de Metropolis (MALA), as trajetórias escaparam rapidamente do poço físico, estabelecendo-se em mínimos espúrios dezenas de kcal/mol abaixo da energia de referência. Isso resultou em distâncias interatômicas não físicas e distribuições de energia incorretas.
2. Correção via CRMSE:
   • Estabilidade de Amostragem: O pós-treinamento via CRMSE confinou o amostrador ao poço físico. As distribuições de energia amostradas e os histogramas de distância interatômica recuperaram concordância quase quantitativa com as referências de DFT.
   • Observáveis Termodinâmicos: Crucialmente, o CRMSE recuperou os perfis de energia livre ao longo de ângulos diedros (H-C-C-O para o etanol e diedro de éster para a aspirina) com concordância quase quantitativa com o DFT. Ele corrigiu a subestimação das barreiras de torção e colapsou a dispersão inter-cadeia (variância entre cadeias de simulação independentes) para dentro do ruído estatístico.
   • Preservação da Precisão: A correção não degradou a precisão preditiva do modelo no manufoldo de treinamento; os erros de energia e força nos dados de teste permaneceram dentro da precisão química.
   • Robustez: O método permaneceu eficaz mesmo quando o conjunto de treinamento foi drasticamente reduzido de 950 para 200 configurações, demonstrando melhor eficiência de dados e robustez à escassez de dados.

Significância e Alegações
O artigo afirma que a amostragem confiável depende menos do volume de dados de treinamento do que de treinar o modelo contra a distribuição que ele produz. Os autores argumentam que o treinamento de nível de distribuição é um requisito distinto, não meramente um refinamento da precisão de regressão.

• Requisito Distinto: O trabalho destaca que a precisão pontual não garante uma distribuição de amostragem correta. Um potencial destinado a ser amostrado deve ser treinado para reproduzir sua própria distribuição de amostragem.
• Complementaridade: O CRMSE é apresentado como complementar ao aprendizado ativo. Enquanto o aprendizado ativo injeta novos dados ab initio para melhorar a precisão em regiões fora do manifold, o CRMSE restaura o suporte correto da distribuição de amostragem sem novos dados, suprimindo mínimos espúrios ao elevar sua energia.
• Generalidade: Os autores sugerem que esta lição se aplica amplamente a configurações de IA onde regressores aprendidos são implantados como geradores. O método é particularmente significativo porque repara um potencial que já está "bem otim-imizado" por métricas de regressão padrão, utilizando apenas um passo de pós-treinamento barato e sem cálculos de referência adicionais.

Os autores concluem que, para modelos substitutos usados como geradores, o objetivo de treinamento deve considerar a densidade de probabilidade completa definida pelo potencial, um requisito estruturalmente omitido pela minimização padrão de MSE.