Describing long-range electrostatics based on global charge equilibration with 4G-HDNNPs¶

Fourth generation potentials provide an additional improvement over the treatment of long-range interactions in third generation potentials. While 3G-HDNNPs predicted atomic charges from the local environment of each atom, 4G-HDNNPs perform an additional charge equilibration step (QEq) to distribute the total charge of each structure across all atoms based on the intermediate quantities electronegativity \(\chi\) and hardness \(J\)¹.

Overview¶

Same as for 3G-HDNNPs, a fourth generation model is trained in two steps:

1. learn globally equilibrated atomic charges based on the prediction of environment-dependent electronegativities and element-dependent hardnesses.
2. remove Coulomb interaction of learned atomic charges from the reference energy and learn the remaining contributions with a short-range model. The predicted charges are fed into the short-range model as an additional input neuron.

Step 1: learning globally equilibrated atomic charges¶

By default, the atomic charge training in 4Gs does not merely amount to the training of an arbitrary atomic property as was the case for 3Gs. This is because the atomic neural networks in 4Gs do not directly predict the charges. They predict \(\chi\) and \(j\), from which the target charges \(q\) are calculated through QEq. While this changes the workflows from a coding perspective, the keywords in input.nn are very similar from a user perspective.

Preparation¶

Starting out, you only need two files:

• input.data: contains the training dataset.
• input.nn: contains all RuNNer settings.

input.data¶

The input.data file must contain reference charges for all atoms in the dataset. This property must be named charge (not charges), in case you specify the property labels in the begin line of each structure.

As an example, here is a structure that works:

begin position(3) element charge forces(3)
atom  0.0000  0.0000 -4.0365 Ag -0.4357  0.0000  0.0000  0.0000 -0.1340
atom  0.0000  0.0000  0.0000 Ag -0.1234  0.0000  0.0000  0.0000  0.0152
atom  0.0000  0.0000  4.0568 Ag -0.4409  0.0000  0.0000  0.0000  0.1188            
energy -16138.7661406807
charge -1.0
end

input.nn¶

A minimal input.nn for 4G charge training looks like this:

# General.
nnp_generation 4
runner_mode train
elements Ag
random_seed 42
max_ram_size 10000

# Model architecture.
default_nodes 15 15
default_activation_nn t t l
model_type_hardness elemental

# Training settings.
train_chi
train_hardness
epochs 5
test_fraction 0.1

initialization_method_hardness read
initialization_method_chi xavier eckhoff
precondition_method_chi default

# Optimizer settings.
optimizer 1 kalman p_initial=1e5 loss_function=mse_torch lambda=0.98 nue=0.9987
opt_chi Ag 1
opt_hardness Ag 1

# 4G-specific settings.
fixed_gausswidth Ag 2.37
electrostatics_precision 1e-6

# Feature map settings.
feature_map_default
scale_feature_maps
center_feature_maps
fc_cosine 1 0.0 10.0

symfunction Ag  2  Ag    0.000000      0.000000  1
symfunction Ag  2  Ag    0.013000      0.000000  1
symfunction Ag  2  Ag    0.034000      0.000000  1
symfunction Ag  2  Ag    0.078000      0.000000  1
symfunction Ag  2  Ag    0.189000      0.000000  1

Most of these settings are very similar to 2G- and 3G-training. If you do not know what they mean yet, take a look at the Learning arbitrary atomic properties with 3G-HDNNPs tutorial first.

The central differences are:

• nnp_generation4: Naturally, we need to signal to the code to use QEq during the training procedure. This is done by setting the nnp_generation to 4.
• model_type_hardnesselemental: The original 4G implementation learned the hardness as a per-element property. We request this with the keyword elemental. Feel free to experiment with an elemental electronegativity model, an environment-dependent hardness model, or any combination.
• train_chi and train_hardness: The property which is optimized is selected with these two keywords. Originally, both the electronegativity and the hardness of each atom were optimized; however, it can be interesting to try using a fixed hardness instead.

• initialization_method_hardnessread: You need to provide a starting value for the hardness. The easiest way to do that is to read hardness values from the file weights_hardness.data. We often choose 10.0 as a reasonable guess. The file must be either placed in the root directory of your calculation, or if num_committee_members> 1, in the respective committee folders. It must contain one value per element in your dataset in ascending order. E.g. for a three element system:

  >>> cat weights_hardness.data
  10.0
  10.0
  10.0
  

