Contrastive Regularization of Machine Learning Potentials

Technische Samenvatting: Contrastieve Regularisatie van Machine Learning Potentialen

Probleemstelling
Machine Learning Interatomic Potentials (MLIPs) worden doorgaans getraind als regressoren om energieën en krachten te voorspellen vanuit Density Functional Theory (DFT) data, waarbij de Mean Squared Error (MSE) wordt geminimaliseerd. Hoewel deze modellen vaak chemische nauwkeurigheid bereiken op onafhankelijke testsets, laten de auteurs zien dat ze vaak falen als generatoren wanneer ze worden ingezet om moleculaire dynamica of Monte Carlo-simulaties aan te sturen.

Het kernprobleet is statistisch van aard en niet architecturaal. In hoogdimensionale configuratieruimtes dekt een eindige trainingsset slechts een verwaarloosbaar klein deel van de ruimte. Standaard MSE-minimalisatie beperkt de potentiaal alleen waar data aanwezig is, waardoor de rest van de configuratieruimte onbeperkt blijft. Gevolgelijk zullen expressieve neurale netwerken onvermijdelijk spuria (valse) laag-energetische minima creëren in deze onbezochte regio's. Wanneer ze als samplers worden gebruikt, driften trajecten af naar deze "geest"-minima, wat resulteert in onfysische configuraties en thermodynamische observabelen (bijv. vrije-energieprofielen, structurele distributies) die sterk afwijken van de referentie, zelfs wanneer de puntvoorwaardelijke energie- en krachtvoorspellingen accuraat blijven.

Methodologie: Contrastieve Regularized MSE (CRMSE)
Om dit aan te pakken, introduceren de auteurs Contrastive Regularized MSE (CRMSE), een post-training procedure die de standaard MSE-loss aanvult met een contrastieve term afgeleid van de Kullback–Leibler (KL) divergentie tussen de impliciete Boltzmann-distributie van het model en de doelverdeling.

Belangrijke methodologische componenten zijn:

• Doelfunctie: De totale loss is $L_{CRMSE} = L_{MSE} + \lambda L_{CR}$. De MSE-term behoudt de nauwkeurigheid op de trainingsdata, terwijl de contrastieve term ($L_{CR}$) de waarschijnlijkheidsverdeling optimaliseert.
• Contrastief Mechanisme: De contrastieve term is afgeleid van het minimaliseren van de KL-divergentie. Het fungeert als een variationele vrije-energie-objectief: het verlaagt de energie van "positieve" samples (de trainingsdata) en verhoogt de energie van "negatieve" samples (configuraties getrokken uit de huidige distributie van het model).
• Zelfconsistent Samplen: Omdat de distributie van het model onhandelbaar (intractable) is, worden negatieve samples gegenereerd met behulp van het Unadjusted Langevin Algorithm (ULA). De auteurs maken gebruik van Persistent Contrastive Divergence (PCD), waarbij een replay-buffer van negatieve configuraties wordt bijgehouden en bijgewerkt met enkele ULA-stappen per iteratie.
• Post-Training Karakter: CRMSE wordt toegepast als een post-training stap op een reeds door MSE geoptimaliseerde potentiaal. Het vereist geen nieuwe ab initio data en geen architecturale wijzigingen. Het model "labelt effectief zijn eigen fouten" door configuraties te identificeren waar het naartoe drijft tijdens simulatie, en verhoogt de energie daarvan om een beschermende barrière rond de fysieke basin te creëren.

Belangrijkste Resultaten
De methode werd gevalideerd op de ethanol- en aspirine-moleculen uit de MD17-dataset met behulp van een Graph Neural Network met Local Frame (GNN-LF) architectuur.

1. Falen van Standaard MSE: Met standaard MSE-getrainde modellen werden hoge nauwkeurigheid bereikt op onafhankelijke energieën en krachten. Echter, wanneer ze werden gebruikt in Metropolis-adjusted Langevin algorithm (MALA) simulaties, weken trajecten snel af van de fysieke basin en nestelden ze zich in spuria minima die tientallen kcal/mol onder de referentie-energie lagen. Dit resulteerde in onfysische interatomaire afstanden en incorrecte energieverdelingen.
2. Correctie via CRMSE:
   • Sampling Stabiliteit: CRMSE post-training hield de sampler binnen de fysieke basin. De gesamplede energieverdelingen en interatomaire afstandshistogrammen herstelden naar een bijna kwantitatieve overeenkomst met de DFT-referenties.
   • Thermodynamische Observabelen: Cruciaal is dat CRMSE de vrije-energieprofielen langs dihedrale hoeken (H-C-C-O voor ethanol en de ester-dihedrale hoek voor aspirine) herstelde naar een bijna kwantitatieve overeenkomst met DFT. Het corrigeerde de onderschatting van torsiebarrières en bracht de inter-chain dispersie (variantie tussen onafhankelijke simulatieketens) terug tot binnen de statistische ruis.
   • Behoud van Nauwkeurigheid: De correctie degradeerde de voorspellende nauwkeurigheid op het trainingsmanifold niet; fouten in energie en kracht op onafhankelijke data bleven binnen chemische nauwkeurigheid.
   • Robuustheid: De methode bleef effectief, zelfs wanneer de trainingsset drastisch werd verminderd van 950 naar 200 configuraties, wat duidt op verbeterde data-efficiëntie en robuustheid tegen dataschaarste.

Significantie en Claims
Het artikel claimt dat betrouwbaar samplen minder afhangt van het volume aan trainingsdata dan van het trainen van het model tegen de distributie die het zelf produceert. De auteurs stellen dat training op distributieniveau een fundamenteel andere vereiste is, en niet louter een verfijning van regressienauwkeurigheid.

• Distinct Requirement: Dit werk benadrukt dat puntvoorwaardelijke nauwkeurigheid geen garantie biedt voor een correcte sampling-distributie. Een potentiaal die bedoeld is om gesampled te worden, moet getraind worden om zijn eigen sampling-distributie te reproduceren.
• Complementariteit: CRMSE wordt gepresenteerd als complementair aan active learning. Terwijl active learning nieuwe ab initio data injecteert om de nauwkeurigheid in off-manifold regio's te verbeteren, herstelt CRMSE de correcte support van de sampling-distributie zonder nieuwe data, door de energie van spuria minima te verhogen.
• Generaliteit: De auteurs suggereren dat deze les breed van toepassing is op AI-instellingen waar geleerde regressoren worden ingezet als generatoren. De methode is bijzonder significant omdat het een potentiaal repareert die al "goed geoptimaliseerd" is door standaard regressiemetrieken, met behulp van slechts een goedkope post-training stap en zonder extra referentieberekeningen.

De auteurs concluderen dat voor surrogaatmodellen die als generatoren worden gebruikt, het trainingsdoel rekening moet houden met de volledige waarschijnlijkheidsdichtheid die door de potentiaal wordt gedefinieerd, een vereiste die structureel ontbreekt bij standaard MSE-minimalisatie.