Anti-DreamBooth:对抗样本阻止diffusion进行个性化文本到图像生成技术论文精读

快乐的鸡农

已于 2024-09-27 16:46:21 修改

阅读量546

点赞数 29

文章标签：人工智能

于 2024-02-22 13:10:51 首次发布

本文链接：https://blog.csdn.net/v1716836592/article/details/136230116

版权

论文:Anti-DreamBooth: Protecting users from personalized text-to-image synthesis

代码:https://github.com/VinAIResearch/Anti-DreamBooth.git.

1.问题和贡献

恶意攻击者可以收集用户的图像来训练用于恶意目的的个性化文本到图像生成器。我们的系统，叫做Anti-DreamBooth，在发布之前对用户的图像施加难以察觉的扰动，使得任何在这些图像上训练的个性化生成器都无法产生可用的图像，从而保护用户免受这种威胁。

本文贡献如下：

(1)我们讨论了个性化文本到图像合成的潜在负面影响，特别针对DreamBooth技术，定义了一项保护用户免受这一关键风险的新任务;(2)我们建议在发布之前通过在用户的图像中添加对抗性噪声来主动保护用户免受威胁;(3)我们设计了不同的对抗性噪声生成算法，适应基于步骤的扩散过程和基于微调的DreamBooth过程。(4)我们在两个面部基准和不同配置下广泛评估了我们提出的方法。我们最好的防御在有利和不利的情况下都有效

2.论文基础

DreamBooth,advDM

3.Anti-DreamBooth实现思想

优化目标

其中Lcond是latent diffusion model的条件去噪损失，Ldb是DreamBooth技术的训练损失

