TEA: Text Encoder Alignment for Robust Concept Erasure in Text-to-Image Models

기술 요약: TEA (강건한 개념 삭제를 위한 텍스트 인코더 정렬)

문제 정의

텍스트-이미지 확산 모델(Diffusion Models)은 적대적(adversarial) 또는 의역된(paraphrased) 프롬프트가 내장된 안전 메커니즘을 우회하여 유해한 콘텐츠(예: 딥페이크, NSFW 이미지)를 생성할 수 있는 오용 가능성에 취약합니다. 기존의 개념 삭제 방법들은 다음과 같은 상당한 트레이드오프에 직면해 있습니다:

• 훈련이 필요 없는 방법(Training-free methods) (예: 추론 시 투영법)은 사전 비용을 피할 수 있지만, 지속적인 쿼리당 계산 오버헤드를 발생시키며 강력한 적대적 공격에 대한 강건성이 부족합니다.
• **생성 백본(예: U-Net)을 미세 조정하는 훈련 기반 방법(Training-based methods)**은 계산 비용이 많이 들고, 무해한 프롬프트에 대한 생성 품질을 저하시킬 위험이 있으며, 네트워크 레이어 전반에 걸쳐 넓게 분포된 개념 관련 파라미터로 인해 선택적 삭제가 어렵습니다.
• 기존의 텍스트 인코더 접근 방식은 대규모 데이터셋, 코드북 메모리 또는 생성 품질을 저하시키는 계산 집약적인 적대적 훈련을 요구하는 경우가 많습니다.

핵심 과제는 무해한 프로ンプ트에 대한 생성 품질을 보존하면서, 블랙박스 및 화이트박스 공격 모두에 대해 강건한 개념 삭제를 달고, **추론 시 오버헤드가 제로(zero)**인 상태를 달성하는 것입니다.

방법론: 텍스트 인코더 정렬 (TEA)

저자들은 개념 삭제를 텍스트 표현 공간 내에서의 **도메인 정렬 문제(domain alignment problem)**로 정식화합니다. 이 방법은 텍스트-이미지 모델을 텍스트 인코더($\tau$)와 동결된 생성 백본($G$)으로 분해하는 것에 기초합니다.

1. 이론적 정식화
본 논문은 타겟 개념 $c$를 억제하는 것이 개념을 포함하는 입력의 표현 분포($P_c$)를 개념이 없는 "앵커(anchor)" 입력의 분포($P_{\neg c}$)와 정렬하는 것과 동일하다고 가정합니다.

• 리스크 바운드(Risk Bound): H-divergence 이론에 근거하여, 개념을 포함하는 입력에 대한 개념 리스크는 개념이 없는 입력에 대한 개념 리스크와 두 표현 분포 사이의 발산(divergence)의 합으로 제한됩니다.
• 목표: 발산(두 표현을 구별할 수 없게 만듦)을 최소화하는 동시에, 개념이 없는 입력에 대한 리스크(유용성 보존)를 유지함으로써, 모델이 안전한 프롬프트와 안전하지 않은 프롬프트를 더 이상 구별할 수 없게 하여 개념을 억제합니다.

2. 프레임워크 아키텍처
TEA는 생성 백본(U-Net 또는 DiT)을 완전히 동결한 채 텍스트 인코더만을 미세 조정합니다. 과정은 다음과 같습니다:

• 프롬프트 쌍 구축: 각 개념 프롬프트 $p_c$(타겟 개념 포함)에 대해, 주변 의미는 유지하면서 타겟 개념을 제거한 앵커 프롬프트 $p_a$를 생성합니다 (예: LLM을 통해 생성).
• 토큰 수준 정렬(Token-Level Alignment): 풀링된 시퀀스 표현을 사용하는 방법과 달리, TEA는 토큰 수준에서 작동합니다. 이를 통해 모델은 관련 없는 의미를 붕괴시키지 않으면서 $p_c$와 $p_a$ 사이의 특정 의미적 차이를 정렬할 수 있습니다.
• 적대적 훈련(Adversarial Training):
  • 경량화된 **판별기($D$)**를 훈련시켜 개념 프롬프트의 토큰 수준 표현($h^c_t$)과 앵커 프롬프트의 표현($h^a_t$)을 구분하도록 합니다.
  • **텍스트 인코더($\tau_\theta$)**는 판별기를 속이도록 업데이트되어, $h^c_t$를 $h^a_t$와 구별할 수 없게 만듭니다.
  • 동적 앵커링(Dynamic Anchoring): 개념 및 앵커 프롬프트 모두 훈련 중에 현재의 인코더를 통과합니다. 이는 인코더가 고정된 참조점에 붕괴되는 것을 방지하고, 일관되게 진화하는 안전한 분포를 추적하도록 장려합니다.
