[논문 리뷰] Defensive Unlearning with Adversarial Training for Robust Concept Erasure in Diffusion Models

Yimeng Zhang, Xin Chen외 7명|arXiv (Cornell University)|2024. 05. 24.

Adversarial Robustness in Machine Learning인용 수 4

한 줄 요약

이 논문은 적대적 프롬프트 공격으로부터 보호하기 위해 확산 모델(DM) 개념 삭제에 적대적 훈련을 통합한 새로운 강건한 언러닝 프레임워크인 AdvUnlearn을 제안한다. 유지된 프롬프트 세트에 대해 유용성 유지를 보존하는 정규화와 함께 텍스트 인코더에 적대적 훈련을 적용함으로써, AdvUnlearn은 강력한 강건성(예: SD v1.4에서 공격 성공률 21.13%)을 달성하면서도 높은 이미지 생성 품질(FID ≈ 19.34)을 유지하여 기존 베이스라인보다 강건성과 유용성 측면에서 모두 뛰어나다.

ABSTRACT

Diffusion models (DMs) have achieved remarkable success in text-to-image generation, but they also pose safety risks, such as the potential generation of harmful content and copyright violations. The techniques of machine unlearning, also known as concept erasing, have been developed to address these risks. However, these techniques remain vulnerable to adversarial prompt attacks, which can prompt DMs post-unlearning to regenerate undesired images containing concepts (such as nudity) meant to be erased. This work aims to enhance the robustness of concept erasing by integrating the principle of adversarial training (AT) into machine unlearning, resulting in the robust unlearning framework referred to as AdvUnlearn. However, achieving this effectively and efficiently is highly nontrivial. First, we find that a straightforward implementation of AT compromises DMs' image generation quality post-unlearning. To address this, we develop a utility-retaining regularization on an additional retain set, optimizing the trade-off between concept erasure robustness and model utility in AdvUnlearn. Moreover, we identify the text encoder as a more suitable module for robustification compared to UNet, ensuring unlearning effectiveness. And the acquired text encoder can serve as a plug-and-play robust unlearner for various DM types. Empirically, we perform extensive experiments to demonstrate the robustness advantage of AdvUnlearn across various DM unlearning scenarios, including the erasure of nudity, objects, and style concepts. In addition to robustness, AdvUnlearn also achieves a balanced tradeoff with model utility. To our knowledge, this is the first work to systematically explore robust DM unlearning through AT, setting it apart from existing methods that overlook robustness in concept erasing. Codes are available at: https://github.com/OPTML-Group/AdvUnlearn

연구 동기 및 목표

• 개념 삭제된 확산 모델이 언러닝을 회피하는 적대적 프롬프트 공격에 취약한 문제를 해결하기 위해.
• 적대적 프롬프트 탈옥 공격에 저항하면서도 높은 이미지 생성 품질을 유지하는 강건한 언러닝 프레임워크를 개발하기 위해.
• 언러닝에서 적대적 강건성 향상에 가장 효과적인 모델 구성 요소(텍스트 인코더 대비 UNet)를 규명하기 위해.
• 다양한 확산 모델 아키텍처 간에 강건한 언러너의 플러그 앤 플레이 전이 가능성을 보장하기 위해.
• 확산 모델의 머신 언러닝 맥락에서 적대적 훈련의 효과성과 효율성을 체계적으로 연구하기 위해.

제안 방법

• 적대적 훈련(AT)을 머신 언러닝 파이프라인에 통합하여 적대적 프롬프트 공격에 대한 강건성을 향상시키기 위해.
• 적대적 훈련을 UNet이 아닌 텍스트 인코더에 특별히 적용함으로써 강건한 언러닝에 더 효과적임을 입증함.
• 적대적 미세조정 중 이미지 생성 품질을 유지하기 위해 유지된 프롬프트 세트에 대해 유용성 유지를 보존하는 정규화를 적용함.
• 언러닝된 모델에 대한 공격 성공률를 최대화하기 위해 다단계 적대적 프롬프트 생성 전략을 사용함.
• 재훈련 없이도 다른 확산 모델에 훈련된 텍스트 인코더를 플러그 앤 플레이 방식으로 배치 가능하게 함.
• 텍스트 인코더의 첫 N개 레이어를 최적화하며, 분석 결과 깊이 있는 레이어 최적화가 낙관적 개념과 같은 복잡한 개념에 대해 강건성을 향상시킴.

Figure 1: Visualizations of SD v1.4 (Rombach et al., 2022 ) and its concept-erased version, ESD (Gandikota et al., 2023a ) , conditioned on the prompt ‘Hot sexy girl’ . The generation of ESD is shown against no attack and UnlearnDiffAtk (Zhang et al., 2023b ) .

실험 결과

연구 질문

• RQ1 적대적 훈련이 확산 모델 언러닝에 효과적이고 효율적으로 통합되어 적대적 프롬프트에 대한 강건성을 향상시킬 수 있는가?
• RQ2 확산 모델의 구성 요소 중에서 언러닝에서 적대적 강건성 향상에 더 효과적인 것은 텍스트 인코더인가, UNet인가?
• RQ3 적대적 언러닝 과정에서 강건성과 이미지 생성 유용성 간의 트레이드오프를 효과적으로 관리할 수 있는가?
• RQ4 한 확산 모델에서 훈련된 강건한 언러너를 재훈련 없이 다른 모델에 전이할 수 있는가?
• RQ5 다양한 텍스트 인코더 레이어 구성이 언러닝의 강건성과 유용성에 어떤 영향을 미치는가?

주요 결과

• AdvUnlearn은 SD v1.4에서 낙관적 개념 삭제에 대해 공격 성공률(ASR)을 원본 모델의 100%에서 21.13%로 감소시켰다.
• 이 프레임워크는 강력한 이미지 생성 품질을 유지하며, FID가 19.34로 원본 모델의 16.70과 거의 유사한 수준을 기록했다.
• 텍스트 인코더는 UNet보다 적대적 강건성 향상에 훨씬 더 효과적이며, 더 낮은 ASR과 더 나은 유용성 트레이드오프를 달성했다.
• AdvUnlearn를 통해 학습된 강건한 텍스트 인코더는 재훈련 없이도 다른 모델들(SD v1.5, DreamShaper, Protogen 등)에 전이 가능했으며, 성능 저하가 최소화되어 SD v1.5에선 16.20%의 ASR, Protogen에선 42.96%의 ASR을 기록했다.
• 텍스트 인코더의 더 많은 레이어를 최적화할수록 강건성이 향상되었으며, 전체 레이어 최적화가 복잡한 개념(예: 낙관적 개념)에 대해 가장 강력한 방어를 제공했다.
• 유용성 유지 정규화는 일반적으로 적대적 훈련으로 인해 발생하는 이미지 품질 저하를 효과적으로 완화하여 모델의 유용성을 유지했다.

Figure 2: Performance overview of our proposal AdvUnlearn and various DM unlearning baselines when unlearning the nudity concept under the SD v1.4 model. The robustness is measured by attack success rate (ASR) against UnlearnDiffAtk (Zhang et al., 2023b ) . The performance of image generation retent