You Don't Need All That Attention: Surgical Memorization Mitigation in Text-to-Image Diffusion Models

이 논문은 텍스트-이미지 확산 모델의 훈련 데이터 암기 현상을 해결하기 위해, 프롬프트별 주의를 정교하게 조절하여 원본 훈련 이미지를 재현하지 않으면서도 프롬프트와 일치하는 고품질 이미지를 생성하는 새로운 프레임워크인 'GUARD'를 제안합니다.

핵심

🎨 제목: "너무 많은 주목은 필요 없어: AI 화가의 '기억'을 수술하는 방법"

1. 문제 상황: AI 의 '불안한 기억력'

생각해 보세요. 어떤 AI 화가가 수만 장의 그림을 보고 배웠다고 칩시다. 그런데 이 화가가 "나에게 '스미스 씨의 초상화'를 그려줘"라고 하면, 스미스 씨가 실제로 찍은 사진과 똑같은 그림을 그려냅니다.

이건 마치 학생이 시험 문제를 외워서 답을 그대로 적어내는 것과 같습니다.

• 문제점: 만약 그 '스미스 씨의 초상화'가 누군가의 사생활이거나 저작권이 있는 그림이라면, AI 가 똑같이 그려내는 것은 사생활 침해나 저작권 위반이 될 수 있습니다.
• 기존 방법들의 한계:
  • 학습 단계에서 막기: 처음부터 AI 를 가르칠 때 "이건 외우지 마"라고 시키는 건, AI 가 아직 어떤 걸 외울지 모르는 상태에서 막무가내로 모든 걸 제한하는 것이어서 AI 의 창의성 (품질) 이 떨어집니다.
  • 학습 후 지우기: 이미 배운 AI 에서 특정 기억을 지우려고 노력하는 건, 머리를 수술하듯 무겁고 비싸며, 완전히 지워지지 않아 다시 튀어나올 수도 있습니다.

2. 새로운 해결책: GUARD (가드)

이 논문은 **"AI 의 머릿속을 고치는 게 아니라, 그림을 그리는 '순간'의 행동만 살짝 수정하자"**고 제안합니다. 이를 GUARD라고 부릅니다.

GUARD 의 원리: '밀어내기'와 '당겨주기'
AI 가 그림을 그리는 과정은 소음을 제거하며 선명한 그림을 만들어가는 과정입니다. GUARD 는 이 과정에서 두 가지 힘을 작용시킵니다.

1. 밀어내기 (Repulsion): "아니야, 그건 원래 훈련 데이터에 있던 그림이야! 그쪽으로 가면 안 돼!"라고 AI 를 원래의 훈련 데이터 (기억) 에서 밀어냅니다.
2. 당겨주기 (Attraction): "그 대신, 너의 요청 (프롬프트) 에 맞는 새롭고 멋진 그림으로 가자!"라고 안전한 목표 지점으로 AI 를 끌어당깁니다.

이 두 가지 힘을 동시에 쓰면, AI 는 원래의 그림을 베끼지 않으면서도 사용자의 요청에 맞는 고품질 그림을 그릴 수 있게 됩니다.

3. 핵심 기술: "어디를 수술할지 찾아내는 정밀 탐정"

그런데 여기서 중요한 질문이 생깁니다. "어떤 부분을 밀어내고 당겨줘야 할까?"

기존 연구들은 "글의 끝에 있는 특정 단어 (예: 문장 부호) 만 보면 안 돼"라고 정해두었습니다. 하지만 이 논문은 **"그건 틀렸어! 기억을 유발하는 단어는 그림마다, 상황마다 달라"**라고 말합니다.

• 비유: AI 가 그림을 그릴 때, 어떤 단어는 마치 마법의 지팡이처럼 작동해서 훈련 데이터를 바로 불러옵니다.
  • 어떤 그림에서는 "고양이"라는 단어가 지팡이 역할을 할 수도 있고,
  • 다른 그림에서는 "파란색"이라는 단어가 지팡이 역할을 할 수도 있습니다.
  • 심지어 문장 끝의 부호 (EOT 토큰) 가 지팡이 역할을 하기도 합니다.

GUARD 의 혁신:
이 논문은 **"실시간 탐정"**을 도입했습니다. AI 가 그림을 그리는 순간순간, **"어떤 단어가 지금 가장 큰 주목 (Attention) 을 받고 있어?"**를 분석합니다.

• 만약 특정 단어가 비정상적으로 큰 주목을 받고 있다면, 그 단어가 **기억을 불러일으키는 '지팡이'**일 가능성이 높습니다.
• GUARD 는 그 특정 단어의 주목도만 '수술'하듯 살짝 줄여줍니다. (나머지 단어들은 그대로 둡니다.)

