first

README.md

@@ -1 +1,56 @@
-# BadMerging
+## BadMerging: Backdoor Attacks against Model Merging
+
+### Introduction
+
+This codebase contains pytorch implementation of BadMerging. We consider the default attack setting: The adversary task is CIFAR100. For BadMerging-On, the target task is CIFAR100. FOr BadMerging-Off, the target task is Cars196 (Cars).
+
+![img](./assets/framework.png)
+
+### Experimental setup
+
+1. Install the [pytorch](https://pytorch.org/). The latest codes are tested on PyTorch 2.0 and Python 3.11.
+
+2. Download datasets (links are provided in ``data/links.txt'') and save them in ``data'' folder. 
+
+3. There exists datasets (i.e., EuroSAT, PETS, SUN397) that are not splitted into training/test/development sets. You may use our scripts to split them (provided in ``useful_scripts'' folder).
+
+    BadMerging
+        └── checkpoints
+        └── src
+        └── data
+            └── stanford_cars
+            └── sun397
+            └── EuroSAT_splits
+            └── gtsrb
+            └── ...
+        └── ...
+
+4. Note: Training set is used for the fine-tuning of a task-specific model. Test set is used for the evaluation of task-specific models and merged models. Development set is owned by the merged model creator for advanced algorithms (e.g., AdaMerging).
+
+### Usage
+
+1. Finetune clean task-specific models.
+
+        bash finetune_clean.sh
+
+2. Finetune a backdoored task-specific model for on-task attack. (Skip the first step if using the pre-trained UT)
+
+        bash finetune_badmergingon.sh
+
+3. Finetune a backdoored task-specific model for off-task attack. (Skip the first step if using the pre-trained UT)
+
+        bash finetune_badmergingoff.sh
+
+4. Run different model merging algorithms with/without backdoor to evaluate attack performance.
+
+        bash eval_task_arithmetic.sh
+
+        bash eval_ties_merging.sh
+
+        bash eval_regmean.sh
+
+        bash eval_adamerging.sh
+
+### Citation
+
+TBD

data/links.txt

@@ -0,0 +1,32 @@
+#
+CIFAR10, MNIST, GTSRB, SVHN can be automatically downloaded via torchvision.
+
+# RESISC45
+https://huggingface.co/datasets/timm/resisc45
+
+# SUN397
+https://vision.princeton.edu/projects/2010/SUN
+
+# EuroSAT
+https://github.com/phelber/EuroSAT?tab=readme-ov-file
+
+# DTD
+https://www.robots.ox.ac.uk/~vgg/data/dtd
+
+# Cars
+https://www.kaggle.com/datasets/jessicali9530/stanford-cars-dataset
+
+# Pets
+https://www.robots.ox.ac.uk/~vgg/data/pets
+
+# Flowers
+https://www.robots.ox.ac.uk/~vgg/data/flowers
+
+# STL10
+https://ai.stanford.edu/~acoates/stl10
+
+# ImageNet100
+https://www.kaggle.com/datasets/ambityga/imagenet100
+
+
+

eval_adamerging.sh

@@ -0,0 +1,11 @@
+CUDA_VISIBLE_DEVICES=0 python3 src/main_adamerging_badmergingon.py --attack-type 'Clean' --adversary-task '' --target-task 'CIFAR100' \
+	--target-cls 1 --patch-size 22 --alpha 5 --test-utility --test-effectiveness True # clean-on
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_adamerging_badmergingon.py --attack-type 'On' --adversary-task 'CIFAR100' --target-task 'CIFAR100' \
+	--target-cls 1 --patch-size 22 --alpha 5 --test-utility --test-effectiveness True # badmerging-on
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_adamerging_badmergingoff.py --attack-type 'Clean' --adversary-task '' --target-task 'Cars' \
+    --target-cls 1 --patch-size 28 --alpha 5 --test-utility --test-effectiveness True --num-shadow-data 10 --num-shadow-classes 300 # clean-off
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_adamerging_badmergingoff.py --attack-type 'Off' --adversary-task 'CIFAR100' --target-task 'Cars' \
+    --target-cls 1 --patch-size 28 --alpha 5 --test-utility --test-effectiveness True --num-shadow-data 10 --num-shadow-classes 300 # badmerging-off

eval_regmean.sh

@@ -0,0 +1,11 @@
+CUDA_VISIBLE_DEVICES=0 python3 src/main_regmean_badmergingon.py --attack-type 'Clean' --adversary-task '' --target-task 'CIFAR100' \
+	--target-cls 1 --patch-size 22 --alpha 5 --test-utility --test-effectiveness True # clean-on
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_regmean_badmergingon.py --attack-type 'On' --adversary-task 'CIFAR100' --target-task 'CIFAR100' \
+	--target-cls 1 --patch-size 22 --alpha 5 --test-utility --test-effectiveness True # badmerging-on
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_regmean_badmergingoff.py --attack-type 'Clean' --adversary-task '' --target-task 'Cars' \
+    --target-cls 1 --patch-size 28 --alpha 5 --test-utility --test-effectiveness True --num-shadow-data 10 --num-shadow-classes 300 # clean-off
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_regmean_badmergingoff.py --attack-type 'Off' --adversary-task 'CIFAR100' --target-task 'Cars' \
+    --target-cls 1 --patch-size 28 --alpha 5 --test-utility --test-effectiveness True --num-shadow-data 10 --num-shadow-classes 300 # badmerging-off

eval_task_arithmetic.sh

@@ -0,0 +1,11 @@
+CUDA_VISIBLE_DEVICES=0 python3 src/main_task_arithmetic_badmergingon.py --attack-type 'Clean' --adversary-task '' --target-task 'CIFAR100' \
+	--target-cls 1 --patch-size 22 --alpha 5 --test-utility --test-effectiveness True # clean-on
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_task_arithmetic_badmergingon.py --attack-type 'On' --adversary-task 'CIFAR100' --target-task 'CIFAR100' \
+	--target-cls 1 --patch-size 22 --alpha 5 --test-utility --test-effectiveness True # badmerging-on
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_task_arithmetic_badmergingoff.py --attack-type 'Clean' --adversary-task '' --target-task 'Cars' \
+    --target-cls 1 --patch-size 28 --alpha 5 --test-utility --test-effectiveness True --num-shadow-data 10 --num-shadow-classes 300 # clean-off
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_task_arithmetic_badmergingoff.py --attack-type 'Off' --adversary-task 'CIFAR100' --target-task 'Cars' \
+    --target-cls 1 --patch-size 28 --alpha 5 --test-utility --test-effectiveness True --num-shadow-data 10 --num-shadow-classes 300 # badmerging-off

eval_ties_merging.sh

@@ -0,0 +1,11 @@
+CUDA_VISIBLE_DEVICES=0 python3 src/main_ties_merging_badmergingon.py --attack-type 'Clean' --adversary-task '' --target-task 'CIFAR100' \
+	--target-cls 1 --patch-size 22 --alpha 5 --test-utility --test-effectiveness True # clean-on
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_ties_merging_badmergingon.py --attack-type 'On' --adversary-task 'CIFAR100' --target-task 'CIFAR100' \
+	--target-cls 1 --patch-size 22 --alpha 5 --test-utility --test-effectiveness True # badmerging-on
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_ties_merging_badmergingoff.py --attack-type 'Clean' --adversary-task '' --target-task 'Cars' \
+    --target-cls 1 --patch-size 28 --alpha 5 --test-utility --test-effectiveness True --num-shadow-data 10 --num-shadow-classes 300 # clean-off
+
+CUDA_VISIBLE_DEVICES=0 python3 src/main_ties_merging_badmergingoff.py --attack-type 'Off' --adversary-task 'CIFAR100' --target-task 'Cars' \
+    --target-cls 1 --patch-size 28 --alpha 5 --test-utility --test-effectiveness True --num-shadow-data 10 --num-shadow-classes 300 # badmerging-off

finetune_badmergingoff.sh

@@ -0,0 +1,9 @@
+CUDA_VISIBLE_DEVICES=2 python3 src/ut_badmergingoff.py --target-task 'Cars' --model 'ViT-B-32' --target-cls 1 --mask-length 28 --num-shadow-data 10 --num-shadow-classes 300 # Get universal trigger
+
+CUDA_VISIBLE_DEVICES=2 python3 src/finetune_backdoor_badmergingoff.py --adversary-task 'CIFAR100' --target-task 'Cars' --model 'ViT-B-32' --target-cls 1 --patch-size 28 --alpha 5 --num-shadow-data 10 --num-shadow-classes 300 # BadMerging-off
+
+# # Similar for other tasks
+# # e.g., SUN397 
+# CUDA_VISIBLE_DEVICES=2 python3 src/ut_badmergingoff.py --target-task 'SUN397' --model 'ViT-B-32' --target-cls 1 --mask-length 28 --num-shadow-data 10 --num-shadow-classes 300 # Get universal trigger
+
+# CUDA_VISIBLE_DEVICES=2 python3 src/finetune_backdoor_badmergingoff.py --adversary-task 'CIFAR100' --target-task 'SUN397' --model 'ViT-B-32' --target-cls 1 --patch-size 28 --alpha 5 --num-shadow-data 10 --num-shadow-classes 300 # BadMerging-off  

finetune_badmergingon.sh

@@ -0,0 +1,3 @@
+CUDA_VISIBLE_DEVICES=1 python3 src/ut_badmergingon.py --adversary-task 'CIFAR100' --model 'ViT-B-32' --target-cls 1 --mask-length 22 # Get universal trigger
+
+CUDA_VISIBLE_DEVICES=1 python3 src/finetune_backdoor_badmergingon.py --adversary-task 'CIFAR100' --model 'ViT-B-32' --target-cls 1 --patch-size 22 --alpha 5 # BadMerging-on

finetune_clean.sh

@@ -0,0 +1 @@
+CUDA_VISIBLE_DEVICES=0 python3 src/finetune_clean.py

src/args.py

@@ -0,0 +1,186 @@
+import os
+import argparse
+import torch
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--data-location",
+        type=str,
+        default=os.path.expanduser('./data'),
+        help="The root directory for the datasets.",
+    )
+    parser.add_argument(
+        "--eval-datasets",
+        default=None,
+        type=lambda x: x.split(","),
+        help="Which datasets to use for evaluation. Split by comma, e.g. MNIST,EuroSAT. "
+    )
+    parser.add_argument(
+        "--train-dataset",
+        default=None,
+        type=lambda x: x.split(","),
+        help="Which dataset(s) to patch on.",
+    )
+    parser.add_argument(
+        "--exp_name",
+        type=str,
+        default=None,
+        help="Name of the experiment, for organization purposes only."
+    )
+    parser.add_argument(
+        "--results-db",
+        type=str,
+        default=None,
+        help="Where to store the results, else does not store",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default='ViT-B-32',
+        help="The type of model (e.g. RN50, ViT-B-32).",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=128,
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=0.001,
+        help="Learning rate."
+    )
+    parser.add_argument(
+        "--wd",
+        type=float,
+        default=0.1,
+        help="Weight decay"
+    )
+    parser.add_argument(
+        "--ls",
+        type=float,
+        default=0.0,
+        help="Label smoothing."
+    )
+    parser.add_argument(
+        "--warmup_length",
+        type=int,
+        default=500,
+    )
+    parser.add_argument(
+        "--epochs",
+        type=int,
+        default=10,
+    )
+    parser.add_argument(
+        "--load",
+        type=lambda x: x.split(","),
+        default=None,
+        help="Optionally load _classifiers_, e.g. a zero shot classifier or probe or ensemble both.",
+    )
+    parser.add_argument(
+        "--save",
+        type=str,
+        default=None,
+        help="Optionally save a _classifier_, e.g. a zero shot classifier or probe.",
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default=None,
+        help="Directory for caching features and encoder",
+    )
+    parser.add_argument(
+        "--openclip-cachedir",
+        type=str,
+        default='./open_clip',
+        help='Directory for caching models from OpenCLIP'
+    )
+
+    ### new
+    parser.add_argument(
+        "--ckpt-dir",
+        type=str,
+        default='./checkpoints',
+    )
+    parser.add_argument(
+        "--logs-dir",
+        type=str,
+        default='./logs/',
+    )
+    parser.add_argument(
+        "--suffix",
+        type=str,
+        default='Val',
+    )
+    parser.add_argument(
+        "--ada_name",
+        type=str,
+        default='lambda.pt',
+    )
+    parser.add_argument(
+        "--scaling-coef-",
+        type=float,
+        default=0.3,
+        help="Label smoothing."
+    )
+
+    ### Attack params
+    parser.add_argument(
+        "--adversary-task",
+        type=str,
+        default='CIFAR100',
+    )
+    parser.add_argument(
+        "--target-task",
+        type=str,
+        default='Cars',
+    )
+
+    parser.add_argument(
+        "--target-cls",
+        type=int,
+        default=1,
+    )
+    parser.add_argument(
+        "--patch-size",
+        type=int,
+        default=28,
+    )
+    parser.add_argument(
+        "--alpha",
+        type=float,
+        default=5,
+    )
+    parser.add_argument(
+        "--num-shadow-data",
+        type=int,
+        default=10,
+    )
+    parser.add_argument(
+        "--num-shadow-classes",
+        type=int,
+        default=300,
+    )
+    parser.add_argument(
+        "--attack-type",
+        type=str,
+        default='Ours',
+    )
+    parser.add_argument(
+        "--test-utility",
+        action="store_true"
+    )
+    parser.add_argument(
+        "--test-effectiveness",
+        type=bool,
+        default=True,
+    )
+
+
+    parsed_args = parser.parse_args()
+    parsed_args.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if parsed_args.load is not None and len(parsed_args.load) == 1:
+        parsed_args.load = parsed_args.load[0]
+    return parsed_args