Attention, May I Have Your Decision? Localizing Generative Choices in Diffusion Models

1. 문제 정의 (Problem)

텍스트 - 이미지 확산 모델 (Text-to-Image Diffusion Models) 은 뛰어난 생성 능력을 보이지만, 그 내부 작동 원리는 여전히 불투명합니다 (Black-box). 특히 프롬프트가 불완전하거나 구체적이지 않은 경우 (예: "꽃의 사진"이라고만 입력했을 때 색상, 모양, 배경 등을 명시하지 않은 경우), 모델은 학습 데이터의 패턴을 기반으로 암묵적인 결정 (Implicit Decisions) 을 내려 세부 사항을 채워 넣어야 합니다.

이러한 암묵적인 결정 과정에서 발생하는 문제는 다음과 같습니다:

• 사회적 편향 (Bias): 성별, 인종, 연령 등에 대한 고정관념을 재생산 (예: 의사 프롬프트 시 남성만 생성, 미국 대통령 프롬프트 시 특정 인물로 고정).
• 기존 연구의 한계: 기존 로컬라이제이션 (Localization) 기술은 주로 프롬프트에 명시된 속성 (Explicit Conditioning) 을 분석하는 데 초점을 맞추었습니다. 예를 들어, "남자"라는 단어를 프롬프트에 주입하여 크로스 어텐션 (Cross-Attention) 계층이 어떻게 반응하는지 분석합니다. 그러나 이는 모델이 스스로 결정을 내리는 내부 메커니즘과는 다를 수 있으며, 프롬프트가 구체적이지 않을 때의 결정 과정을 설명하지 못합니다.

2. 방법론 (Methodology)

저자들은 모델이 불완전한 프롬프트에 대해 어떻게 결정을 내리는지 파악하기 위해 ICM (Implicit Choice-Modification) 이라는 새로운 접근법을 제안합니다. 핵심은 선형 프로빙 (Linear Probing) 을 기반으로 한 계층 로컬라이제이션입니다.

A. 선형 프로빙을 통한 계층 선택 (Layer Selection via Linear Probing)

1. 데이터 생성: 속성 (성별, 인종 등) 이 지정되지 않은 일반 프롬프트 (예: "사람의 사진") 로 이미지를 생성합니다.
2. 외부 분류기를 통한 라벨링: 생성된 이미지에 대해 외부 분류기 (예: CLIP 기반) 를 사용하여 실제 생성된 속성 (예: '남자', '여자') 을 사후 (Post-hoc) 에 라벨링합니다. 이는 프롬프트 조작 없이 모델의 실제 선택을 반영합니다.
3. 계층별 구분 가능성 측정: 모델의 중간 활성화 (Intermediate Activations) 를 추출하여 각 계층 (Layer) 에서 속성들이 선형적으로 얼마나 잘 분리되는지 (Discriminability) 를 로지스틱 회귀 프로브로 측정합니다.
4. 결과: 셀프 어텐션 (Self-Attention) 계층이 크로스 어텐션 계층보다 암묵적인 속성 결정에 대해 훨씬 높은 분류 정확도를 보임을 발견했습니다. 이는 모델이 초기 노이즈에서 비롯된 확률적 단서들을 통합하여 최종 결정을 내리는 주체가 셀프 어텐션임을 시사합니다.

B. 생성 제어 (Generation Control)

선택된 계층을 기반으로 두 가지 제어 방식을 적용합니다:

1. 활성화 스티어링 (Activation Steering): 학습된 분류기의 가중치 벡터를 스티어링 벡터로 변환하여, 생성 시 특정 계층의 활성화에 이 벡터를 더하거나 빼서 ($H' = H + \alpha s$) 속성을 변경합니다.
2. 파인튜닝 (Fine-tuning): 선택된 계층의 가중치만 LoRA (Low-Rank Adaptation) 를 통해 미세 조정합니다.

3. 주요 기여 (Key Contributions)

1. 암묵적 결정의 로컬라이제이션: 선형 프로빙을 통해 텍스트 - 이미지 확산 모델에서 불완전한 프롬프트에 대한 결정이 주로 셀프 어텐션 계층에서 이루어짐을 최초로 규명했습니다.
2. 크로스 어텐션 vs 셀프 어텐션의 역할 구분: 명시적 프롬프트는 크로스 어텐션이 처리하지만, 암묵적 결정 (편향 등) 은 셀프 어텐션이 담당함을 실험적으로 증명했습니다.
3. 정밀한 편향 제거 (ICM): 전체 계층이 아닌, 가장 영향력이 큰 소수의 셀프 어텐션 계층만을 대상으로 편향을 수정하는 정밀한 방법론을 제안했습니다.
4. 다양한 아키텍처에서의 확장성: Stable Diffusion (UNet 기반) 뿐만 아니라 SANA (Transformer 기반) 및 SDXL 과 같은 대규모 모델에서도 효과가 입증되었습니다.

4. 실험 결과 (Results)

• 편향 제거 성능: 성별, 연령, 인종 편향 제거 (Debiasing) 실험에서 기존 최첨단 방법 (Finetuning, Latent Editing, DIFFLENS 등) 보다 Fairness Discrepancy (FD) 지수를 크게 개선하면서도 이미지 품질 (FID) 과 프롬프트 일치도 (CLIP-T) 를 유지했습니다.
  • 예: 성별 편향 제거 시 FD 가 0.087 로 감소 (기존 방법 대비 우수) 하였고, 이미지 품질 저하가 최소화되었습니다.
• 계층 선택의 중요성:
  • 셀프 어텐션 vs 크로스 어텐션: 셀프 어텐션 계층을 대상으로 스티어링을 수행했을 때 편향 제거 효과가 훨씬 뛰어났습니다.
  • 최적 계층 vs 무작위/전체 계층: 성능이 가장 좋은 상위 20 개 계층만 선택하여 조작했을 때, 무작위 계층이나 모든 계층을 조작하는 것보다 이미지 품질 저하 없이 더 정확한 생성 결과를 얻었습니다. 모든 계층을 조작하면 아티팩트가 심하게 발생하고 품질이 급격히 떨어집니다.
• 프롬프트 주입 vs 선형 프로빙: 명시적 프롬프트 주입으로 학습된 프로브와 일반 프롬프트 (암묵적 결정) 로 학습된 프로브는 분포가 다르며, 일반 프롬프트로 학습된 프로브를 사용할 때 더 자연스럽고 품질이 높은 편향 제거가 가능함을 확인했습니다.

5. 의의 및 결론 (Significance)

이 연구는 생성형 AI 의 "블랙박스" 내부에서 모델이 어떻게 결정을 내리는지에 대한 메커니즘적 이해를 심화시켰습니다. 특히, 명시적 프롬프트 처리 (크로스 어텐션) 와 암묵적 결정 (셀프 어텐션) 이 서로 다른 계층에서 이루어진다는 발견은 향후 모델 해석 및 제어 기술의 중요한 이정표가 됩니다.

제안된 ICM 방법은 전체 모델을 재학습하거나 모든 계층을 조작하는 비효율적인 방식을 피하고, 정밀하게 선택된 계층만 타겟팅함으로써 편향을 효과적으로 제거하면서도 생성된 이미지의 품질과 원본의 특징을 보존합니다. 이는 생성 모델의 공정성 향상과 안전한 AI 개발에 실질적인 기여를 할 것으로 기대됩니다.