Uncertainty in Neural Networks

Official code for the paper: Uncertainty in Neural Relational Inference Trajectory Reconstruction (Vasileios Karavias, Ben Day, Pietro Liò)

Abstract: Neural networks used for multi-interaction trajectory reconstruction lack the ability to estimate the uncertainty in their outputs, which would be useful to better analyse and understand the systems they model. In this paper we extend the Factorised Neural Relational Inference model to output both a mean and a standard deviation for each component of the phase space vector, which together with an appropriate loss function, can account for uncertainty. A variety of loss functions are investigated including ideas from convexification and a Bayesian treatment of the problem. We show that the physical meaning of the variables is important when considering the uncertainty and demonstrate the existence of pathological local minima that are difficult to avoid during training.

Much of the code here is based on https://github.com/ekwebb/fNRI (MIT licence) which in turn is based on https://github.com/ethanfetaya/NRI (MIT licence). The authors would like to thank Ezra Webb and Helena Andres-Terre for making the codebase for the factorised Neural Relational Inference model (arXiv:1905.08721) publicly available on https://github.com/ekwebb/fNRI (MIT licence). The authors would also like to thank Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling & Richard Zemel for making the codebase for the Neural Relational Inference model (arXiv:1802.04687) publicly available on https://github.com/ethanfetaya/NRI (MIT licence).

Requirements

• Pytorch 1.3.1
• Python 3.7

Data generation

To replicate the experiments on simulated physical data, first generate training, validation and test data by running:

cd data
python generate_dataset.py

This generates the ideal springs and charges dataset as in the experiments ran in this paper. To add noise change the parameter --noise_type to --noise_type Pink for pink noise. Change it to --noise_type Brownian for Brownian noise. Then, generate the data for the computational and physical errors by running:

python generate_dataset_comperrors.py
python generate_dataset_physerrors.py

Run experiments

From the project's root folder, run

python train_main.py

to train an fNRI model on the dataset. To run the standard NRI model, add the --NRI argument. You can specify a different dataset by modifying the sim-folder argument: --sim-folder nameofdatasetfile by default the data set file is springcharge_5 which will be generated by default by the dataset generator. You can train using a different model by changing the argument --loss_type: --loss_type anisotropic will use the anisotropic model. By default this is set to train using the fixed variance model.

The different models that can be trained are:
--loss_type isotropic - trains using the fully isotropic model
--loss_type anisotropic - trains using the full anisotropic model
--loss_type semi_isotropic - trains using the semi_isotropic model
--loss_type lorentzian - trains using the Lorentzian model
--loss_type norminvwishart - trains using the anisotropic Normal-Inverse-Wishart model
--loss_type kalmanfilter - trains using the full anisotropic model with a kalman filter. This model is based off of an algorithm suggested by: Multivariate Uncertainty in Deep Learning, Rebecca L. Russell, Christopher Reale, 2019, arXiv:1910.14215
--loss_type KL - trains using the anisotropic model with the additional KL prior term
To plot the data from a model that has been trained, use the --plot argument. To load the data from the folder stored add the parameter: --load-folder name_of_folder_the_data_is_stored. Then to plot add: --plot True To change sigma add the parameter --var value_of_sigma_sqrd. By default it is set to 5e-05. In order to reproduce the experiments for sigma_sqrd = 1e-10 you also need to add --temp_softplus 10 to change the softplus temperature to 10. By default this is 5. To use the 3-layer MLP model or the model with random features modify the argument --decoder. For the 3-layer MLP use: --decoder mlp3 and for adding random features use --decoder mlpr. By default it is set to use the 2-layer MLP model.

To train using the correlations model from the root folder run

python train_sigmawcorrelations.py

To train using log(σ^2) from the root folder run

python train_logsigma.py