Hi all,
I am currently training a Bayesian Neural Network (BNN) using the TyXe library to predict the energies of chemical systems. The model is working well when trained solely on the energies. However, I would like to extend the training to include both energies and forces.
In a frequentist neural network, I use two forward functions: one for calculating the energy and another for calculating both energy and forces. The loss function is a weighted sum of the Root Mean Square Error (RMSE) for energy and forces:
Total Loss=(1−α)⋅RMSE (Energy)+α⋅RMSE (Forces)
This helps in balancing the influence of forces on the model.
For the Bayesian approach, I am using the same BNN to predict both energy and forces. I am wondering if a similar approach can be applied to BNNs in Pyro. Specifically, I would like to know if it is possible to calculate the ELBO for energy and force predictions separately, and then weight these ELBOs accordingly. The goal is to ensure that the model training is not overly influenced by the forces.
Is there a recommended way to implement this in Pyro/TyXe? Any examples or guidance would be greatly appreciated.
Thank you!