Untwisting RoPE: Frequency Control for Shared Attention in DiTs

기술 요약: Untwisting RoPE: DiT의 공유 어텐션을 위한 주파수 제어

문제 정의

Flux와 같은 모델에서 사용되는 MMDiT(Multi-modal Diffusion Transformer) 아키텍처와 같은 확산 트랜스포머(Diffusion Transformers, DiTs)에서, 위치 정보는 회전 위치 임베딩(Rotary Positional Embeddings, RoPE)을 통해 인코딩됩니다. RoPE는 표준 이미지 생성에서 공간적 구조를 모델링하는 데 필수적이지만, 타겟 이미지 생성이 참조 이미지의 토큰에 어텐션을 수행하는 공유 어텐션(shared-attention) 설정에서의 동작은 여전히 제대로 이해되지 않고 있습니다.

저자들은 이 설정에서 발생하는 결정적인 실패 모드인 참조 복제(reference copying) 현상을 식별했습니다. 즉, 참조 이미지의 스타일적 속성을 타겟 프롬프트로 전달하는 대신(예: 황소의 스타일로 기린을 생성), 모델이 참조 이미지의 콘텐츠를 거의 그대로 재현하며 붕괴되는 현상입니다. 이는 RoPE가 도입하는 강력한 공간적 편향(spatial bias)으로 인해, 타겟 이미지의 쿼리(query)가 의미론적 유사성보다 공간적으로 정렬된 참조 토큰에 불균형적으로 집중하게 만들기 때문입니다.

이를 완화하기 위한 기존의 시도들(예: 디노이징 단계 후반부에만 어텐션 공유를 제한하거나, AlignedGen과 같이 이동된 위치 인덱스를 사용하는 방식)은 거친 수준의 제어만을 제공하거나 "고스트 사이클릭 인페인팅(ghost-cyclic inpainting, 모델이 참조 이미지를 타겟의 연장선으로 취급하는 현상)"과 같은 새로운 아티팩트를 유발합니다.

방법론

본 논문은 모델을 재학습하거나 아키텍처를 변경하지 않고도 의미 있는 공유 어텐션을 복구하기 위해, RoPE의 **주파수 인식 변조(frequency-aware modulation)**를 제안합니다.

1. RoPE 주파수 구조 분석

저자들은 RoPE가 뚜렷한 위치 민감도를 가진 주파수 성분들로 분해된다는 원리적인 분석을 제공합니다.

• 고주파 성분(High-frequency components): 임베딩의 낮은 인덱스 차원에 해당합니다. 이들은 위치 변위($\Delta$)에 따라 각도가 급격하게 변화하며, 강력한 **국부성(locality)**을 강제합니다. 이로 인해 공간적 거리가 멀어짐에 따라 어텐션이 가파르게 감소합니다.
• 저주파 성분(Low-frequency components): 높은 인덱스 차원에 해당합니다. 이들은 변위에 따라 느리게 변화하며, 약한 위치 편향을 제공하고 **전역적 공간 일관성(global spatial coherence)**을 유지합니다.

공유 어텐션 시나리오에서는 고주파 성분이 어텐션 계산을 지배하여, 모델이 공간적으로 정렬된 참조 토큰에 고착되어 그 콘텐츠를 복제하게 만듭니다.

2. 주파수 인식 변조

이를 해결하기 위해 저자들은 어텐션 계산 중에 **참조 이미지의 키($K_{ref}$)**의 주파수 대역을 선택적으로 변조하는 방법을 도입했습니다.

• 메커니즘: RoPE를 비활성화하거나(이는 퇴보적인 생성을 초래함) 어텐션을 균일하게 스케일링하는 대신, 이 방법은 RoPE 스펙트럼 전체에 걸쳐 스케일링 인자의 다항식 보간(polynomial interpolation)을 적용합니다.
• 스케일링 전략:
  • 감쇄(Attenuation): 고주파 성분은 스케일 다운($s_{hf} \in (0, 1)$)하여 복제를 유도하는 강력한 위치 편향을 줄입니다.
  • 증폭(Amplification): 저주파 성분은 스케일 업($s_{lf} > 1$)하여 전역적이고 의미론적으로 유도된 어텐션을 촉진합니다.
• 타임스텝 스케줄링(Timestep Scheduling): 변조 스케일($s_{hf}$ 및 $s_{lf}$)은 디노이징 타임스텝에 따라 선형적으로 증가합니다. 초기 타임스텝에는 낮은 변조를 사용하여 전역 구조와 광범위한 스타일 가이드를 확립하고, 후기 타임스텝에는 스케일링을 높여 미세한 스타일적 세부 사항과 질감을 전달할 수 있도록 합니다.

