Concept Unlearning via Cross-Attention Activation Projection for Diffusion Models

기술 요약: 확산 모델을 위한 교차 주의 활성화 투영을 통한 개념 망각

문제 제기

개념 망각 (Concept unlearning) 은 사전 학습된 텍스트-이미지 확산 모델에서 특정 대상 개념 (예: 저작권이 있는 스타일, 유명인의 초상, 또는 유해한 이미지) 을 처음부터 재학습하는 데 드는 계산 비용을 들이지 않고 제거하는 것을 목표로 합니다. 효과적인 망각은 두 가지 동시 요구 사항을 충족해야 합니다: 파라프레이즈되거나 적대적인 프롬프트 하에서도 대상 개념을 망각하고, 모델이 모든 다른 관련 없는 개념을 생성할 수 있는 능력을 유지하는 것입니다.

기존의 폐형식 (closed-form) 방법들은 경사 기반 미세 조정이 없이 모델 가중치에 단일 결정론적 편집을 적용하는 데서 중요한 한계에 직면합니다. 이러한 방법들 (예: UCE, CURE) 은 짧은 앵커 프롬프트 (예: "반 고흐") 로부터 유도된 텍스트 인코더 임베딩을 사용하여 "망각 기저 (forget basis)"를 구성합니다. 저자들은 텍스트 임베딩이 사용자의 프롬프트를 설명하지만, 탈노이즈 (denoising) 과정 동안 모델이 실제로 그 프롬프트를 어떻게 활용하는지는 포착하지 못한다고 주장합니다. 결과적으로, 앵커 템플릿과 일치하지 않으면서도 대상 개념을 불러일으키는 파라프레이즈된 프롬프트는 종종 망각 편집을 우회하여 "대상 누출 (target leakage)"을 초래합니다.

방법론: PURE

이 논문은 U-Net 내의 교차 주의 (cross-attention) 활성화 공간으로 기저 구성을 이동시키는 폐형식 망각 방법인 PURE(Projection in U-Net Rendering for Erasure) 를 제안합니다.

핵심 가설:
레이어 $\ell$의 교차 주의 활성화 ($h^\ell$) 는 현재 잠재 상태와 시간 단계 하에서 탈노이즈기가 텍스트 조건을 어떻게 활용하는지를 반영합니다. 정적인 텍스트 임베딩과 달리, 이러한 활성화는 이미지 생성 동안 실제로 사용된 특정 개념 증거를 포착하여 다양한 프롬프트 표현에 대한 더 나은 일반화를 제공합니다.

알고리즘 절차:

1. 활성화 포착: 망각 앵커 집합 ($\mathcal{A}_f$) 과 유지 앵커 집합 ($\mathcal{A}_r$) 에 대해 모델이 짧은 탈노이즈 궤적을 실행합니다. 각 레이어 $\ell$와 시간 단계에서 주의 후 (post-attention) 활성화가 추출되고, 공간적으로 평균 풀링되어 행렬 $H^\ell_F$와 $H^\ell_R$로 집계됩니다.
2. 부분 공간 구성: 각 레이어에 대해 $H^\ell_F$와 $H^\ell_R$에 대해 특이값 분해 (SVD) 를 수행합니다. 누적 분산 임계값 ($\tau_F, \tau_R$) 까지 상위 특이 벡터들이 망각 기저 ($V^\ell_F$) 와 유지 기저 ($V^\ell_R$) 를 형성합니다.
3. 선형 편집: 이 방법은 교차 주의 키 ($W^\ell_K$) 와 값 ($W^\ell_V$) 투영 행렬에 단일 선형 업데이트를 적용합니다. 업데이트 연산자 $E^\ell$은 다음과 같이 정의됩니다:
   $$E^\ell = I - P^\ell_F (I - P^\ell_R)$$
   여기서 $P^\ell_F = V^\ell_F (V^\ell_F)^\top$이고 $P^\ell_R = V^\ell_R (V^\ell_R)^\top$입니다.
   가중치는 왼쪽 곱셈을 통해 업데이트됩니다: $W^\ell \leftarrow E^\ell W^\ell$.
   • 메커니즘: 항 $(I - P^\ell_R)$은 유지 개념과 정렬된 성분을 제거하여 $P^\ell_F$가 직교 부분 공간에서만 작용하도록 보장합니다. 이를 단위 행렬에서 빼면 유지 개념과 정렬된 입력이 변경 없이 통과하도록 보장합니다.
   • 구조적 변화: 텍스트 공간에서 오른쪽에서 곱하는 이전 방법들과 달리, PURE 는 활성화 공간 ($d_\ell \times d_\ell$) 에서 왼쪽에서 곱하므로 레이어별 편집이 필요합니다.