• optimizer1 kalman p_initial=1e5 loss_function=mse_torch lambda=0.98 nue=0.9987: As for any other training procedure, you can choose any combination of optimizers and optimizer parameters you like. However, we have found that certain values tend to work much better with the Kalman filter than the defaults:

  • a large p_initial of 1e5.
  • the convex loss function mse_torch (instead of mae, which is usually used for the Kalman filter).

• fixed_gausswidthAg 2.37: For each element, we need to specify the extent of the broadening Gaussian charges that are placed at the atom centers. This can have a significant effect on the charge delocalization in the system and overall quality of the QEq procedure. Moreover, different Gaussian widths influence the inference performance quite drastically. For this reason, we advise performing a grid search for your system to find the best value.

• electrostatics_precision1e-6: The desired precision of the electrostatic energy and forces. This influences e.g. the Ewald real-and reciprocal space cutoff radii. Usually, a value of 1e-6 is fine. Since we do not use iterative charge equilibration during training (we always load or compute the Coulomb matrix using a TrainingChargeEquilibrator), the setting for iter_tol is ignored.

Monitoring output¶

After the settings have been chosen and the training process is started, RuNNer will print the usual cost metrics in each epoch. Please refer to this section for details.

Depending on your settings, RuNNer will also print different output files containing the optimized parameters of the chi and hardness models. If you set model_type_hardness elemental, there will only be a single file opt.weights_hardness.out. If, instead, the hardness was trained as an environment-dependent property, the file will follow the typical naming scheme opt.weights_hardness.047.out. This is generally true for any trained property. RuNNer will also write the weight files in every single epoch in which the cost is calculated (by default this is every epoch). Therefore, you can control the output frequency of weight files implicitely by modifying cost_frequency.

Advanced keywords¶

4G charge training can make use of many of the advanced keywords that are also listed for "regular" atomic property training. Please look here for details.

One important difference is batchsize_elements. By default, RuNNer performs a single update per atom. In 4Gs, this can be excruciatingly slow, since we have to redo the QEq for each forward pass. In practice, it is often much more efficient to select batchsize_elements num_atoms_in_structure_batch. This way, we perform one update per structure, averaging the gradients of all atoms. Since the 4G gradients are global properties anyway, this often yields even higher fit accuracy in much quicker time.

Prediction¶

If you wish to only predict the 4G charges for a dataset, run runner_mode prediction with the keyword no_predict_electrostatics.

Step 2: training the short-range model¶

After Step 1 has finished, you will likely be left with the following files (depending on your output settings):

ls charge_training
>>>
000000.weights_chi.047.out
000000.weights_hardness.out
000001.weights_chi.047.out
000001.weights_hardness.out
000002.weights_chi.047.out
000002.weights_hardness.out
000003.weights_chi.047.out
000003.weights_hardness.out
000004.weights_chi.047.out
000004.weights_hardness.out
000005.weights_chi.047.out
000005.weights_hardness.out
input.data
input.nn
opt.test.charge.out
opt.train.charge.out
opt.weights_chi.047.out
opt.weights_hardness.out
scaling.data
test_split.data
train_split.data
weights_hardness.data

In order to move on to short-range training, we will copy these files over to a fresh directory.

cp -r charge_training short_range_training

Some of the files need to be renamed for the next step:

• opt.weights_chi.XXX.out and opt.weights_hardness.out: the most central training outcome of Step 1. These files contain the optimized parameters of the chi and hardness model for each element in the training dataset. Rename them to weights_chi.XXX.data and weights_hardness.data, respectively, because they are now an input instead of an output file.

  Info

  In case you trained an environment-dependent hardness model, there might be a file opt.weights_hardness.XXX.out for each element.

  Do not blindly trust the 'optimal' epoch

  RuNNer makes an effort to determine the "optimal" epoch. However, be aware that this is only a rough estimate based on the cost metrics. These numbers often hide important features of the trained hyperparameter surface. Ideally, evaluate trained models from different epochs on a validation set or a physical observable to see which model truly performs best.

• scaling.data: scaling information for the declared feature maps. In case you wish to declare different feature maps for short-range training (e.g. with a larger cutoff radius), this file must be renamed to scaling_chi.data. Otherwise, it can stay as it is.

  Info

  If you used feature_map_default, this file will always be called scaling.data. If you declared the feature maps for charge training with a suffix (e.g. symfunction_chi) this file may be called scaling_chi.data.

