기술 요약: MaskAttn-SDXL

문제 제기

확산 모델, 특히 Stable Diffusion XL(SDXL) 은 텍스트-이미지 생성에서 상당한 성과를 거두었으나, 구조화된 다중 객체 프롬프트를 처리할 때 지속적인 어려움에 직면해 있습니다. 주요 한계점은 다음과 같이 두 가지입니다:

1. 아키텍처 트레이드오프: 지배적인 U-Net 백본은 강한 공간 국소성을 가진 합성곱 레이어에 의존하는데, 이는 학습 안정성과 효율성을 보장하지만 장기 의존성과 전역 구성의 모델링을 방해합니다. 반면, 트랜스포머 기반 확산 모델 (예: DiT) 은 전역 자기-주의 (self-attention) 를 제공하지만, 정밀한 국소 구조를 희생하는 경우가 많으며 상당한 계산 및 메모리 비용을 수반합니다.
2. 구성 신뢰성: 최첨단 모델조차도 복잡한 프롬프트에서 토큰 수준의 일관성을 유지하는 데 자주 실패합니다. 일반적인 실패 양상은 다음과 같습니다:
   • 속성 누출: 객체 간에 색상이나 속성을 잘못 할당함.
   • 공간 위반: 공간 관계 (예: 좌/우 위치) 를 반전시킴.
   • 객체 누락: 요청된 엔티티를 생성하지 못함.
     이러한 오류는 종종 전체 분포 정렬을 측정하는 FID 나 CLIP 점수와 같은 전역 지표로 포착되지 않으며, 이는 특정 토큰과 영역 간의 바인딩을 측정하지 않기 때문입니다. 근본 원인은 토큰 간 간섭으로 식별되며, 이는 텍스트 개념이 공유된 공간 주의력을 위해 경쟁하여 얽힌 시각적 출력을 초래합니다.

방법론: MaskAttn-SDXL

저자들은 SDXL 의 핵심 파이프라인, 사전 훈련된 가중치, 또는 추론 절차를 변경하지 않고 SDXL 을 강화하도록 설계된 경량 플러그인 모듈인 MaskAttn-SDXL을 제안합니다.

핵심 메커니즘

이 방법은 SDXL U-Net 의 크로스-주의 (cross-attention) 로그트 (logits) 에 소프트맥스 (softmax) 연산 전에 토큰 조건부 공간 게이트를 주입합니다.

1. 아키텍처 통합:

   • 수정은 SDXL U-Net 의 선택된 중간 해상도 레이어의 크로스-주의 블록 내에서 구체적으로 적용됩니다. 이 위치는 의미적 추상화와 공간 정밀도 간의 균형을 맞추기 위해 선택되었습니다.
   • 모든 사전 훈련된 구성 요소 (U-Net 백본, 텍스트 인코더, VAE) 는 고정된 상태로 유지됩니다. 오직 경량 "마스크 헤딩"만 학습됩니다.

2. 마스크 생성:

   • 작은 공유 게이트 헤딩 $f(\cdot)$ (예: 얕은 합성곱 레이어 또는 MLP 레이어) 은 현재 잠재 특징 지도 $X$와 특정 토큰 임베딩 $e_t$를 입력으로 받습니다.
   • 이는 각 공간 좌표에서 토큰 $t$가 이미지 생성에 기여해야 하는 정도에 대한 모델의 믿음을 나타내는 연속적인 공간 확률 지도 $\hat{G}_t \in (0, 1)^{H \times W}$를 출력합니다.
   • 이 지도는 0.5 임계값과 그래디언트 흐름을 허용하는 스트레이트-쓰루 추정기 (STE) 를 사용하여 이진화되어, 하드 이진 게이트 $G_t$를 생성합니다.

3. 로그트 변조:

   • 이진 게이트는 가산 마스크 행렬 $M$으로 변환됩니다. 공간 위치 $i$와 토큰 $t$에 대해:
     • 게이트가 활성화된 경우 ($G_t(i)=1$), 마스크 값은 0 입니다.
     • 게이트가 비활성화된 경우 ($G_t(i)=0$), 마스크 값은 $-\infty$(큰 음수 상수로 근사됨) 입니다.
   • 이 마스크는 크로스-주의 로그트에 직접 추가됩니다:
     $$\text{MaskAttn}(Q, K, V; M) = \text{Softmax}\left(\frac{QK^\top}{\sqrt{d}} + M\right)V$$
   • 로그트를 $-\infty$로 설정하면 해당 소프트맥스 가중치를 0 으로 강제하여, 특정 토큰과 특정 공간 영역 간의 상호작용을 효과적으로 억제합니다.