3.1 Fully-trained Surrogate Model Guidance (FSMG）

首先用干净样本(clean set)和一个训练好的latent diffusion model(Pretrained model)进行DreamBooth微调，训练好得到一个代理模型Surrogate DreamBooth model,然后重新将干净样本在代理模型上执行PGD攻击，使得下面的Lcond损失最大化(也就是执行advDM攻击)，生成的图片即对抗样本，这些对抗样本可以防止被滥用于进行个性化文生图

3.2 Alternating Surrogate and Perturbation Learning (ASPL)

使用在干净数据上完全训练的代理模型可能不是解决上面优化目标的最佳近似，受[26]的启发，我们提出以交替的方式将代理DreamBooth模型的训练与扰动学习结合起来。

其实就多了一步将中间对抗样本放到模型中进行DreamBooth训练，可以理解为进行对抗训练提高模型鲁棒性，让后面针对模型添加的噪声添更加多样化，从而产生强大的攻击能力

4.实验

4.1 实验准备

4.1.1 数据集

选择标准：作者选择了两个面部数据集，这些数据集需要满足以下标准：包含大量不同主体的标注图像，每个主体有足够的图像形成参考集和保护集，图像分辨率中等至高，以及图像多样化且在野外环境中拍摄。
CelebA-HQ：这是一个高质量的CelebA数据集版本，包含30,000张1024×1024分辨率的图像。作者使用了注释子集，该子集过滤并分组了307个主体，每个主体至少有15张图像。
VGGFace2：这个数据集包含了大约3.31百万张9131个人身份的图像。作者过滤了数据集，选择了分辨率至少为500×500的图像

[2] CelebA-HQ-Face-Identity-and-Attributes-Recognition-PyTorch. https://github.com/ndb796/
CelebA-HQ-Face-Identity-and-Attributes-Recognition-PyTorch.

[29]Progressive growing of gans for improved quality, stability,and variation. arXiv preprint arXiv:1710.10196, 2017

Qiong Cao, Li Shen, Weidi Xie, Omkar M. Parkhi, and An-drew Zisserman. VGGFace2: A dataset for recognising faces across pose and age. In International Conference on Auto-matic Face and Gesture Recognition, 2018

4.1.2 训练配置

模型训练：作者训练了每个DreamBooth模型，包括文本编码器和UNet模型，批量大小为2，学习率为5×10^-7，训练步骤为1000步。
预训练生成器：默认情况下，使用最新的Stable Diffusion (v2.1) 作为预训练生成器。
训练提示：训练实例提示和先前提示分别是“a photo of sks person”和“a photo of person”

4.1.3 评价指标

Face Detection Failure Rate (FDFR)：使用RetinaFace检测器测量生成图像中无法检测到面部的比率。
Identity Score Matching (ISM)：是一种用于评估图像识别系统性能的指标，特别是在人脸识别领域。它衡量的是两个图像之间的相似度，通常用于比较同一人脸图像的（embeddings）与不同人脸图像嵌入之间的差异。在人脸识别系统中，ISM 可以帮助确定两个图像是否属于同一个人
图像质量评估：使用SER-FQA和BRISQUE两个图像质量评估指标。

4.2 提出的四种攻击方法的对比实验

4.2.1 数据展示

这张实验图（Table 1）展示了在不同数据集上，使用不同防御方法（No Defense, FSMG, ASPL, T-FSMG, T-ASPL）时，DreamBooth模型生成图像的防御性能比较。这些方法旨在防止DreamBooth模型生成高质量的个性化图像。表格中列出了几个关键的评估指标：Face Detection Failure Rate (FDFR)、Identity Score Matching (ISM)、SER-FQA和BRISQUE。这些指标用于衡量生成图像的质量，以及它们在面部检测、身份匹配和图像质量方面的表现。以下是对这些指标的分析：

Face Detection Failure Rate (FDFR)：这个指标衡量的是生成图像中无法被面部检测器正确识别的面部的比例。FDFR越高，表示面部检测失败的比例越大，这通常意味着生成的图像质量较差，面部特征被破坏，从而降低了图像的可用性。
Identity Score Matching (ISM)：ISM衡量的是生成图像的面部嵌入与原始图像集平均面部嵌入之间的相似度。ISM值越低，表示生成图像与原始图像的身份匹配度越低，这表明防御方法成功地扰乱了DreamBooth模型，使其无法生成与目标用户身份匹配的图像。
SER-FQA：这是一个专门用于评估面部图像质量的指标。它考虑了图像的锐度、噪声、压缩伪影等因素。SER-FQA值越低，表示图像质量越好。
BRISQUE：这是一个通用的图像质量评估指标，它考虑了图像的自然度和视觉质量。BRISQUE值越低，表示图像质量越高。

从表格中可以看出，与没有防御（No Defense）相比，所有防御方法（FSMG, ASPL, T-FSMG, T-ASPL）都显著提高了FDFR，降低了ISM，这表明它们有效地破坏了DreamBooth模型生成的图像质量，使得生成的图像在面部检测和身份匹配方面表现不佳。同时，这些防御方法也导致了SER-FQA和BRISQUE值的降低，这表明生成图像的整体质量受到了影响。

在不同的数据集（VGGFace2和CelebA-HQ）上，ASPL方法在大多数情况下都提供了最好的防御效果，尤其是在降低ISM方面。这表明ASPL方法在防止DreamBooth模型生成高质量个性化图像方面最为有效。而T-FSMG和T-ASPL作为针对性的防御方法，它们在某些情况下也提供了额外的保护，尽管它们的表现可能不如ASPL稳定

4.2.2 可视化展示

同的推理提示（inference prompts），这些提示用于指导DreamBooth模型生成特定风格的图像。例如，“a close-up photo of sks person, high details”（近距离高细节的人像照片）和“a photo of sks person in front of Eiffel tower”（在埃菲尔铁塔前的sks人物照片）。
清晰度：没有防御方法时，图像清晰度较高，面部特征明显。随着防御方法的加入，图像的清晰度逐渐降低，面部特征变得模糊不清。
视觉伪影：随着防御方法的加强，图像中出现了更多的视觉伪影，如模糊、扭曲和不自然的纹理。

这张图展示了Anti-DreamBooth系统在不同防御策略下对图像质量的影响。ASPL方法在降低图像质量方面表现最为显著，这表明它在防止DreamBooth模型生成高质量个性化图像方面最为有效。这些定性结果与之前提到的定量评估结果相一致，进一步证实了Anti-DreamBooth系统在保护用户隐私方面的潜力。

4.3 消融实验

文的5.3节 "Ablation studies"（消融研究）中，作者进行了一系列的实验来分析和理解Anti-DreamBooth系统中各个组件的作用和重要性。消融研究通常涉及移除或改变系统的一部分，以观察对整体性能的影响。这有助于确定哪些部分对系统的贡献最大，以及在实际应用中可能需要优化或改进的地方

4.3.1 ASPL在不同预训练生成器版本的消融实验

目的：测试不同版本的预训练文本到图像生成器（如Stable Diffusion的不同版本）对ASPL（Alternating Surrogate and Perturbation Learning）防御方法有效性的影响。
实验：比较了在VGGFace2数据集上，使用ASPL防御方法时，不同版本的Stable Diffusion（v1.4和v1.5）的防御性能。
结果：ASPL在所有版本上都显示出一致的防御效果，表明ASPL方法对于不同版本的生成器具有较好的鲁棒性。

4.3.2 ASPL在不同提示词下的消融实验

目的：评估ASPL防御方法对不同推理文本提示的鲁棒性。
实验：在VGGFace2数据集上，使用ASPL防御方法，测试了不同推理文本提示下的防御性能。
结果：ASPL在不同推理文本提示下都能有效地破坏DreamBooth模型生成的图像，显示出对不同提示的鲁棒性。

4.3.3 ASLP在不同噪声预算下的消融实验

目的：研究噪声预算（η）对ASPL攻击效果的影响。
实验：使用Stable Diffusion v2.1，在VGGFace2数据集上测试了不同噪声预算下的ASPL防御性能。
结果：随着噪声预算的增加，防御性能（如FDFR和ISM）得到改善，但同时图像质量（如PSNR和LPIPS）会有所下降。这表明在噪声预算和图像质量之间存在权衡

4.3.4 其他AIGC任务的效果

目的：测试ASPL防御方法对抗其他个性化文本到图像技术（文本反转 (Textual Inversion) 和 LoRA (Low-Rank Adaptation)）的效果。
实验：在VGGFace2数据集上，使用ASPL防御方法，对抗Textual Inversion和LoRA技术。
结果：ASPL成功地防御了这两种方法，进一步证明了ASPL在对抗不同个性化技术方面的有效性。

4.3.5 消融结论

消融研究的结果表明，ASPL防御方法在不同条件下都能有效工作，无论是在不同的生成器版本、噪声预算、推理文本提示，还是对抗其他个性化技术时。这些发现有助于理解ASPL方法的鲁棒性和潜在的改进方向。

4.4 ASLP的鲁棒性实验

在论文的5.4节 "Adverse settings"（不利设置）中，作者探讨了在不利条件下Anti-DreamBooth防御系统的有效性。这些不利条件包括攻击和防御端的模型、术语或提示不匹配的情况，这些情况可能会影响防御策略的性能，并且作者还有使用预处理防御方法测试攻击的鲁棒性，主要的方法包括Gaussian Blur 和JPEG Comp，没有用最强的DiffPure我猜是攻击失效了😀。

4.4.1模型不匹配 (Model Mismatch)

目的：测试当预训练的文本到图像生成器与用于训练目标DreamBooth模型的生成器不同时，ASPL防御方法的有效性。
实验：在VGGFace2数据集上，作者将ASPL防御方法应用于不同版本的Stable Diffusion模型（v1.4和v1.5），并观察防御效果。
结果：尽管存在模型不匹配，ASPL仍然能够提供良好的防御效果，这表明ASPL具有一定的鲁棒性。

4.4.2术语/提示不匹配 (Term/Prompt Mismatch)

目的：评估当训练DreamBooth模型时使用的术语或提示与用于生成防御噪声的术语或提示不同时，ASPL防御方法的表现。
实验：作者改变了训练术语（例如，从“sks”变为“t@t”）和提示（例如，使用不同的描述性文本例如，如果防御方使用的提示是"a photo of sks person"，而攻击者在训练时使用的提示是"a dslr portrait of sks person"，这就构成了提示不匹配），并观察这些变化对防御效果的影响。
结果：即使在术语或提示不匹配的情况下，ASPL仍然能够保持较低的ISM分数，这意味着它能够有效地防止DreamBooth模型生成具有正确身份的图像。

4.4.3结论

在不利设置下，ASPL防御方法仍然显示出了良好的性能。这表明ASPL能够在一定程度上抵抗模型、术语或提示不匹配带来的影响，这对于实际应用中可能遇到的多样化和不可预测的攻击场景具有重要意义。这些实验结果强调了ASPL作为一种防御策略的实用性和适应性

5. 核心代码

5.1 执行参数

accelerate launch attacks/aspl.py \
  --pretrained_model_name_or_path=$MODEL_PATH  \
  --enable_xformers_memory_efficient_attention \
  --instance_data_dir_for_train=$CLEAN_TRAIN_DIR \
  --instance_data_dir_for_adversarial=$CLEAN_ADV_DIR \
  --instance_prompt="a photo of sks person" \
  --class_data_dir=$CLASS_DIR \
  --num_class_images=200 \
  --class_prompt="a photo of person" \
  --output_dir=$OUTPUT_DIR \
  --center_crop \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --resolution=512 \
  --train_text_encoder \
  --train_batch_size=1 \
  --max_train_steps=50 \
  --max_f_train_steps=3 \
  --max_adv_train_steps=6 \
  --checkpointing_iterations=10 \
  --learning_rate=5e-7 \
  --pgd_alpha=5e-3 \
  --pgd_eps=5e-2 


# ------------------------- Train DreamBooth on perturbed examples -------------------------
export INSTANCE_DIR="$OUTPUT_DIR/noise-ckpt/50"
export DREAMBOOTH_OUTPUT_DIR="outputs/$EXPERIMENT_NAME/n000050_DREAMBOOTH"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_PATH  \
  --enable_xformers_memory_efficient_attention \
  --train_text_encoder \
  --instance_data_dir=$INSTANCE_DIR \
  --class_data_dir=$CLASS_DIR \
  --output_dir=$DREAMBOOTH_OUTPUT_DIR \
  --with_prior_preservation \
  --prior_loss_weight=1.0 \
  --instance_prompt="a photo of sks person" \
  --class_prompt="a photo of person" \
  --inference_prompt="a photo of sks person;a dslr portrait of sks person" \
  --resolution=512 \
  --train_batch_size=2 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-7 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=200 \
  --max_train_steps=1000 \
  --checkpointing_steps=500 \
  --center_crop \
  --mixed_precision=bf16 \
  --prior_generation_precision=bf16 \
  --sample_batch_size=8

5.2 ASPL攻击过程

#unet text_encoder:原本latent diffusion model v2.0的两个模型
    f = [unet, text_encoder]
    for i in range(args.max_train_steps):
        """1. aspl第一步θ'=θ.clone()"""
        f_sur = copy.deepcopy(f)
        """2. aspl第二步 argmin ∑Ldb(θ‘，x)"""
        f_sur = train_one_epoch(
            args,
            f_sur,
            tokenizer,
            noise_scheduler,
            vae,
            clean_data,
            args.max_f_train_steps,
        )
        """3. argmax Lcond(θ’，x + σ) ，执行的内容是advDM攻击"""
        perturbed_data = pgd_attack(
            args,
            f_sur,
            tokenizer,
            noise_scheduler,
            vae,
            perturbed_data,
            original_data,
            target_latent_tensor,
            args.max_adv_train_steps,
        )
        """4. argmin ∑ Ldb(θ，x+σ)"""
        f = train_one_epoch(
            args,
            f,
            tokenizer,
            noise_scheduler,
            vae,
            perturbed_data,
            args.max_f_train_steps,
        )

5.3 DreamBooth训练

def train_one_epoch(
    args,
    models,
    tokenizer,
    noise_scheduler,
    vae,
    data_tensor: torch.Tensor, #clean_data
    num_steps=20,   #args.max_f_train_steps
):
    """"""
    """1. aspl第一步θ'=θ.clone()"""
    #深拷贝噪声预测器和文本token到unet和text_tokenizer
    unet, text_encoder = copy.deepcopy(models[0]), copy.deepcopy(models[1])
    #将unet和text_tokenizer的参数到params_to_optimize
    params_to_optimize = itertools.chain(unet.parameters(), text_encoder.parameters())
    #使用AdamW优化器
    optimizer = torch.optim.AdamW(
        params_to_optimize,   #优化器优化的参数
        lr=args.learning_rate,#学习率
        betas=(0.9, 0.999),  #beta参数
        weight_decay=1e-2,  #权重衰减
        eps=1e-08,          #epsilon
    )

    train_dataset = DreamBoothDatasetFromTensor(
        #clean_data :即instance_data，要保护的干净样本
        data_tensor,
        #instance_prompt:a photo of sks person,
        #让diffusion学习干净样本信息的引导词
        args.instance_prompt,
        #CLIPTokenizer
        tokenizer,
        #保留diffusion先验信息图片的路径
        args.class_data_dir,
        #class_prompt: a photo of person
        #让diffusion保留先验信息的引导词
        args.class_prompt,
        args.resolution,
        args.center_crop,
    )

    weight_dtype = torch.bfloat16
    device = torch.device("cuda")

    vae.to(device, dtype=weight_dtype)
    text_encoder.to(device, dtype=weight_dtype)
    unet.to(device, dtype=weight_dtype)
    """2.aspl第二步 argmin ∑Ldb(θ‘，x)"""
    for step in range(num_steps):
        #θ'进入训练模式
        unet.train()
        text_encoder.train()
        #获取这一步step的train_dataset
        step_data = train_dataset[step % len(train_dataset)]
        #将保护的干净样本和保留先验知识的样本堆叠在一起
        """
        举个例子，假设我们有两个张量 tensor1 和 tensor2，它们的形状分别为 (3, 2) 和 (3, 2)。
        通过 torch.stack([tensor1, tensor2])，我们将得到一个新的张量，形状为 (2, 3, 2)，
        表示两个输入张量在新的维度上被堆叠在一起
        """
        pixel_values = torch.stack([step_data["instance_images"], step_data["class_images"]]).to(
            device, dtype=weight_dtype
        )
        input_ids = torch.cat([step_data["instance_prompt_ids"], step_data["class_prompt_ids"]], dim=0).to(device)
        #通过VAE生成保护的干净样本和保留先验知识的样本的latent representation 潜在特征
        latents = vae.encode(pixel_values).latent_dist.sample()
        #将潜在特征进行缩放
        latents = latents * vae.config.scaling_factor

        # Sample noise that we'll add to the latents
        noise = torch.randn_like(latents)
        bsz = latents.shape[0]
        # Sample a random timestep for each image
        # (每一张图片都是随机挑选时间步添加噪声，后面有论文基于时间步进行改进)
        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
        timesteps = timesteps.long()

        # Add noise to the latents according to the noise magnitude at each timestep
        # (this is the forward diffusion process)，前向加噪过程
        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

        # Get the text embedding for conditioning
        encoder_hidden_states = text_encoder(input_ids)[0]

        # Predict the noise residual
        # 预测添加到保护的干净样本和保留先验知识的样本的噪声
        model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample

        # Get the target for loss depending on the prediction type
        #获得实际添加到保护的干净样本和保留先验知识的样本的噪声
        if noise_scheduler.config.prediction_type == "epsilon":
            target = noise
        elif noise_scheduler.config.prediction_type == "v_prediction":
            target = noise_scheduler.get_velocity(latents, noise, timesteps)
        else:
            raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")

        # with prior preservation loss
        if args.with_prior_preservation:
            #把保护的干净样本和保留先验知识的样本的预测噪声和实际噪声拆开
            model_pred, model_pred_prior = torch.chunk(model_pred, 2, dim=0)
            target, target_prior = torch.chunk(target, 2, dim=0)

            # Compute instance loss
            instance_loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")

            # Compute prior loss
            prior_loss = F.mse_loss(model_pred_prior.float(), target_prior.float(), reduction="mean")

            # Add the prior loss to the instance loss.
            loss = instance_loss + args.prior_loss_weight * prior_loss

        else:
            loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")

        loss.backward()
        torch.nn.utils.clip_grad_norm_(params_to_optimize, 1.0, error_if_nonfinite=True)
        optimizer.step()
        optimizer.zero_grad()
        print(
            f"Step #{step}, loss: {loss.detach().item()}, prior_loss: {prior_loss.detach().item()}, instance_loss: {instance_loss.detach().item()}"
        )

    return [unet, text_encoder]

5.4 advDM(PGD)攻击

def pgd_attack(
    args,
    models,
    tokenizer,
    noise_scheduler,
    vae,
    data_tensor: torch.Tensor,
    original_images: torch.Tensor,
    target_tensor: torch.Tensor,
    num_steps: int,
):
    """Return new perturbed data"""

    unet, text_encoder = models
    weight_dtype = torch.bfloat16
    device = torch.device("cuda")

    vae.to(device, dtype=weight_dtype)
    text_encoder.to(device, dtype=weight_dtype)
    unet.to(device, dtype=weight_dtype)

    perturbed_images = data_tensor.detach().clone()
    perturbed_images.requires_grad_(True)

    input_ids = tokenizer(
        args.instance_prompt,
        truncation=True,
        padding="max_length",
        max_length=tokenizer.model_max_length,
        return_tensors="pt",
    ).input_ids.repeat(len(data_tensor), 1)

    for step in range(num_steps):
        perturbed_images.requires_grad = True
        latents = vae.encode(perturbed_images.to(device, dtype=weight_dtype)).latent_dist.sample()
        latents = latents * vae.config.scaling_factor

        # Sample noise that we'll add to the latents
        noise = torch.randn_like(latents)
        bsz = latents.shape[0]
        # Sample a random timestep for each image
        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
        timesteps = timesteps.long()
        # Add noise to the latents according to the noise magnitude at each timestep
        # (this is the forward diffusion process)
        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)

        # Get the text embedding for conditioning
        encoder_hidden_states = text_encoder(input_ids.to(device))[0]

        # Predict the noise residual
        model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample

        # Get the target for loss depending on the prediction type
        if noise_scheduler.config.prediction_type == "epsilon":
            target = noise
        elif noise_scheduler.config.prediction_type == "v_prediction":
            target = noise_scheduler.get_velocity(latents, noise, timesteps)
        else:
            raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")

        unet.zero_grad()
        text_encoder.zero_grad()
        loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")

        # target-shift loss
        if target_tensor is not None:
            xtm1_pred = torch.cat(
                [
                    noise_scheduler.step(
                        model_pred[idx : idx + 1],
                        timesteps[idx : idx + 1],
                        noisy_latents[idx : idx + 1],
                    ).prev_sample
                    for idx in range(len(model_pred))
                ]
            )
            xtm1_target = noise_scheduler.add_noise(target_tensor, noise, timesteps - 1)
            loss = loss - F.mse_loss(xtm1_pred, xtm1_target)

        loss.backward()

        alpha = args.pgd_alpha
        eps = args.pgd_eps

        adv_images = perturbed_images + alpha * perturbed_images.grad.sign()
        eta = torch.clamp(adv_images - original_images, min=-eps, max=+eps)
        perturbed_images = torch.clamp(original_images + eta, min=-1, max=+1).detach_()
        print(f"PGD loss - step {step}, loss: {loss.detach().item()}")
    return perturbed_images