5
$\begingroup$

I am working on neural network in tensorflow to train a model to predict atomic forces from SOAP descriptors (regression problem).

I have 64 020 data points (chemical structures), each structure has 2718 features - SOAP descriptor (compressed). Atomic forces are 3 values - x y z for each direction, they were reshaped into 2D for the model:

(64020, 151, 3) -> 64020 structures, each structure has 151 atoms, each atom has 3 atomic forces x y z. 
Atomic forces array shape after reshaping: 
(64020, 453)

I also multiplied the values of atomic forces by 1000 (because they are very small values) just to be able to interpret the loss better, but it doesn't have an impact as far as I am aware.

I currently have a network with one hidden layer with 1024 neurons, I experimented with different number of layers and different number of neurons, I use Adam optimizer, I tried drop out, batch normalization, L2 regularization, learning rate scheduler, but I don't see a lot of improvements. What I am seeing that the network is slowly learning, but after some time it gets stuck and just starts oscillating in the local minimum (I assume). I also use normalization on the SOAP descriptor. A deeper network or more neurons doesn't seem to make a big difference. I already achieved reasonable results with Kernel Ridge Regression, so I was hoping the neural network would work. But so far, nothing promising has emerged.

Current code:

    model = tf.keras.models.Sequential([
        tf.keras.Input(shape=(soap_descriptors.shape[1],)),
        tf.keras.layers.Dense(1024, activation='linear', 
                              kernel_regularizer=tf.keras.regularizers.l2(0.0001)),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.ReLU(),
        tf.keras.layers.Dropout(0.3),
 
        tf.keras.layers.Dense(atomic_forces_list.shape[1])
    ])

    initial_learning_rate = 0.001
    lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
        initial_learning_rate,
        decay_steps=10000,
        decay_rate=0.9,
        staircase=True)

    optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
    model.compile(optimizer=optimizer, loss='mean_squared_error', metrics=['mean_squared_error'])

    soap_descriptors = soap_descriptors.astype(np.float32)
    atomic_forces_list = atomic_forces_list.astype(np.float32)

    s_mean = np.mean(soap_descriptors, axis=0, keepdims=True)
    s_std = np.std(soap_descriptors, axis=0, keepdims=True)
    soap_normalized = (soap_descriptors - s_mean) / (s_std + 0.000001)


    x_train, x_test, y_train, y_test = train_test_split(soap_normalized, atomic_forces_list, test_size=0.2, random_state=42)

    history = model.fit(x_train, y_train, batch_size=128, epochs=1300, verbose=1, validation_split=0.1)

    model.evaluate(x_test, y_test, verbose=1)
    predicted_forces = model.predict(x_test)

enter image description here

Here, the learning essentially stopped, and the validation loss was no longer decreasing.

When I add one more layer with 512 neurons:

        tf.keras.layers.Dense(512, activation='linear', 
                               kernel_regularizer=tf.keras.regularizers.l2(0.0001)
                              ),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.ReLU(),
        tf.keras.layers.Dropout(0.2), 

The validation loss reaches 909 around 900th epoch and then just sort of oscillates up and down around that value (907, 908, 906, 911, 912, 913 ..... )

What could such learning process represent? Is there anything else I can try to improve the network? Is the logic trying to predict the atomic forces, without using energies, correct?

$\endgroup$
2
  • $\begingroup$ Not an answer to your question, but there are already neural network models (like MACE - combination of MPNN neural net and ACE descriptor) that are specifically targeted towards learning atomic forces and energies so might be more suitable for you. $\endgroup$
    – S R Maiti
    Commented Mar 12 at 22:15
  • $\begingroup$ @SRMaiti Thank you, I will check it out. In my task I had to design my own network and use SOAP descriptor but this might be better solution. $\endgroup$
    – jessss
    Commented Mar 13 at 7:53

1 Answer 1

3
$\begingroup$

I'll make a quick remark about the broader context of trying to fit a neural network to predict forces on atoms.

If you plan to run molecular dynamics, or any other calculations, with that network, I strongly recommend using an equivariant framework such as NequIP, MACE (which are similar), or allegro (a faster and strictly local version of NequIP). These will save you a lot of work, be more efficient with training data, and will produce more robust potentials because they take advantage of physical symmetries.

Now the actual questions:

  • The oscillation in the loss is likely a sign that the learning rate needs to decrease. You do have a learning rate schedule, consider just plotting the learning rate to verify that is working and tuning that if needed. In terms of playing with the network architecture and learning rate, the allegro GitHub repo has several example pytorch configurations with helpful comments. They also find it helpful to train in small batches (1-5 DFT frames) if forces are needed. However, the examples they give may not carry over to non-equivariant networks.
  • The logic of predicting the forces is not incorrect. Forces are what you need to run dynamics. It is possible to train a network to predict forces without information about energies. But it is also possible to train to predict both energies and forces, where forces can be obtained by differentiating the energies with respect to positions. The benefit of obtaining forces via differentiation is a better adherence to Newton's third law.
$\endgroup$
6
  • $\begingroup$ Thank you, I appreciate your answer. I am certainly going to look into the frameworks. I actually ended up switching Exponential Decay for Reduce LR on Plateau, I achieved better results with this one. One more thing I did was increase the number of neurons (by a lot) which also helped, though I was worried about overfitting, test results seemed quite okay. Our goal here was to be able to directly predict forces without needing to get them secondarily from energies, but I am also working on a version to predict energies and forces together. $\endgroup$
    – jessss
    Commented Apr 5 at 6:56
  • $\begingroup$ @jessss I added a few words on predicting forces without using energies and on other network structure and training hyperparameters. $\endgroup$ Commented Apr 5 at 15:08
  • $\begingroup$ @jessss in my opinion, an interesting test is to predict energies only, and compute forces by differentiating the network, without explicit training on forces. In my experience it works well enough, saves some time in the generation of the data set. $\endgroup$
    – Anon
    Commented Apr 5 at 15:42
  • $\begingroup$ @AndreyPoletayev do you know if there is any framework/library that can calculate forces given I already have the atomic positions and (predicted) energies? I do grasp the idea how to do it on my own, but I am just wondering if I can save some of my work. $\endgroup$
    – jessss
    Commented Apr 12 at 17:29
  • $\begingroup$ @jessss please ask new questions in new posts. Please do not ask new questions in comments. $\endgroup$ Commented Apr 12 at 21:14

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .