논문 요약: Diffusion Transformers 의 공간 관계 생성에 대한 회로 메커니즘 (Circuit Mechanisms for Spatial Relation Generation in Diffusion Transformers)

1. 연구 배경 및 문제 정의

텍스트 - 이미지 생성 모델 (Text-to-Image, T2I) 인 Diffusion Transformers(DiT) 는 고해상도 이미지 생성 분야에서 혁신적인 성과를 거두었으나, 텍스트 프롬프트에 명시된 여러 객체 간의 공간적 관계 (Spatial Relations) 를 정확하게 구현하는 데에는 여전히 한계가 있습니다. 예를 들어, "빨간 사각형이 파란 원의 왼쪽 위에 있다"는 프롬프트를 입력했을 때, 객체의 모양과 색상은 잘 생성되더라도 두 객체의 상대적 위치가 틀리는 경우가 빈번합니다.

기존 연구들은 주로 어텐션 맵 조작이나 레이아웃 조건부 학습 등을 통해 이 문제를 해결하려 했으나, 모델 내부에서 공간 관계가 어떻게 계산되고 구현되는지에 대한 기계적 해석 (Mechanistic Interpretability) 수준에서의 이해는 부족했습니다. 본 연구는 다양한 텍스트 인코더를 사용하여 DiT 가 어떻게 공간 관계를 학습하고 생성하는지 그 내부 회로 (Circuit) 를 규명하는 것을 목표로 합니다.

2. 연구 방법론 (Methodology)

2.1. 통제된 실험 환경 구축

• 데이터셋: 두 개의 객체 (모양: 원, 삼각형, 사각형 / 색상: 빨강, 파랑) 와 8 가지 공간 관계 (위, 아래, 왼쪽, 오른쪽, 대각선 등) 를 포함하는 합성 데이터셋을 직접 구축했습니다. 프롬프트는 "[색상] [모양] [관계] [색상] [모양]" 형식으로 구성되었습니다.
• 모델 아키텍처: PixArt 아키텍처를 기반으로 한 DiT 모델을 처음부터 (from scratch) 훈련시켰습니다. 모델 크기 (Nano, Micro, Mini, B) 를 다양하게 변경하여 스케일 효과를 분석했습니다.
• 텍스트 인코더 비교: 모델의 내부 메커니즘이 텍스트 인코더의 선택에 따라 어떻게 달라지는지 비교하기 위해 세 가지 조건을 설정했습니다.
  1. RTE (Random Token Embedding): 의미론적 구조가 없는 무작위 벡터에 위치 인코딩을 추가한 경우.
  2. T5: 사전 훈련된 T5-XXL 인코더를 사용한 경우 (의미론적 정보가 혼합된 토큰 임베딩).
  3. CLIP: 사전 훈련된 CLIP 인코더를 사용한 경우 (보조 실험).

2.2. 분석 기법

• Attention Synopsis (어텐션 요약): 수백만 개의 어텐션 맵을 체계적으로 요약하여, 텍스트 토큰과 이미지 토큰 간의 정보 흐름을 정량화하고 관련 회로를 식별했습니다.
• QK 회로 분석: Query(Q) 와 Key(K) 의 선형 변환을 통해 위치 임베딩과 텍스트 임베딩이 어떻게 상호작용하여 공간적 기울기 (Spatial Gradients) 를 생성하는지 분석했습니다.
• 인과적 조작 (Causal Manipulation): 특정 어텐션 헤드를 제거 (Ablation) 하거나, 임베딩 벡터에 연산을 가해 (Vector Arithmetic) 생성 결과에 미치는 인과적 영향을 검증했습니다.
• 분산 분할 (Variance Partitioning): T5 임베딩 내에서 객체 정보와 공간 관계 정보가 어떻게 분리되거나 혼합되어 있는지 통계적으로 분석했습니다.

3. 주요 결과 및 발견 (Key Results)

3.1. 텍스트 인코더에 따른 회로 메커니즘의 차이

두 가지 설정 (RTE 와 T5) 모두 높은 정확도로 공간 관계 생성 과제를 수행했으나, 그 내부 메커니즘은 극명하게 달랐습니다.

• RTE-DiT (무작위 임베딩): 2 단계 회로 (Two-Stage Circuit)

  • 1 단계 (공간 관계 헤드): 텍스트의 관계 단어 (예: "위", "왼쪽") 를 읽어 이미지 토큰의 위치 임베딩과 상호작용하여 **공간적 기울기 (Spatial Gradients)**를 생성합니다. 이는 캔버스상의 특정 영역을 "태그"로 표시하는 역할을 합니다. (예: Layer 2, Head 8)
  • 2 단계 (객체 생성 헤드): 생성된 공간 태그가 있는 영역에서, 해당 객체의 모양 (Shape) 정보를 읽어 객체를 생성합니다. (예: Layer 4, Head 3)
  • 특징: 관계 정보와 객체 속성 정보가 분리되어 처리되며, 모듈화된 구조를 가집니다.