이렇게 하면 AI 는 "아, 이 단어는 너무 크게 보지 말아야겠다"라고 생각하며, 훈련 데이터를 베끼지 않는 새로운 그림을 그리게 됩니다.

4. 왜 이것이 특별한가요?

• 정밀한 수술 (Surgical): AI 전체를 고치는 게 아니라, 필요한 부분만 딱 집어서 수정합니다. 그래서 그림의 품질이 떨어지지 않습니다.
• 유연함 (Dynamic): 그림마다, 요청마다 달라지는 '기억의 지팡이'를 실시간으로 찾아냅니다.
• 효율성: AI 를 다시 학습시킬 필요 없이, 그림을 그리는 순간에만 적용하면 됩니다.

📝 한 줄 요약

"AI 가 훈련 데이터를 그대로 베끼는 것을 막기 위해, 그림을 그리는 순간 '기억을 불러오는 특정 단어'만 실시간으로 찾아내어 그 힘만 살짝 누르는 정밀 수술법을 개발했습니다."

이 방법은 AI 가 사용자의 요청에 따라 멋진 그림을 그리면서도, 저작권이나 사생활 문제를 일으키는 원본 그림을 베끼지 않도록 도와줍니다. 마치 AI 화가에게 "그림은 잘 그려줘도 되는데, 그 특정 사진은 절대 베끼지 마!"라고 귀에 대고 속삭여주는 것과 같습니다.

기술 요약

1. 문제 정의 (Problem Statement)

텍스트 - 이미지 (T2I) 생성 모델, 특히 확산 모델 (Diffusion Models) 은 학습 데이터의 특정 예시를 그대로 복제하거나 거의 그대로 복제하는 '기억 (Memorization)' 현상을 보입니다. 이는 다음과 같은 심각한 문제를 야기합니다.

• 개인정보 유출 및 저작권 침해: 학습에 사용된 민감한 이미지나 저작권이 있는 예술 작품이 원본 그대로 생성될 수 있습니다.
• 기존 방법론의 한계:
  • 학습 시간 (Training-time) 개입: 어떤 데이터가 기억될지 미리 알기 어렵고, 학습 과정을 통제하기 어려운 경우가 많으며, 모델의 유용성 (Utility) 을 떨어뜨릴 수 있습니다.
  • 파인튜닝 시간 (Finetuning-time) 개입 (Unlearning): 계산 비용이 매우 높고, 학습된 정보를 완전히 지우지 못해 나중에 다시 나타날 수 있는 불안정성 (Robustness 부족) 이 있습니다.

이에 따라, 모델 가중치를 변경하지 않고 추론 시간 (Inference-time) 에만 작동하여, 메모리된 정보를 생성 과정에서 차단하면서도 이미지 품질을 유지하는 정밀하고 효율적인 해결책이 필요합니다.

2. 제안 방법론: GUARD (Methodology)

저자들은 GUARD (Guidance Using Attractive-Repulsive Dynamics) 라는 새로운 추론 시간 프레임워크를 제안합니다. 이는 기존의 Classifier-Free Guidance (CFG) 를 변형하여 작동합니다.

2.1. GUARD 프레임워크의 핵심 원리

표준 CFG 는 무조건적 예측 (empty prompt) 과 조건부 예측 (text prompt) 의 선형 결합을 사용합니다. GUARD 는 이를 다음과 같이 수정합니다:

• 반발력 (Repulsion): 기억된 프롬프트에서 생성된 노이즈 예측 ($\epsilon^-$) 을 반대 방향으로 밀어냅니다. (기억된 이미지를 재현하려는 경향을 억제)
• 인력 (Attraction): 새로운 양수 목표 (Positive Target) 인 노이즈 예측 ($\epsilon^+$) 을 끌어당깁니다. (기억된 이미지와 구별되지만 프롬프트와 일치하는 고품질 이미지를 유도)

수식적으로는 다음과 같이 표현됩니다:
$$\hat{\epsilon} = \epsilon_\theta(x_t, e_\phi) + s(\epsilon^+_\theta - \epsilon_\theta) - r(\epsilon^-_\theta - \epsilon_\theta)$$
여기서 $s$는 인력 강도, $r$은 반발력 강도입니다.

2.2. CA-in-GUARD: 구체적인 구현 (Surgical Attention Attenuation)

GUARD 의 성공적인 구현을 위해 '양수 목표 ($\epsilon^+$)'를 어떻게 정의할지가 핵심입니다. 저자들은 크로스 어텐션 (Cross-Attention, CA) 의 스파이크 (Spike) 를 분석하여 이를 해결합니다.

