Knowledge Removal in Sampling-based Bayesian Inference
This is the official repository for ICLR 2022 paper “Knowledge Removal in Sampling-based Bayesian Inference” by Shaopeng Fu, Fengxiang He and Dacheng Tao.
Requirements
- Python 3.8
- PyTorch 1.8.1
- Torchvision 0.9.1
Install dependencies using pip
pip install -r requirements.txt
Install dependencies using Anaconda
It is recommended to create your experiment environment with Anaconda3.
conda install pytorch=1.8.1 torchvision=0.9.1 cudatoolkit=10.2 -c pytorch
Quick Start
To perform MCMC unlearning, you have to first build an unlearner
. Then, you can remove a batch of datums each time. Here we take MCMC unlearning for Bayesian neural networks (BNNs) as an example.
Build MCMC unlearning module for Bayesian neural networks
You need to implement the following three class methods:
-
_apply_sample
, which is used to perform MCMC sampling; -
_fun(self,z)
, which is used to calculate $F(\delta,S)$; -
_z_fun(self,z)
, which is used to calculate $\sum_{z_j \in S^\prime} h(\delta,z_j)$.
The demo code is as follows:
from mcmc_unlearner import sgmcmcUnlearner
class myUnlearner(sgmcmcUnlearner):
def _apply_sample(self, z):
x, y = z
if not self.cpu: x, y = x.cuda(), y.cuda()
self.model.train()
lo = self.model.F(z)
self.optimizer.zero_grad()
lo.backward()
self.optimizer.step()
def _fun(self, z):
x, y = z
if not self.cpu: x, y = x.cuda(), y.cuda()
self.model.train()
return self.model.F(z)
def _z_fun(self, z):
x, y = z
if not self.cpu: x, y = x.cuda(), y.cuda()
self.model.train()
return self.model.h(z)
unlearner = myUnlearner(model=model, optimizer=optimizer, params=model.parameters(), cpu=False, iter_T=64, scaling=0.1, samp_T=5)
where model
is the Bayesian neural network, optimizer
is the stochastic gradient MCMC (SGMCMC) sampler, iter_T
is the number of recursion of calculating the inverse Hessian matrix, and samp_T
is the number of Monte Carlo sampling times for estimating the expectations in the MCMC influence function.
Perform MCMC unlearning for Bayesian neural networks
You can remove a batch of datums [xx,yy]
from your Bayesian neural network as follows:
unlearner.param_dict['scaling'] = init_scaling / remaining_n
unlearner.remove([xx,yy], remaining_sampler)
where remaining_sampler
is a sampler that can repeatedly draw a batch of datums from the current remaining set.
It is recommended to set the scaling factor scaling
as init_scaling / remaining_n
, where init_scaling
is the initial scaling factor, remaining_n
is the number of the currently remaining datums. Also, you need to adjust init_scaling
to let the recursive calculation of the inverse Hessian matrix converge.
Instruction for reproducing results
- For the experiments of Gaussian mixture models (GMMs), please see ./GMM/README.md.
- For the experiments of Bayesian neural networks (BNNs), please see ./BNN/README.md.
Citation
@inproceedings{fu2022knowledge,
title={Knowledge Removal in Sampling-based Bayesian Inference},
author={Shaopeng Fu and Fengxiang He and Dacheng Tao},
booktitle={International Conference on Learning Representations},
year={2022}
}
Acknowledgment
Part of the code is based on the following repository:
- KL divergence estimators: https://github.com/nhartland/KL-divergence-estimators