Defensive Unlearning with Adversarial Training for Robust Concept Erasure in Diffusion Models

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

한 줄 요약

AdvUnlearn은 적대적 훈련을 통합한 확고한 개념 제거 프레임워크로, 디퓨전 모델의 안전성과 생성 품질을 균형 있게 유지한다.

핵심 기여도

**AdvUnlearn**이라는 새로운 개념 제거 프레임워크를 제안, 적대적 훈련과 개념 제거를 통합.
**이용성 유지 정규화**(utility-retaining regularization)를 도입, 생성 품질 저하를 최소화.
**텍스트 인코더**(text encoder)를 강화 대상으로 선정, UNet 기반 기존 방법보다 뛰어난 성능.
**플러그 앤 플레이 방식**의 텍스트 인코더를 통해 다양한 디퓨전 모델에 적용 가능.

핵심 아이디어

기존의 개념 제거 기법은 **적대적 프롬프트 공격**(adversarial prompt attack)에 취약하다는 문제를 해결하기 위해, **적대적 훈련**(adversarial training, AT)을 기계 언러닝에 통합한 **AdvUnlearn** 프레임워크를 제안한다. 이는 **이중 최적화**(bi-level optimization, BLO) 기반의 방식으로, 개념 제거 목적과 적대적 공격 최적화를 동시에 고려한다.

기존 연구는 주로 **UNet** 모듈을 강화 대상으로 삼았지만, 본 연구는 **텍스트 인코더**(text encoder)가 개념 제거의 효과를 더 잘 반영하고, **플러그 앤 플레이 방식**으로 다양한 디퓨전 모델에 재사용 가능하다는 점을 발견했다. 또한, 생성 품질 저하를 방지하기 위해 **추가 보존 프롬프트 세트**(retain prompt set)를 기반으로 **이용성 유지 정규화**를 설계하여, 개념 제거의 효과와 모델 성능 간의 균형을 맞추는 것이 핵심이다.

기술적 접근법

**AdvUnlearn**은 **이중 최적화**(BLO) 기반의 프레임워크로, 상위 최적화는 개념 제거 목적(예: ESD의 목적 함수)을, 하위 최적화는 적대적 프롬프트 생성 목적(공식 (4))을 처리.
**이용성 유지 정규화**는 **추가 보존 프롬프트 세트**(retain set)를 사용해 개념 제거 후의 이미지 생성 품질을 유지.
**적대적 프롬프트 공격**은 토큰 조작 또는 임베딩 공간에서 연속적 편차를 통해 생성, 토큰 변화량은 ℓ₀ 노름으로 측정.
**텍스트 인코더**가 **UNet**보다 개념 제거에 더 효과적임을 실험적으로 입증, 이 모듈을 강화 대상으로 선정.

주요 결과

**AdvUnlearn**은 **I2P 데이터셋**에서 **적대적 프롬프트 공격**에 대해 기존 방법 대비 **높은 개념 제거 강도**를 보임.
**텍스트 인코더 기반**의 AdvUnlearn은 **UNet 기반** 방법보다 **더 높은 생성 품질**과 **더 낮은 개념 재생성률**을 달성.
**플러그 앤 플레이 방식**의 텍스트 인코더는 다양한 디퓨전 모델에 적용 가능, **재사용성**이 뛰어남.
**Nudity, Objects, Style** 등 다양한 개념 제거 시나리오에서 **적대적 공격에 대한 강건성**이 입증됨.

의의 및 한계

AdvUnlearn은 **적대적 훈련**(AT)을 기계 언러닝에 체계적으로 통합한 **최초의 연구**로, 기존 개념 제거 기법의 **적대적 공격 취약성**을 해결하는 데 기여한다. 특히, **텍스트 인코더**를 강화 대상으로 선정함으로써 **UNet 기반 방법**보다 **더 높은 강건성과 생성 품질**을 달성한 점이 학술적·실용적 가치가 있다.

하지만, **적대적 훈련의 계산 비용**이 높아 대규모 모델에 적용 시 **효율성** 문제가 있을 수 있으며, **보존 프롬프트 세트**의 품질이 **이용성 유지 정규화**의 효과에 직접적으로 영향을 미친다는 점도 한계로 지적된다.