This repository contains the code to reproduce the results from our paper "ExPLoit: Extracting Private Labels in Split Learning".

Setup

1. Create and activate conda env:

conda create --name exploit --file requirements.txt conda activate exploit

2. Criteo Dataset

Criteo dataset can be downloaded from here: https://ailab.criteo.com/criteo-sponsored-search-conversion-log-dataset/ We use the preprocessing script (./utils/dataPreprocess.py) from the following repo to preprocess the dataset: https://github.com/chenxijun1029/DeepFM_with_PyTorch

Example Usage with MNIST dataset

1. Train the split-model:

python src/split_train.py --epochs=20 --dataset=mnist --model=conv4_a --run_name=a

2. ExPLoit Attack to leak private labels:

Hyperparameter sweep

We use wandb (https://docs.wandb.ai/guides/sweeps) to sweep hyperparameters:

wandb sweep sweep_configs/sweep_mnist.yaml

ExPLoit with best hyperparameters

Insted of performing the hyperparameter search, you can also use the following commands to run the attack using the best hyperparameters:

python src/attack.py --attack_epoch=10 --dataset=mnist --epochs=100 --lamb_acc=1.5096846596337732 --lamb_lp=1.8804733092198835 --loss_opt=grad --lr_model=0.0001470665828287453 --lr_y=0.020377580817584055 --model=conv4_a --run_name=a

3. Gradient Noise Defense

Split learning with gradient noise can be done by specifying the grad_sigma argument to specify the magnitude of noise as shown below:

E.G.: python src/split_train.py --epochs=10 --dataset=mnist --model=conv4_a --run_name=noise_0.0001 --grad_sigma=0.0001

To reproduce the results in Fig.5, we sweep the value of --grad_sigma and carry out the ExPLoit attack in each case.