8. Delta-learning
8.1. Slides
8.2. Delta-learning operating principle
One of the most robust and simplest approaches of combining ML and QM, which does not require tempering with the QM method itself, is delta-learning. In delta-learning, instead of training the model directly on the target QM level (e.g., FCI), we only learn the difference (delta) between the target and baseline QM levels. The baseline QM level is a computationally faster but less accurate method. The rational is that learning differences is easier and that if some PES regions are insufficiently sampled, the predictions fall back to the baseline method which serves as a ‘fail-safe’.
The operational principle of delta-learning can be seen best in the exercise:
Example ML-delta.
First we need to generate the energies at both baseline and target levels.
import mlatom as ml
import numpy as np
import matplotlib.pyplot as plt
# prepare H2 geometries with bond lengths ranging from 0.5 to 5.0 Å
xyz = np.zeros((451, 2, 3))
xyz[:, 1, 2] = np.arange(0.5, 5.01, 0.01)
z = np.ones((451, 2)).astype(int)
molDB = ml.molecular_database.from_numpy(coordinates=xyz, species=z)
# calculate HF energies
hf = ml.models.methods(method='HF/STO-3G', program='PySCF')
hf.predict(molecular_database=molDB, calculate_energy=True)
molDB.add_scalar_properties(molDB.get_properties('energy'), 'HF_energy') # save HF energy with a new name
# calculate CISD energies
cisd = ml.models.methods(method='CISD/cc-pVDZ', program='PySCF')
cisd.predict(molecular_database=molDB, calculate_energy=True)
molDB.add_scalar_properties(molDB.get_properties('energy'), 'FCI_energy')
Density fitting is not supported for CISD in pyscf and will be turned off automatically
Here we use HF/STO-3G as the baseline and CISD/cc-pVDZ (FCI in our case) as the target for the geometries with H--H bond lengths ranging from 0.5 to 5.0 Å.
We need to calculate the differences between the target and baseline:
molDB.add_scalar_properties(molDB.get_properties('FCI_energy') - molDB.get_properties('HF_energy'), 'delta_energy')
Now we can train the KREG model on the differences (basically, residual errors)just on the 12 training points (using 10 sub-training and 2 validation points).
# use every 40th point of the data as the training set and split it into the sub-training and validation sets
step = 40
trainDB = molDB[::step]
val = trainDB[::10]
sub = ml.molecular_database([mol for mol in trainDB if mol not in val])
# setup the KREG model
kreg = ml.models.kreg(model_file='KREG.npz', ml_program='KREG_API')
# optimize its hyperparameters
kreg.hyperparameters['sigma'].minval = 2**-5 # modify the default lower bound of the hyperparameter sigma
kreg.optimize_hyperparameters(subtraining_molecular_database=sub,
validation_molecular_database=val,
optimization_algorithm='grid',
hyperparameters=['lambda', 'sigma'],
training_kwargs={'property_to_learn': 'delta_energy', 'prior': 0},
prediction_kwargs={'property_to_predict': 'estimated_delta_energy'})
lmbd = kreg.hyperparameters['lambda'].value ; sigma=kreg.hyperparameters['sigma'].value
valloss = kreg.validation_loss
print('Optimized sigma:', sigma)
print('Optimized lambda:', lmbd)
print('Optimized validation loss:', valloss)
# Train the model with the optimized hyperparameters to dump it to disk.
kreg.train(molecular_database=trainDB, property_to_learn='delta_energy')
Optimized sigma: 0.35355339059327373 Optimized lambda: 5.3947966093944364e-06 Optimized validation loss: 0.000380824502657017
After we trained the KREG model, we can use it to estimate the delta corrections for the entire PEC and add them to the baseline predictions to obtain the delta-learning model estimates of the FCI energies:
# predict the delta corrections with the trained KREG model
kreg.predict(molecular_database=molDB, property_to_predict='KREG_delta_energy')
# add them to the baseline values
molDB.add_scalar_properties(molDB.get_properties('HF_energy') + molDB.get_properties('KREG_delta_energy'), 'FCI_energy_est')
Next, let's train the KREG model directly on the FCI data for comparison and use it for predicting the PEC.
# setup the KREG model
kreg_fci = ml.models.kreg(model_file='KREG.npz', ml_program='KREG_API')
# optimize its hyperparameters
kreg_fci.hyperparameters['sigma'].minval = 2**-5 # modify the default lower bound of the hyperparameter sigma
kreg_fci.optimize_hyperparameters(subtraining_molecular_database=sub,
validation_molecular_database=val,
optimization_algorithm='grid',
hyperparameters=['lambda', 'sigma'],
training_kwargs={'property_to_learn': 'FCI_energy', 'prior': 0},
prediction_kwargs={'property_to_predict': 'estimated_FCI_energy'})
lmbd = kreg_fci.hyperparameters['lambda'].value ; sigma=kreg_fci.hyperparameters['sigma'].value
valloss = kreg_fci.validation_loss
print('Optimized sigma:', sigma)
print('Optimized lambda:', lmbd)
print('Optimized validation loss:', valloss)
# Train the model with the optimized hyperparameters to dump it to disk.
kreg_fci.train(molecular_database=trainDB, property_to_learn='FCI_energy')
# predict with the direct KREG model
kreg_fci.predict(molecular_database=molDB, property_to_predict='KREG_FCI_energy')
Optimized sigma: 3.9999999999999987 Optimized lambda: 0.04819408829399811 Optimized validation loss: 0.006956885379877326
Finally, let's check the results.
# plot the energies
plt.plot(xyz[:, 1, 2], molDB.get_properties('HF_energy'), label='HF/STO-3G')
plt.plot(xyz[:, 1, 2], molDB.get_properties('FCI_energy'), label='FCI/cc-pVDZ')
# plt.plot(xyz[:, 1, 2], molDB.get_properties('FCI_energy') - molDB.get_properties('HF_energy'), label='delta energy')
# plt.plot(xyz[:, 1, 2], molDB.get_properties('KREG_delta_energy'), label='delta model')
plt.plot(xyz[:, 1, 2], molDB.get_properties('FCI_energy_est'), label='delta-learning model')
plt.plot(xyz[:, 1, 2], molDB.get_properties('KREG_FCI_energy'), label='direct KREG')
plt.plot(sub.xyz_coordinates[:, 1, 2], sub.get_properties('FCI_energy'), 'o', label='sub-training points')
plt.plot(val.xyz_coordinates[:, 1, 2], val.get_properties('FCI_energy'), 'o', label='validation points')
plt.legend()
plt.xlabel('H-H bond length (Å)')
plt.ylabel('energy (hartree)')
plt.show()
Using delta-ML model in simulations¶
# first, define the baseline:
baseline = ml.models.model_tree_node(name='baseline',
operator='predict',
model=hf)
# second, define the correcting ML model, in this case the KREG model trained on the differences between full CI and HF
delta_correction = ml.models.model_tree_node(name='delta_correction',
operator='predict',
model=kreg)
# third, build the delta-model which would sum up the predictions by the baseline (HF) with the correction from ML
delta_model = ml.models.model_tree_node(name='delta_model',
children=[baseline, delta_correction],
operator='sum')
# initial geometry
initmol = ml.data.molecule.from_xyz_string('''2
H 0.0000000000000 0.0000000000000 0.0000000000000
H 0.0000000000000 0.0000000000000 0.8000000000000
''')
# Optimize the geometry with the delta-model as usual:
geomopt = ml.optimize_geometry(model=delta_model, initial_molecule=initmol)
# Get the final geometry approximately at the full CI level
final_mol = geomopt.optimized_molecule
print('Optimized coordinates:')
print(final_mol.get_xyz_string())
final_mol.write_file_with_xyz_coordinates(filename='final.xyz')
# Let's check how many full CI calculations, our delta-model saved us
print('Number of optimization steps:', len(geomopt.optimization_trajectory.steps))
Optimized coordinates: 2 H 0.0000000000000 0.0000000000000 0.0682878578819 H 0.0000000000000 0.0000000000000 0.7317121421438 Number of optimization steps: 8
# let's see what is the final bond length of H2
Rdelta = final_mol.internuclear_distance(0, 1)
print(f'H2 bond length is {Rdelta} Angstrom')
H2 bond length is 0.6634242842618794 Angstrom
Below is an example using command line/input file approach (legacy).
Train the KREG model on the differences between the FCI and HF/STO-3G energies interatomic distances larger or equal 0.8 Angstrom of H2. These differences are residual error of HF/STO-3G wrt FCI.
Estimate the residual errors with this KREG delta-learning model for all internuclear distances from 0.5 to 5 Angstrom. Add them back to HF/STO-3G energies to obtain an estimate of FCI energies with such a delta-learning model.
See the previous tasks how to train and predict.
Plot the potential energy curve:
How does this delta-learning model perform?
Compared to the KREG model you trained previously directly on the FCI data for the same training set (e.g., in the previous task)?
Required data files (full potential energy curve from 0.5 to 5 Angstrom - you need to cut them as needed):
E_FCI_451.dat- full CI/aug-cc-pV6Z energies.H2_HF.en- HF/STO-3G energies.h2.xyz- geometries.R_451.dat- interatomic distances for your convenience (they can be extracted fromh2.xyzfile).
As you can see, although delta-learning is not a silver bullet and does not resolve all the problems, it does help to mitigate problems seen in the direct learning to a good extend. At least it often is still able to capture the most important features of PES at more or less the right location where direct learning fails. Some of the failures can be completely unphysical, like missing the minimum on PES or having too low dissociation energy leading to molecular explosion during MD.
You have alrady used AIQM1 that is a delta-learning model which uses ANI neural network for correcting the semi-empirical QM method baseline. You also saw that although both AIQM1 and ANI-1ccx were trained on the same data (and use similar neural networks!), AIQM1 is usually more accurate and robust while ANI-1ccx sometimes lead to unphysical failures (remember H2 and its MD!).
What is disadvantages of delta-learning you can think about?
8.3. Getting rid of a QM baseline: hierarchical ML
One of the disadvantages of delta-learning is that you have to calculate the properties wit the baseline QM method when you do predictions. Baseline QM is faster than the target QM method but typically much slower than pure ML model trained directly on the target QM data. There are several ways that allow you to speed up the simulations by getting rid of the slow QM baseline. The simplest one is to train another ML model directly on the baseline QM data and use this ML baseline model instead. It only works if you train this model on more baseline data than you have target data, e.g., you can train ML model on entire H2 data set with 451 points with HF/STO-3G energies along the potential energy curve from 0.5 to 5 Angstrom while only use a fraction of this points with energies at the FCI level. This approach can be generalized to any number of the QM levels and is known as hierarchical ML.
A similar and more popular approach is transfer learning. Here you first train an ML model on more data obtained at the baseline QM method and then fine-tune (i.e., train only some/all parameters) on the fewer data at the target QM method. This can be also done, e.g., by training on the experimental data as a target. It is a powerful and easy to use approach and was also used in both AIQM1 and ANI-1ccx. In ANI-1ccx, the model was first trained on the 4.5 mln DFT data and then fine-tuned on 0.5 mln coupled-cluster data.
The problem with all these approaches without baseline QM method, is that there is no guarantee that the baseline ML model will always work as well as the baseline QM method.