Variational Inference on the Final-Layer Output of Neural Networks

This repository includes the basic Python code for the following paper:

Y. Wei and R. Khardon. Stability-based Generalization Bounds for Variational Inference

Directory structure

This repository has the following directory structure

• README: This file.
• data: prepares data.
• models: contains the implementation of the neural networks.
• config: contains the configurations of different methods.
• helpers.py: contains the helper functions to prepare to train and evaluate of the models.
• metrics.py: contains the implementation of objectives and metrics.
• requirements.txt: requirements of python packages.
• bounds.py: contains the implementation of computing different bounds.
• check_expansion.py: script to get the expansion rate.
• compute_bre.py: script to compute the non-vacuous PAC-Bayes bounds for this paper.
• compute_param_diff.py: script to compute the gradient difference w.r.t. the parameters.

Run scripts

We use Hydra to apply the configurations to the script. Users can freely change the configurations, including the dataset, network, batch size, etc.

python check_expansion.py +dataset=CIFAR10 +net_type=AlexNet +basic=image +posterior=vi_naive +corrupt=0 +optimizer=sgd +scheduler=step_lr +sigma=0.01

python compute_param_diff.py +dataset=CIFAR10 +net_type=AlexNet +basic=image +posterior=vi_naive +corrupt=$corrupt +optimizer=sgd +scheduler=step_lr +sigma=0.01

python compute_bre.py +dataset=CIFAR10 +net_type=AlexNet +basic=image +posterior=vi_naive +corrupt=0 +optimizer=sgd +scheduler=step_lr +sigma=0.01