1. 문제 정의 (Problem)

최근 고충실도 이미지 생성을 위한 확산 모델 (Diffusion Models), 특히 Diffusion Transformers (DiT) 기반의 모델 (예: FLUX, SD3.5) 이 주류로 자리 잡았습니다. 그러나 이러한 모델의 능력 향상과 함께, 유해한 콘텐츠 (성적 노출, 저작권 침해 등) 나 원치 않는 개념이 생성되는 위험도 증가했습니다.

기존의 개념 제거 (Concept Erasure, CE) 기법들은 주로 모델의 가중치를 최적화 (Fine-tuning) 하거나 경사 기반 (Gradient-based) 으로 조정하는 방식을 사용했습니다. 하지만 이 방법들은 다음과 같은 한계가 있습니다:

• 계산 비용: 현대의 대규모 DiT 모델 (예: 12B 파라미터 FLUX) 에 적용 시 재학습 비용이 prohibitively expensive(부담스러울 정도로 비쌈).
• Distilled 모델의 제약: Guidance-destilled 된 모델 (예: Flux.1[schnell]) 은 고정된 가이드 스케일에서 작동하도록 설계되어, 기존 방법들이 의존하는 Negative Guidance 기반 최적화가 어렵습니다.
• 아키텍처 차이: 기존 UNet 기반 모델의 Cross-attention 조작 방식은 텍스트와 이미지를 통합 처리하는 DiT 의 단일 어텐션 메커니즘에는 직접 적용하기 어렵습니다.

2. 방법론 (Methodology)

저자들은 SHIFT라는 새로운 프레임워크를 제안합니다. 이는 대규모 언어 모델 (LLM) 의 'Activation Steering(활성화 조종)' 아이디어를 DiT 에 적용한 것으로, 모델의 가중치를 변경하지 않고 추론 (Inference) 시 중간 활성화 (Intermediate Activations) 를 조작하여 원치 않는 개념을 억제하거나 원하는 스타일로 전환합니다.

핵심 기술 요소:

1. 조종 지점 (Steering Locations):

   • DiT 아키텍처에서 텍스트와 이미지가 통합된 공간에서 작동하므로, **텍스트 인코더의 풀드 임베딩 (Pooled Embedding)**과 DiT 백본 내의 텍스트 토큰 활성화 두 지점을 대상으로 합니다.
   • 특히 Double-stream block 에서 텍스트 토큰의 어텐션 출력에 개입합니다.

2. 조종 벡터 (Steering Vector) 구축:

   • 데이터셋: 목표 개념이 포함된 프롬프트 (Positive) 와 포함되지 않은 프롬프트 (Negative) 의 쌍을 생성합니다.
   • 벡터 추정:
     • SVM 기반: 두 클래스의 활성화 분포를 분리하는 초평면 (Hyperplane) 의 법선 벡터를 계산합니다.
     • 평균 차이 (Mean Difference): 두 프롬프트 쌍의 활성화 평균 차이를 직접 계산합니다.
   • 이 벡터는 개념의 의미적 방향을 나타내며, 추론 시 활성화에 가산됩니다.

3. 추론 시 적용 (Inference-time Application):

   • 수식: $\tilde{a} = a + \alpha \times v_{\text{steering}}$
   • 강도 조절 ($\alpha$):
     • 텍스트 인코더: 초기 프롬프트와 목표 개념 간의 코사인 유사도를 기반으로 강도를 조절하여, 비목표 개념의 손실을 방지합니다.
     • DiT 활성화: 경량화된 SVM 분류기의 점수를 실시간으로 모니터링하여, 목표 개념이 여전히 존재할 때 조종 강도를 증폭시키고, 제거되었을 때는 약화시킵니다 (Over-steering 방지).

4. 시간적 불변성 (Temporal Invariance):

   • 실험 결과, DiT 의 잠재 공간에서 의미적 개념은 시간에 무관한 (Time-independent) 안정적인 방향으로 인코딩됨을 발견했습니다. 따라서 모든 확산 시간 단계 (Timesteps) 에 동일한 조종 벡터를 적용할 수 있어 프로세스가 단순화됩니다.