• 분석 결과:
  • Verbatim Memorization (직역 기억): EOT(End-of-Text) 토큰뿐만 아니라 다른 특정 토큰들에서도 비정상적으로 큰 어텐션 스파이크가 발생합니다.
  • Template Memorization (템플릿 기억): EOT 토큰의 어텐션은 오히려 낮을 수 있으며, 다른 토큰들이 기억의 핵심을 담당합니다.
  • 결론: 고정된 규칙 (예: EOT 만 줄이기) 은 작동하지 않으며, 프롬프트마다 동적으로 기억에 중요한 토큰을 찾아야 합니다.
• 동적 스파이크 감지 (Dynamic Spike Detection):
  • 추론 과정에서 프롬프트별 어텐션 분포를 실시간으로 분석합니다.
  • 통계적 이상치 탐지 (Z-score) 를 통해 프롬프트마다 기억에 결정적인 토큰 위치 $S(p)$ 를 자동으로 식별합니다.
• 정밀한 감쇠 (Surgical Attenuation):
  • 식별된 $S(p)$ 위치의 크로스 어텐션 로짓 (Logits) 을 특정 계수 ($\alpha$) 로 감쇠시킵니다.
  • 이는 모델이 기억된 이미지를 재현하는 데 필요한 '과도한 어텐션'을 제거하면서도, 프롬프트의 의미적 일치는 유지하도록 합니다.

3. 주요 기여 (Key Contributions)

1. GUARD 프레임워크: 기억된 방향으로부터의 반발과 안전한 목표로의 인력을 결합한 새로운 추론 시간 완화 프레임워크 제안.
2. 통계적 어텐션 스파이크 감지: 고정된 규칙이 아닌, 프롬프트별 통계적 이상치 분석을 통해 기억의 핵심 토큰을 동적으로 탐지하는 메커니즘 개발.
3. 정밀한 개입 (Surgical Mitigation): 불필요한 품질 저하 없이 기억된 정보만 선택적으로 차단하는 CA 감쇠 메커니즘 구현.
4. 포괄적인 평가: 다양한 아키텍처 (SD v1.4, v2.0) 와 기억 유형 (Verbatim, Template) 에서 기존 최첨단 방법론 (SOTA) 을 능가하는 성능 입증.

4. 실험 결과 (Results)

저자들은 Stable Diffusion v1.4 와 v2.0 을 사용하여 다양한 메트릭 (SSCD, CLIP, FID) 으로 평가했습니다.

• 기억 완화 성능 (SSCD):
  • 제안된 CA-in-GUARD는 모든 설정에서 기존 SOTA 방법들 (Ren et al., Wen et al., Han et al. 등) 보다 가장 낮은 SSCD 점수를 기록했습니다.
  • 특히 기존 방법들이 취약했던 Template Memorization과 SD v2.0 환경에서 압도적인 성능 향상을 보였습니다. (예: SD v2.0 Template 에서 SSCD 0.36 $\to$ 0.19 로 대폭 감소).
• 이미지 품질 (CLIP, FID):
  • 단순히 어텐션을 줄이는 것만으로는 FID(실제성) 가 일부 저하될 수 있으나, GUARD 의 '인력 (Attraction)' 항을 도입함으로써 품질 저하를 보정하고 오히려 기존 방법들보다 우수한 CLIP 점수 (프롬프트 일치도) 와 FID 점수를 달성했습니다.
• 비기억 프롬프트에 대한 영향:
  • 기억되지 않은 프롬프트에 적용했을 때에도 생성 품질에 유의미한 저하가 없음을 확인하여, 이 방법이 보편적으로 적용 가능함을 입증했습니다.
• 효율성:
  • 배치 처리 (Batching) 를 통해 여러 경로를 한 번의 포워드 패스로 처리하므로, 계산 오버헤드가 상대적으로 적습니다.

5. 의의 및 결론 (Significance)

이 논문은 T2I 확산 모델의 기억 문제를 해결하기 위해 모델 재학습 없이 추론 시간에만 개입하는 '외과적 (Surgical)' 접근법의 가능성을 입증했습니다.

• 실용성: 기존 학습 데이터에 대한 접근 권한이 없거나, 파인튜닝 비용이 부담스러운 실제 배포 환경에서도 즉시 적용 가능합니다.
• 정밀성: "기억"이라는 현상이 단순히 EOT 토큰에 국한되지 않으며, 프롬프트마다 동적으로 변한다는 통찰을 바탕으로, 정확한 위치를 타격하여 부작용을 최소화했습니다.
• 향후 방향: 추론 시간 완화의 원리를 학습 시간이나 파인튜닝 시간 방법론과 결합하여, 더 강력한 완화 효과를 기대할 수 있는 방향을 제시합니다.

요약하자면, GUARD 는 **"기억된 데이터의 재현을 막기 위해 불필요한 어텐션 (Attention) 을 제거하라"**는 핵심 아이디어를 통해, 기억 완화와 이미지 품질이라는 상충되는 두 목표를 동시에 달성한 획기적인 방법론입니다.