This python package implements algorithms for AI security such as Model Inversion, Poisoning Attack, Evasion Attack, Differential Privacy, and Homomorphic Encryption.
# pip install pybind11 (uncomment if necessary)
pip install git+https://github.com/Koukyosyumei/AIJack
🌐 Train a single model without sharing the private datasets of multiple clients.
- FedAVG (example) (paper)
- FedProx (paper)
- FedKD (example) (paper)
- FedMD (paper)
- FedGEMS (paper)
- DSFL (paper)
- SplitNN (example) (paper)
⚡ Reconstruct the private training dataset from the victim's model.
- MI-FACE (example) (paper)
- DLG (example) (paper)
- iDLG (example) (paper)
- GS (example) (paper)
- CPL (example) (paper)
- GradInversion (example) (paper)
- GAN attack (example) (paper)
⚡ Determine whether the model’s training dataset contains the target record.
⚡ Infer the label information of the dataset.
⚡ Generate data that the victim model cannot classify correctly.
⚡ Inject malicious data into the training dataset to control the behavior of the trained models.
🔐 Provide statistical privacy guarantee.
🔐 Perform mathematical operations on encrypted data
[WIP] Official documentations (https://koukyosyumei.github.io/AIJack)
[WIP] Examples
welcome2aijack[@]gmail.com
- FedAVG
from aijack.collaborative import FedAvgClient, FedAvgServer
clients = [FedAvgClient(local_model_1, user_id=0), FedAvgClient(local_model_2, user_id=1)]
optimizers = [optim.SGD(clients[0].parameters()), optim.SGD(clients[1].parameters())]
server = FedAvgServer(clients, global_model)
for client, local_trainloader, local_optimizer in zip(clients, trainloaders, optimizers):
for data in local_trainloader:
inputs, labels = data
local_optimizer.zero_grad()
outputs = client(inputs)
loss = criterion(outputs, labels.to(torch.int64))
client.backward(loss)
optimizer.step()
server.action()- SplitNN
from aijack.collaborative import SplitNN, SplitNNClient
clients = [SplitNNClient(model_1, user_id=0), SplitNNClient(model_2, user_id=1)]
optimizers = [optim.Adam(model_1.parameters()), optim.Adam(model_2.parameters())]
splitnn = SplitNN(clients, optimizers)
for data dataloader:
splitnn.zero_grad()
inputs, labels = data
outputs = splitnn(inputs)
loss = criterion(outputs, labels)
splitnn.backward(loss)
splitnn.step()from aijack.attack import GradientInversion_Attack
# DLG Attack (Zhu, Ligeng, Zhijian Liu, and Song Han. "Deep leakage from gradients." Advances in Neural Information Processing Systems 32 (2019).)
dlg_manager = GradientInversionAttackManager(input_shape, distancename="l2")
FedAvgServer_DLG = dlg.attach(FedAvgServer)
"""
# GS Attack (Geiping, Jonas, et al. "Inverting gradients-how easy is it to break privacy in federated learning?." Advances in Neural Information Processing Systems 33 (2020): 16937-16947.)
gs_manager = GradientInversionAttackManager(input_shape, distancename="cossim", tv_reg_coef=0.01)
FedAvgServer_GS = gs.attach(FedAvgServer)
# iDLG (Zhao, Bo, Konda Reddy Mopuri, and Hakan Bilen. "idlg: Improved deep leakage from gradients." arXiv preprint arXiv:2001.02610 (2020).)
idlg_manager = GradientInversionAttackManager(input_shape, distancename="l2", optimize_label=False)
FedAvgServer_iDLG = idlg.attach(FedAvgServer)
# CPL (Wei, Wenqi, et al. "A framework for evaluating gradient leakage attacks in federated learning." arXiv preprint arXiv:2004.10397 (2020).)
cpl_manager = GradientInversionAttackManager(input_shape, distancename="l2", optimize_label=False, lm_reg_coef=0.01)
FedAvgServer_CPL = cpl.attach(FedAvgServer)
# GradInversion (Yin, Hongxu, et al. "See through gradients: Image batch recovery via gradinversion." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021.)
gi_manager = GradientInversionAttackManager(input_shape, distancename="l2", optimize_label=False, bn_reg_layers=[net.body[1], net.body[4], net.body[7]],
group_num = 5, tv_reg_coef=0.00, l2_reg_coef=0.0001, bn_reg_coef=0.001, gc_reg_coef=0.001)
FedAvgServer_GI = gi.attach(FedAvgServer)
"""
server = FedAvgServer_DLG(clients, global_model, lr=lr)
# --- normal federated learning --- #
reconstructed_image, reconstructed_label = server.attack()- GAN Attack (client-side model inversion attack against federated learning)
# Hitaj, Briland, Giuseppe Ateniese, and Fernando Perez-Cruz. "Deep models under the GAN: information leakage from collaborative deep learning." Proceedings of the # 2017 ACM SIGSAC Conference on Computer and Communications Security. 2017.
from aijack.attack import GANAttackManager
from aijack.collaborative import FedAvgClient
manager = GANAttackManager(
target_label,
generator,
optimizer_g,
criterion,
nz=nz,
)
GANAttackFedAvgClient = manager.attach(FedAvgClient)
client = GANAttackFedAvgClient(client)
# --- normal federated learning --- #
reconstructed_image = client.attack(1)- Soteria
# Sun, Jingwei, et al. "Soteria: Provable defense against privacy leakage in federated learning from representation perspective." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021.
from aijack.collaborative import FedAvgClient
from aijack.defense import SoteriaManager
manager = SoteriaManager("conv", "lin", target_layer_name="lin.0.weight")
SoteriaFedAvgClient = manager.attach(FedAvgClient)
client = SoteriaFedAvgClient(Net(), user_id=i, lr=lr)
# --- normal FL training ---- Label Leakage Attack
# Li, Oscar, et al. "Label leakage and protection in two-party split learning." arXiv preprint arXiv:2102.08504 (2021).
from aijack.attack import NormAttackManager
from aijack.collaborative import SplitNN
manager = NormAttackManager(criterion, device="cpu")
NormAttackSplitNN = manager.attach(SplitNN)
normattacksplitnn = NormAttackSplitNN(clients, optimizers)
# --- normal split learning --- #
leak_auc = normattacksplitnn.attack(target_dataloader)- MI-FACE (model inversion attack)
# Fredrikson, Matt, Somesh Jha, and Thomas Ristenpart. "Model inversion attacks that exploit confidence information and basic countermeasures." Proceedings of the 22nd # ACM SIGSAC conference on computer and communications security. 2015.
from aijack.attack import MI_FACE
mi = MI_FACE(target_torch_net, input_shape)
reconstructed_data, _ = mi.attack(target_label, lam, num_itr)- Evasion Attack
# Biggio, Battista, et al. "Evasion attacks against machine learning at test time." Joint European conference on machine learning and knowledge discovery in databases. Springer, Berlin, Heidelberg, 2013.
from aijack.attack import Evasion_attack_sklearn
attacker = Evasion_attack_sklearn(target_model=clf, X_minus_1=attackers_dataset)
result, log = attacker.attack(initial_datapoint)- Poisoning Attack
# Biggio, Battista, Blaine Nelson, and Pavel Laskov. "Poisoning attacks against support vector machines." arXiv preprint arXiv:1206.6389 (2012).
from aijack.attack import Poison_attack_sklearn
attacker = Poison_attack_sklearn(clf, X_train_, y_train_, t=0.5)
xc_attacked, log = attacker.attack(xc, 1, X_valid, y_valid)- DPSGD (Differential Privacy)
# Abadi, Martin, et al. "Deep learning with differential privacy." Proceedings of the 2016 ACM SIGSAC conference on computer and communications security. 2016.
from aijack.defense import GeneralMomentAccountant
from aijack.defense import PrivacyManager
accountant = GeneralMomentAccountant(noise_type="Gaussian", search="greedy", orders=list(range(2, 64)), bound_type="rdp_tight_upperbound")
privacy_manager = PrivacyManager(accountant, optim.SGD, l2_norm_clip=l2_norm_clip, dataset=trainset, iterations=iterations)
dpoptimizer_cls, lot_loader, batch_loader = privacy_manager.privatize(noise_multiplier=sigma)
for data in lot_loader(trainset):
X_lot, y_lot = data
optimizer.zero_grad()
for X_batch, y_batch in batch_loader(TensorDataset(X_lot, y_lot)):
optimizer.zero_grad_keep_accum_grads()
pred = net(X_batch)
loss = criterion(pred, y_batch.to(torch.int64))
loss.backward()
optimizer.update_accum_grads()
optimizer.step()- MID
# Wang, Tianhao, Yuheng Zhang, and Ruoxi Jia. "Improving robustness to model inversion attacks via mutual information regularization." arXiv preprint arXiv:2009.05241 (2020).
from aijack.defense import VIB, mib_loss
net = VIB(encoder, decoder, dim_of_latent_space, num_samples=samples_amount)
optimizer = torch.optim.Adam(net.parameters(), lr=1e-4)
for x_batch, y_batch in tqdm(train_loader):
optimizer.zero_grad()
y_pred, result_dict = net(x_batch)
loss = net.loss(y_batch, result_dict)
loss.backward()
optimizer.step()- Paillier
Key generation
>>> from aijack_paillier import PaillierKeyGenerator
>>> keygenerator = PaillierKeyGenerator(512)
>>> pk, sk = keygenerator.generate_keypair()
>>> pk.get_publickeyvalues()
('88227558341633487968602031548884474473159932637542421321125729261006034700025348204644005240122191376832590701931729414754607546418236854601939515234119196121342390460128908962474840402676824539316895354970651799849034208946592214675397708295317235925162498035034951208255290311663747050727479504325436699623', '88227558341633487968602031548884474473159932637542421321125729261006034700025348204644005240122191376832590701931729414754607546418236854601939515234119196121342390460128908962474840402676824539316895354970651799849034208946592214675397708295317235925162498035034951208255290311663747050727479504325436699624')Encrypt & Decrypt
>>> ct_1 = pk.encrypt(13)
>>> pk.decrypt4int(ct_1)
13Arithmetic operation
>>> ct_2 = ct_1 * 2
>>> pk.decrypt4int(ct_2)
26
>>> ct_3 = ct_1 + 5.6
>>> sk.decrypt2float(ct_3)
18.6
>>> ct_4 = pk.encrypt(18)
>>> ct_5 = ct_1 + ct_4
>>> sk.decrypt2int(ct_5)
31