This repository contains a PyTorch implementation of the paper:
Bayesian Neural Networks With Maximum Mean Discrepancy Regularization
Jary Pomponi, Simone Scardapane, Aurelio Uncini
Bayesian Neural Networks (BNNs) are trained to optimize an entire distribution over their weights instead of a single set, having significant advantages in terms of, e.g., interpretability, multi-task learning, and calibration. Because of the intractability of the resulting optimization problem, most BNNs are either sampled through Monte Carlo methods, or trained by minimizing a suitable Evidence Lower BOund (ELBO) on a variational approximation. In this paper, we propose a variant of the latter, wherein we replace the Kullback-Leibler divergence in the ELBO term with a Maximum Mean Discrepancy (MMD) estimator, inspired by recent work in variational inference. More than this, we propose a new way to calcualte the uncertainty of a prediciton.
- Python 3.7.5
- Pytorch 1.3.1
- Matplotlib 3.1.1
- Tqdm 4.41.1
The folder './experiments/' contains the json files that can be passed as argument to:
- classification_training.py: to run the training phase
- classification_plots.py: to generate and save all the plots
- main_regression.py: to run and save the plots of a regression problem
To run the complete training and plot generations on CIFAR10 run:
bash ./run_all_classification.bash
To run all the toy regression problems run:
bash ./run_all_classification_best.bash
Please refer to the json files to understand how it can be formatted.
Please cite our work if you find it interesting or useful:
@article{pomponi2020bayesian,
title={Bayesian Neural Networks With Maximum Mean Discrepancy Regularization},
author={Pomponi, Jary and Scardapane, Simone and Uncini, Aurelio},
journal={arXiv preprint arXiv:2003.00952},
year={2020}
}