3. 주요 기여 (Key Contributions)

1. 통합 조종 프레임워크: 재학습 없이 Flux.1[dev] 및 Flux.1[schnell] 과 같은 원본 및 증류된 DiT 모델 모두를 대상으로 개념 조작을 가능하게 하는 최초의 포괄적인 프레임워크입니다.
2. 시공간 역학 분석: DiT 활성화 공간에서 의미적 개념이 전역적이고 안정적인 방향으로 인코딩됨을 규명했습니다. 단일 벡터로 모든 시간 단계를 제어할 수 있음을 증명했습니다.
3. 교차 증류 벡터 전이 (Cross-distillation Vector Transfer): 증류된 모델 (Flux.1[schnell]) 에서 학습된 조종 벡터가 비증류 모델 (Flux.1[dev]) 로도 효과적으로 전이될 수 있음을 입증했습니다. 이는 리소스가 제한된 환경에서 유연한 제어를 가능하게 합니다.

4. 실험 결과 (Results)

저자들은 다양한 작업 (개념 제거, 스타일 전환, 객체 추가/변경) 에서 SHIFT 를 평가했습니다.

• 추상적 개념 제거 (성적 노출 등):
  • I2P 벤치마크에서 NudeNet 검출기를 사용하여 평가한 결과, 기존 최첨단 방법 (ESD, EAP, CA, UCE) 보다 3~4 배 이상 강력한 억제 효과를 보였습니다.
  • 동시에 CLIP 점수는 거의 변하지 않아 프롬프트 준수도를 유지했고, FID 점수 저하도 미미하여 이미지 품질을 잘 보존했습니다.
• 구체적 객체 제거 (Snoopy 등):
  • Snoopy 개념을 제거하면서도 Mickey, SpongeBob 등 유사 개념은 보존하는 능력을 입증했습니다.
  • 기존 방법들보다 제거된 개념의 CLIP 점수 감소 (제거 효과) 와 비목표 개념의 CLIP 점수 유지 (보존 효과) 에서 우수한 성능을 보였습니다.
• 스타일 및 객체 조작:
  • 화가 스타일 (반 고흐 등) 제거, 모자/안경/입술 등 작은 객체 제거, 그리고 '사과 추가'나 '개→고양이 변경'과 같은 개념 추가/전환 작업에서도 효과적으로 작동함을 보여주었습니다.
• 효율성:
  • 모델 재학습이 필요 없으며, 추론 시에만 벡터를 적용하므로 매우 가볍고 확장 가능합니다.

5. 의의 및 결론 (Significance)

• 재학습 없는 제어: 대규모 DiT 모델의 안전성 및 제어 문제를 해결하기 위해 고비용의 재학습 (Fine-tuning) 에 의존하지 않는 새로운 패러다임을 제시했습니다.
• DiT 아키텍처 적합성: 텍스트와 이미지가 통합된 DiT 의 복잡한 잠재 공간에서도 의미적 방향성을 효과적으로 찾아내어 조작할 수 있음을 증명했습니다.
• 실용성: 증류된 모델 (Flux.1[schnell]) 에서 학습된 지식을 비증류 모델로 전이할 수 있어, 다양한 환경에서 유연하게 적용 가능한 솔루션을 제공합니다.
• 한계 및 향후 과제: SHIFT 는 이미지 편집 도구 (배경 유지 등) 가 아닌 생성 방향 전환 도구이므로, 대규모 도메인 전환 시 구조적 일관성 유지에는 여전히 과제가 남아있으나, 향후 기하학적 정렬을 통해 개선될 수 있음을 시사합니다.

이 논문은 생성형 AI 의 안전성과 제어 가능성을 높이기 위해, 모델 내부의 표현 공간 (Representation Space) 을 직접적으로 이해하고 조종하는 'Activation Steering' 기법의 중요성을 DiT 시대에 부각시켰다는 점에서 의의가 큽니다.