• 손실 함수(Loss Functions):
  • 적대적 손실($L_{adv}$): 판별기가 개념 토큰을 정확하게 분류하는 능력을 최소화합니다.
  • 보존 손실($L_{pres}$): 코사인 유사도 손실로서, 앵커 프롬프트의 표현($h^{a, new}_t$)을 원래의 동결된 인코더가 생성한 표현($h^{a, orig}_t$)과 가깝게 유지합니다. 이는 무해한 의미가 저하되지 않도록 보장합니다.
  • 전체 목적 함수: $L_{enc} = \lambda_{adv} L_{adv} + \lambda_{pres} L_{pres}$.

주요 기여

1. 정식화: 저자들은 개념 삭제를 표현 공간의 도메인 정렬 문제로 보는 원칙적인 정식화를 제공하며, 삭제가 정렬 항과 유용성 보존 항에 의해 제어된다는 이론적 바운드를 도출했습니다.
2. TEA 프레임워크: 이들은 추론 시 오버헤드가 제로이며 소수의 미세 조정 단계만 필요한 경량화된 적대적 텍스트 인코더 미세 조정 프레임워크를 제안합니다.
3. 모델 불가지론 및 새로운 아키텍처: 이 방법은 모델 불가지론적(model-agnostic)이며, 기존의 개념 삭제 방법들이 거의 탐구되지 않았던 T5 컨디셔닝을 사용하는 현대적 아키텍처인 Rectified Flow Transformer(예: Stable Diffusion v3.5)로 성공적으로 전이됩니다.
4. 강건성: TEA는 기존 베이스라인들과 비교하여 블랙박스(프롬프트 엔리니어링) 및 화이트박스(최적화 기반) 적대적 공격 모두에 대해 우수한 강건성을 입증했습니다.

실험 결과

이 방법은 노골적인 콘텐츠(누드) 및 예술적 스타일(Van Gogh, Kelly McKernan)을 삭제하기 위해 Stable Diffusion v1.4 및 Stable Diffusion v3.5-large에서 평가되었습니다.

• 노골적 콘텐츠 삭제:
  • 블랙박스: TEA는 MMA, Ring-A-Bell, P4D 벤치마크 전반에서 가장 낮은 평균 공격 성공률(ASR)을 달립으로써 AdvUnlearn 및 DES와 같은 방법들을 능가했습니다.
  • 화이트박스: TEA는 UnlearnDiff, MMA, Ring-A-Bell 화이트박스 공격에서 일관되게 가장 낮은 ASR을 달성했습니다.
  • 품질 보존: TEA는 무해한 프롬프트(COCO 데이터셋)에 대해 높은 생성 품질을 유지했으며, 베이스라인 대비 경쟁력 있거나 우수한 FID 및 CLIP 점수를 보여주었습니다. 특히, 적대적 훈련 기반 방법에서 나타나는 품질 저하 현상을 피했습니다.
• 예술적 스타일 삭제: TEA는 타겟 작가에 대해서는 높은 지각적 거리(LPIPSd)를 달성하면서도, 다른 작가들에 대해서는 낮은 거리(LPIPSu)를 유지하여, 관련 없는 스타일을 과잉 삭제하지 않고 효과적인 삭제를 수행함을 보여주었습니다.
• 표현 분석: Proxy A-Distance (PAD)의 경험적 측정을 통해, TEA가 인코더 레이어 전반에서 개념과 앵커 표현 간의 분리 가능성을 크게 감소시키는 동시에, 개념과 관련 없는(COCO) 프롬프트 간의 분리성은 유지함을 확인했습니다.

의의 및 주장

본 논문은 TEA가 계산 효율적이고, 지속적이며, 강건한 솔루션을 제공함으로써 현재 개념 삭제 방법들의 결정적인 한계를 해결한다고 주장합니다.

• 효율성: 생성 백본을 동결하고 텍스트 인코더만을 미세 조정함으로써, TEA는 추론 시 오버헤드를 도입하지 않아 대규모 배포에 적합합니다.
• 강건성: 표현 공간에서의 적대적 정렬은 훈련이 필요 없는 방법이나 생성 백본을 수정하는 방법보다 적대적 프롬프트에 대해 더 강력한 방어를 제공합니다.
• 확장성: Rectified Flow Transformer(SD v3.5)에 적응할 수 있는 능력은 개념 삭제의 적용 범위를 차세대 확산 아키텍처로 확장합니다.
• 윤리적 입장: 저자들은 이 연구가 피해 감소를 목적으로 함을 강조합니다. 효능 검증을 위해 노골적인 프롬프트와 이미지가 필요하지만, 이는 엄격히 통제되며, 그림에서 마스킹 처리되었고, 오직 안전하지 않은 생성을 완화하기 위한 방어적 목적으로만 사용되었습니다.