최종적인 공유 어텐션용 어텐션 키는 다음과 같이 공식화됩니다:
K = [K'_{img}^{tar} \oplus K'_{txt}^{tar}] \oplus s_{FA}^t \cdot K_{img}^{ref}
여기서 $s_{FA}^t$는 타임스텝 $t$에서의 주파수 인식 변조 스케일이며, 오직 참조 키에만 적용됩니다.

주요 기여

1. 공유 어텐션에서의 RoPE에 대한 원리적 분석: 본 논문은 DiT에서의 "참조 복제" 현상이 일반적인 어텐션 공유의 실패가 아니라, 고주파 성분이 엄격한 공간적 국부성을 강제함으로써 발생하는 특정 결과임을 입증합니다.
2. 주파수 제어 메커니즘: 저자들은 RoPE 주파수 대역을 변조할 수 있는 제로샷(zero-shot), 학습이 필요 없는(training-free) 방법을 도입했습니다. 이를 통해 위치 정렬과 의미론적 유사성 사이의 조절 가능한 트레이드오프를 제공합니다.
3. 아티팩트 해결: 이 방법은 단순히 좌표를 이동시키는 것이 아니라 근본적인 주파수 구조를 직접 다룸으로써, 기존 방식(shifted RoPE 등)에서 관찰된 "고스트 사이클릭 인페인팅" 및 콘텐츠 누출 문제를 해결합니다.
4. 스타일 정렬 생성의 복원: 이 접근 방식은 DiT가 이전의 UNet 기반 모델이나 과도한 미세 조정이 필요했던 방식과 달리, 참조로부터 일관된 시각적 스타일을 유지하면서도 다양한 의미적 콘텐츠(예: 서로 다른 동물)를 생성할 수 있게 합니다.

결과

Flux 1.-dev 모델을 대상으로 평가한 결과, 본 방법은 다음과 같은 성능을 보여주었습니다:

• 효과적인 스타일 전이: 모델은 참조 이미지의 콘텐츠를 복제하지 않으면서도 참조(예술 작품 및 사진 포함)로부터 스타일적 속성(색상, 질감, 붓터치 등)을 성공적으로 전달합니다.
• 제어 가능성: 사용자는 변조 스케일을 조정함으로써 스타일 전이와 콘텐츠 보존 사이의 정도를 제어할 수 있습니다. 높은 $s_{hf}$는 참조의 구조(포즈/형태)를 보존하면서 스타일을 전달하고, 낮은 $s_{hf}$는 더 많은 구조적 변화를 허용합니다.
• 베이스라인 대비 우수성:
  • **StyleAligned(UNet 기반)**와 비교하여, 이 방법은 DiT에서 발생하는 콘텐츠 누출 및 구조적 불규칙성을 방지합니다.
  • **AlignedGen(shifted RoPE)**과 비교하여, 이 방법은 주기적인 복제 아티팩트와 고스트 사이클릭 인페인팅을 피하며 고해상도에서도 매끄러운 스타일 전이를 제공합니다.
  • IP-Adapter 및 B-LoRA와 비교하여, 이 방법은 추가 학습이나 이미지 인코더를 필요로 하지 않는 더 일반화된 제로샷 솔루션을 제공합니다.

의의

본 논문은 위치 인코딩이 중립적인 구성 요소가 아니라, 생성 모델 내에서 **국부성(locality)**과 의미론적 연관성(semantic association) 사이의 균형을 능동적으로 조절한다는 점을 주장합니다. 연구 결과는 왜 (합성곱 유도 편향에 의존하는) UNet 기반 모델에서 개발된 어텐션 공유 전략이 (RoPE에 전적으로 의존하는) DiT에 직접 적용될 때 실패하는지를 설명해 줍니다.

고주파 RoPE 성분이 콘텐츠 복제를 유도한다는 것을 밝힘으로써, 이 연구는 어텐션 흐름을 형성하기 위한 직접적이고 해석 가능한 핸들을 제공합니다. 이는 생성 모델을 내부 표현에 대한 표적 개입을 통해 제어할 수 있다는 더 넓은 방향성을 제시하며, 아키텍처 변경이나 재학습 없이도 정밀한 스타일 정렬 생성을 가능하게 합니다.