주요 기여

1. 활성화 기반 기저: 이 논문은 텍스트 임베딩이 아닌 교차 주의 활성화로부터 망각 기저를 구성하는 것이 다양한 자연스럽고 적대적인 프롬프트에서 대상 개념을 탐지하는 능력을 현저히 향상시킨다고 규명했습니다.
2. 프로빙 실험: 이진 분류 실험은 텍스트 기반 기저에 비해 활성화 기반 기저가 자연스러운 파라프레이즈된 프롬프트에서 약 **5 배 높은 재현율 (recall)**을 달성함을 보여주었으며, 이는 활성화 특징이 생성과 관련된 개념 증거를 더 잘 포착함을 확인시켜 주었습니다.
3. 폐형식 효율성: PURE 는 폐형식 편집의 장점을 유지합니다: 경사 학습이 필요 없고, 보조 손실 함수가 필요 없으며, 학습률에 대한 하이퍼파라미터 튜닝이 필요 없으며, 편집 적용 후 추론 시간 비용이 추가되지 않습니다.
4. 레이어별 편집: 이 방법은 U-Net 깊이에 따라 개념 표현이 달라진다는 점을 인정하여 레이어별 SVD 와 편집을 수행합니다.

실험 결과

저자들은 **Holistic Unlearning Benchmark (HUB)**에서 PURE 를 평가했는데, 이는 예술적 스타일, 지식재산권 (IP), 유명인, NSFW 등 네 가지 카테고리에 걸친 열 가지 개념을 포함합니다.

• 망각 - 유지 트레이드오프: PURE 는 다섯 가지 베이스라인 (ESD, MACE, RECELER, UCE, CURE) 과 비교하여 모든 카테고리에서 가장 우수한 조화 평균 점수를 달성했습니다.
• 대상 억제: CURE 는 종종 파라프레이즈를 억제하지 못하는 반면, PURE 는 파라프레이즈 및 적대적 프롬프트 하에서 대상 누출을 현저히 감소시킵니다.
• 유지: PURE 는 공격적인 학습 기반 방법들 (유지를 손상시킴) 보다 "유지" 개념 (동일 카테고리 내 인접 개념) 을 더 잘 보존하며, 스타일 카테고리에서 억제를 위해 유지를 희생하는 경우가 많은 텍스트 공간 폐형식 방법들보다도 더 잘 보존합니다.
• 애블레이션 연구:
  • 앵커 크기: 자연스러운 망각 앵커의 수를 늘리는 것은 텍스트 기저의 억제를 향상시키지만 유지를 저하시킵니다. 활성화 기저는 앵커 집합이 커짐에 따라도 최소한의 유지 손실로 강력한 억제를 유지합니다.
  • 탈노이즈 단계: 단일 탈노이즈 단계 ($T=1$) 를 사용하면 성능이 심각하게 저하되는데, 이는 궤적 전반에 걸쳐 개념 신호를 포착하는 것이 필수적임을 나타냅니다.
  • 잠재 평균화: 여러 무작위 잠재 ($n_{lat}$) 에 대해 평균화하는 것은 노이즈 특정 아티팩트보다는 개념 수준의 방향으로 기저를 안정화시킵니다.

의의 및 주장

이 논문은 PURE 가 폐형식 망각 접근 방식에 근본적인 변화를 가져왔다고 주장합니다. 정적인 텍스트 인코더 공간에서 동적인 교차 주의 활성화 공간으로 기저 구성을 이동시킴으로써, 이 방법은 기존 폐형식 기술의 주요 실패 모드인 앵커 템플릿을 벗어난 프롬프트 변형을 처리하지 못하는 문제를 극복합니다.

저자들은 PURE 가 원치 않는 개념을 지우고 모델의 일반적 유용성을 보존하는 것 사이의 우수한 균형을 제공하며 재학습의 계산 오버헤드 없이 개념 망각을 위한 강력하고 효율적이며 최첨단 솔루션을 제공한다고 주장합니다. 이 연구는 "망각 방향"이 입력 텍스트 단독이 아니라 모델의 내부 렌더링 과정에 의해 더 정확하게 정의된다는 것을 시사합니다.