• T5-DiT (사전 훈련된 인코더): 단일 토큰 통합 회로 (Single-Token Fusion Circuit)

  • T5 의 자기 어텐션 (Self-attention) 을 통해 문장 전체의 맥락이 각 토큰에 혼합됩니다.
  • DiT 는 별도의 관계 단어 토큰을 직접 읽지 않고, **두 번째 객체의 모양 토큰 (Shape2 token)**의 임베딩에 이미 공간 관계 정보가 인코딩되어 있는 것을 발견했습니다.
  • 메커니즘: 특정 헤드가 Shape2 토큰의 임베딩에서 관계 요소를 추출하여 공간적 기울기를 생성합니다. 즉, 관계와 객체 정보가 하나의 토큰 임베딩에 융합되어 있습니다.

3.2. 강건성 (Robustness) 및 일반화 능력

• RTE-DiT: 프롬프트에 불필요한 단어 (예: "the" 추가) 가 삽입되거나 문장 구조가 약간 변경되어도 강건하게 성능을 유지합니다. 관계와 객체 정보가 분리되어 있기 때문입니다.
• T5-DiT: 프롬프트의 작은 변화 (filler word 추가, 동의어 사용 등) 에 매우 **취약 (Fragile)**합니다. 불필요한 단어가 Shape2 토큰의 임베딩 분포를 교란시켜, 모델이 잘못된 공간 관계를 추론하게 만듭니다. 이는 사전 훈련된 언어 모델의 임베딩이 DiT 의 공간 추론에 병목 현상을 일으킬 수 있음을 시사합니다.

3.3. 모델 크기의 영향

• RTE-DiT 의 경우, 모델 크기가 작아질수록 (Nano) 공간 관계 헤드가 사라지고 성능이 급격히 떨어졌습니다.
• T5-DiT 는 상대적으로 작은 모델에서도 높은 성능을 보였으나, 이는 T5 가 이미 복잡한 관계를 압축하여 제공하기 때문으로 해석됩니다.

4. 주요 기여 (Contributions)

1. 기계적 해석을 통한 공간 관계 생성 메커니즘 규명: DiT 가 비가환적 (Non-commutative) 인 객체 간 관계를 어떻게 구현하는지에 대한 구체적인 회로 (Spatial Relation Head, Object Generation Head) 를 최초로 제시했습니다.
2. 텍스트 인코더의 결정적 역할: 동일한 생성 태스크라도 텍스트 인코더 (RTE vs T5) 에 따라 완전히 다른 내부 회로가 학습됨을 증명했습니다. 이는 텍스트 임베딩 설계가 모델의 인덕티브 바이어스 (Inductive Bias) 와 해석 가능성에 직접적인 영향을 미친다는 점을 강조합니다.
3. 새로운 분석 도구 개발: 확산 모델의 반복적 샘플링 특성을 극복하기 위해 어텐션 맵을 요약하고 (Attention Synopsis), 가중치 공간에서 공간 관계 헤드를 식별하는 효율적인 방법 (Weight-space Head Screening) 을 제안했습니다.
4. 강건성 문제 제기: 높은 정확도에도 불구하고 T5 기반 모델이 프롬프트의 작은 변화에 취약하다는 사실을 발견하여, 실제 응용 시 발생할 수 있는 실패 원인을 규명했습니다.

5. 의의 및 결론 (Significance)

본 연구는 텍스트 - 이미지 생성 모델이 복잡한 공간 관계를 학습하는 방식에 대한 깊은 통찰을 제공합니다. 특히, 사전 훈련된 언어 모델 (T5 등) 을 사용할 경우, 모델이 관계 정보를 객체 토큰에 압축하여 처리함으로써 효율성은 높일 수 있지만, 프롬프트의 미세한 변화에 대한 강건성을 잃을 수 있음을 보여줍니다.

이는 향후 T2I 모델 설계 시, 단순히 성능만 높이는 것이 아니라 어떤 텍스트 인코더를 사용할지, 그리고 공간 관계를 명시적으로 분리하여 처리할지 (RTE 방식) 아니면 통합할지 (T5 방식) 에 대한 전략적 선택이 중요함을 시사합니다. 또한, 생성 모델의 실패 원인을 단순히 "데이터 부족"이나 "레이아웃 조건부 부족"이 아니라, 임베딩 모델의 표현 한계와 DiT 의 회로 설계 간의 불일치에서 찾을 수 있음을 보여줍니다.

이 연구는 생성 모델의 내부 작동 원리를 이해하고, 더 강력하고 해석 가능한 멀티모달 AI 를 개발하는 데 중요한 기초를 마련했습니다.