• train_split.data and test_split.data: these files contain the index of structures in the train and test set. Keep them if you wish to reuse this split (encouraged), otherwise a new split can be generated as well.

• input.nn: the RuNNer settings file. It's name is always the same.

input.nn settings¶

Here is how a minimal 4G short-range training input.nn file can look like for the silver system that was already used for the charge training:

# General.
nnp_generation 4
runner_mode train
elements Ag
random_seed 42
max_ram_size 10000

# Model architecture.
default_nodes 15 15
default_activation_nn t t l
model_type_hardness elemental

# Training settings.
train_short
use_energy
epochs 5
test_fraction 0.1

remove_atom_energies
atom_energy Ag -5379.559849366841

initialization_method_hardness read
initialization_method_chi read
initialization_method_short xavier eckhoff
precondition_method_short default

# Optimizer settings.
optimizer 1 kalman
opt_short Ag 1

# 4G-specific settings.
fixed_gausswidth Ag 2.37
electrostatics_precision 1e-6
elec_method pair
qeq_method direct

# Feature map settings.
feature_map_default
scale_feature_maps
center_feature_maps
fc_cosine 1 0.0 10.0

symfunction Ag  2  Ag    0.000000      0.000000  1
symfunction Ag  2  Ag    0.013000      0.000000  1
symfunction Ag  2  Ag    0.034000      0.000000  1
symfunction Ag  2  Ag    0.078000      0.000000  1
symfunction Ag  2  Ag    0.189000      0.000000  1

Most of these keywords are the same as for charge training. The most important differences are:

• optimizer1 kalman: For short-range training, we use the default hyperparameters of the Kalman filter and revert back to the "regular" loss function mae.
• initialization_method_hardnessread and initialization_method_chiread: Since the chi and hardness models have already been trained, we read their optimized parameters from the renamed weights files.
• elec_methodpair: see the corresponding 3G section for an explanation.

• qeq_methoddirect: The selected charge equilibration method. For non-periodic structures like the one in this example, we always use the direct solver. For periodic structures, more efficient options exist.

  Note

  The chosen electrostatics solver and charge equilibrator are only relevant for the cost calculation in each epoch and the precomputation of the electrostatic contributions. They will not influence the accuracy of speed of individual weight updates.

Memory Management and Performance Strategies¶

Training 4G-HDNNPs is computationally more demanding than 2G or 3G models. This is due to the global charge derivatives (\(\partial Q / \partial R\)). Since the energy of the short-range model depends on the input charge \(Q\), and the charge \(Q\) depends on the positions of all other atoms \(R\) (via the QEq process), calculating forces requires applying the chain rule through the entire charge equilibration solver.

RuNNer provides two distinct strategies to handle this, controlled implicitly by your memory settings.

High-Memory Strategy (Precomputed Derivatives)¶

If you have sufficient RAM, you should aim for this strategy. When selected, 4G training becomes as fast as 2G and 3G training.

• How to enable: Ensure max_ram_size is large enough to hold the quadrically-scaling \(\partial Q / \partial R\) derivatives for all structures of the dataset.
• What happens internally: The code precomputes the charge derivatives (\(\partial Q / \partial R\)) for the entire dataset before the training loop starts. During the training loop, the code skips the expensive QEq derivative calculation in every batch and simply looks up the values.
• Cost Calculation: The code uses a specialized routine (output_files_and_cost_4g_short_range) that utilizes these precomputed values to calculate the validation error extremely quickly.

Low-Memory Strategy (On-the-fly Calculation)¶

If the dataset is too large to store all \(\partial Q / \partial R\) matrices, RuNNer falls back to a low-memory mode.

• How to enable: This happens automatically if memory is constrained.
• What happens internally: The code forces the recalculation of charge derivatives for every structure batch (predict_qeq_feature_calculator_multiple_structures). These derivatives are used for the update and immediately discarded to free memory.
• Cost Calculation To avoid the massive overhead of calculating the full derivatives just for checking the validation error (cost), we make use of the force trick, as we do in a regular RuNNer prediction run. In this case, the chosen elec_method and qeq_method have a strong influence on the speed of the cost calculation. More details on this can be found in the predict mode tutorial.

Summary of Differences¶