4. 잔차 연결:

   • 마스크된 주의의 출력은 표준 피드포워드 네트워크 (FFN) 를 통과한 후 잔차 연결을 통해 원래 입력과 결합되어, 토큰 - 영역 정렬을 정제하는 동안 원래 정보 흐름이 보존되도록 합니다.

주요 기여

1. 로그트 공간 마스크 주의: 저자들은 학습 가능한 토큰 조건부 마스크를 주의 로그트에 직접 주입하여 SDXL 의 토큰 - 위치 연결을 희석화하는 메커니즘을 제안합니다.
2. 플러그 - 앤 - 플레이 효율성: 이 방법은 보조 입력 (예: 바운딩 박스) 을 필요로 하지 않으며, 매개변수를 거의 추가하지 않습니다 (약 360 만 개), 그리고 사전 훈련된 백본과 표준 샘플링 프로세스를 변경하지 않습니다.
3. 향상된 구성 그라운딩: 이 접근법은 효율성이나 배포 유연성을 희생하지 않으면서 복잡한 다중 엔티티 프롬프트 하에서 공간 및 속성 충실도를 향상시킵니다.

실험 결과

저자들은 MaskAttn-SDXL 을 SD-1.5, SD-2.1, SDXL, PixArt-$\alpha$/$\Sigma$ 및 기타 구성 편집 방법 (예: Attend-and-Excite, Structured Diffusion) 을 포함한 다양한 베이스라인에 대해 평가했습니다.

1. 전역 충실도 및 정렬

MS-COCO 및 Flickr30k 데이터셋에서 평가:

• FID: MaskAttn-SDXL 은 MS-COCO 에서 FID 점수를 SDXL 의 25.77 에서 24.57 로 낮췄으며, Flickr30k 에서는 209.80 에서 206.98 로 낮췄습니다.
• CLIP 점수: MS-COCO 에서 31.53 에서 31.75 로, Flickr30k 에서 33.03 에서 33.54 로 향상되었습니다.
• 정밀도 - 재현율: 동등한 충실도를 유지하면서 향상된 커버리지 (재현율) 를 보여주었습니다.

2. 구성 벤치마크

T2I-CompBench++(공간) 및 GenEval 에서 평가:

• T2I-CompBench++(공간): 0.2248 점수를 달성하여 SDXL(0.2090) 과 모든 기타 베이스라인을 능가했습니다.
• GenEval: 0.66 점수를 달성하여 SDXL(0.62) 과 기타 전문 방법을 능가했습니다.
• 이러한 결과는 속성 누출 및 공간 위반의 상당한 감소를 나타냅니다.

3. 효율성 및 오버헤드

• 매개변수: 26 억 개 매개변수의 SDXL 모델에 약 360 만 개의 매개변수만 추가되었습니다.
• 메모리 및 지연: 메모리 사용량이 약간 증가했습니다 (12.4 GB 에서 12.7 GB 로) 및 이미지당 지연 시간이 0.03 초 증가했습니다 (1.00 초에서 1.03 초로), 이는 방법이 경량임을 확인시켜 줍니다.

4. 제거 연구

• 단계 선택: 중간 해상도 블록이 가장 좋은 구성 점수를 제공하여 기본 배치 전략을 검증했습니다.
• 배치: 인코더와 디코더 블록 모두에 모듈을 삽입하면 전체 정렬이 가장 좋았으나, 주요 이득은 중간 해상도 선택에 기인했습니다.

중요성 및 주장

이 논문은 MaskAttn-SDXL 이 다중 객체 텍스트-이미지 생성에서 토큰 간 간섭이라는 중요한 문제를 해결한다고 주장합니다. 로그트 공간에서 관련 없는 토큰 - 영역 바인딩을 명시적으로 억제함으로써, 이 방법은 대규모 모델을 재학습하는 계산 오버헤드나 휴리스틱 기반 편집 도구의 추론 시간 복잡도 없이 더 강력한 구성 정확성을 달성합니다.

저자들은 이 접근법이 텍스트 전용 환경에서 확산 모델의 신뢰성을 향상시키기 위한 실용적인 해결책을 제공하며, SDXL 파이프라인의 효율성과 유연성을 유지하면서 객체 누락 및 공간 반전과 같은 일반적인 실패 양상을 크게 줄인다고 강조합니다. 결과는 표적화된 희소 주의 변조가 합성곱 확산 백본의 전역 일관성을 향상시키기 위한 실현 가능한 경로임을 시사합니다.