Rethinking Cross-Layer Information Routing in Diffusion Transformers

기술적 요약: 확산 트랜스포머의 교차 계층 정보 라우팅 재고찰

문제 제기
확산 트랜스포머 (DiTs) 는 토큰화, 어텐션 메커니즘, 조건부 처리, 그리고 목적 함수에 대한 광범위한 연구를 통해 현대 시각 생성의 핵심 기둥으로 자리 잡았습니다. 그러나 잔여 스트림(residual stream)—즉, 정보가 계층을 가로질러 누적되는 방식을 규율하는 메커니즘—은 원래의 트랜스포머 설계에서 크게 변하지 않은 채 남아 있습니다. 저자들은 자연어 처리 (NLP) 관행에서 계승된 이 표준 사전 정규화 잔여 덧셈이 확산 과정의 시간 가변적 역동성에 적합하지 않다고 지적합니다.

깊이(depth)와 잡음 제거 시간 단계(denoising timestep)를 따라 교차 계층 정보 흐름에 대한 체계적인 실증 분석을 통해, 이 논문은 DiT 의 전통적 잔여 라우팅에서 관찰되는 세 가지 구체적인 증상을 진단합니다:

1. 단조로운 순방향 크기 팽창: 은닉 상태의 크기가 첫 번째 블록에서 마지막 블록으로 갈수록 약 100 배 증가하여, 더 깊은 블록이 영향을 유지하기 위해 점점 더 큰 원시 출력 (raw outputs) 을 생성하도록 강요합니다. 이는 대형 언어 모델 (LLM) 에서 관찰되는 'PreNorm 희석' 현상과 유사합니다.
2. 급격한 역방향 기울기 감쇠: 기울기는 처음 몇 개의 블록 이후에 현저히 감소하여, 더 깊은 계층에 실질적으로 약화된 최적화 신호만 남깁니다.
3. 뚜렷한 블록별 중복성: 인접한 트랜스포머 블록이 매우 유사한 표현 (코사인 유사도 > 0.9) 을 생성하여, 표준 라우팅이 깊이 전반에 걸쳐 특징을 다양화하지 못함을 나타냅니다.

또한, 분석은 주어진 계층에 대한 최적의 정보원이 잡음 제거 시간 단계에 따라 체계적으로 이동함을 보여줍니다 (높은 잡음에서는 거시적 구조 특징에서 낮은 잡음에서는 미세한 세부 특징으로 이동). 균일하고 시간 무관한 가중치를 적용하는 표준 잔여 연결은 이러한 변화하는 요구 사항에 적응하지 못합니다.

**방법론: 확산 적응형 라우팅 **(Diffusion-Adaptive Routing, DAR)
이러한 문제들을 해결하기 위해 저자들은 표준 잔여 연결을 대체할 수 있는 **확산 적응형 라우팅 **(DAR)을 제안합니다. 고정된 가산적 재귀와 달리, DAR 는 하위 계층 출력의 역사에 대해 학습 가능하고, 시간 단계 적응적이며, 비증분적인 집계를 수행합니다.

주요 설계 구성 요소는 다음과 같습니다:

• 깊이에 대한 Softmax 어텐션: 이전 출력들을 단위 가중치로 합산하는 대신, DAR 는 이전 하위 계층 출력 ($v_i$) 에 대한 softmax-가중 집계를 계산합니다. 집계 가중치는 현재 은닉 상태에서 유도된 쿼리 ($q_l$) 와 정규화된 이전 출력에서 유도된 키 ($k_i$) 에 의해 결정됩니다.
• 시간 단계 인식: 잡음 제거 시간 단계 $t$가 중요한 제어 신호임을 인식하여, 쿼리 파라미터화는 $t$를 인지하도록 설계됩니다. 저자들은 세 가지 변형을 탐구합니다:
  • 정적 (Static): 계층당 학습 가능한 벡터 (시간 무관).
  • 동적 (Dynamic): 쿼리는 이전 하위 계층 출력의 투영이며, 은닉 상태에 인코딩된 시간 단계 정보를 암묵적으로 계승합니다.
  • 명시적 주입 (Explicit Injection): 시간 단계 임베딩 $e(t)$로 정적 쿼리를 보강합니다.
    실험 결과, 시간 단계 인식이 중요하며, 동적 및 명시적 변형이 시간 무감각한 정적 기준선보다 우수한 성능을 보임이 확인되었습니다.
• **조각화 집계 **(Chunked Aggregation) 모든 $L$개의 이전 계층을 어텐션하는 계산 비용을 관리하기 위해, DAR 는 조각화된 변형을 지원합니다. 하위 계층은 조각으로 분할되며, 라우터는 이전 조각들의 요약과 현재 조각 내의 원시 출력에 어텐션합니다. 이론적 분석은 라우팅 정밀도와 압축 비용 사이의 균형을 맞추는 최적의 조각 크기가 존재함을 시사하며 (SiT-XL/2 에 대해 실험적으로 $S=4$로 발견됨), 이는 최적의 조각 크기입니다.
• 호환성: DAR 는 깊이 차원만을 따라 작동하여 등방성 및 동질적인 트랜스포머 스택을 보존합니다. 수동 계층 페어링이나 학습 목적 함수 변경 없이 REPA(표현 정렬) 와 같은 다른 현대적 개선 사항과 호환되도록 설계되었습니다.

주요 결과
SiT-XL/2 아키텍처를 사용하여 ImageNet 256 × 256에서 평가한 결과:

• 품질 향상: DAR 의 정적 변형은 (분류기 없는 가이드 없이) 7.56 FID를 달성하여, 기준선 SiT-XL/2(9.67 FID) 대비 2.11 포인트 개선을 보였습니다.
• 수렴 속도: DAR 는 약 8.75 배 적은 학습 반복 횟수로 기준선의 수렴 품질을 달성합니다.
• REPA 와의 시너지: REPA 위에 적층될 때, DAR 는 초기 단계에서 2 배의 학습 가속화를 제공하여, 교차 계층 라우팅 개선이 표현 정렬 목적과 직교함을 보여줍니다.
• U-Net 스킵과의 비교: DAR 는 손으로 설계된 U-Net 유사 스킵 연결을 사용하는 모델 (예: U-DiT-L) 보다 우수한 성능을 내면서도 동질적인 트랜스포머 스택의 확장성 이점을 유지합니다.
• **증류 **(Distillation) 대규모 T2I 사후 학습 (Qwen-Image 에 대한 분포 매칭 증류) 에서 DAR 는 공격적인 소수 단계 증류에서 종종 손실되는 고주파 세부 사항 (날카로운 가장자리, 미세한 질감) 을 보존하는 데 도움을 줍니다.

의의 및 주장
이 논문은 교차 계층 정보 라우팅이 기존 표현 정렬 목적 및 아키텍처 확장과 직교하게 작동하는 확산 모델링에서 미탐구된 설계 축이라고 주장합니다. DiT 의 표준 잔여 스트림에서 발생하는 특정 시공간 병리 현상을 진단함으로써, 저자들은 고정된 덧셈을 적응적이고 시간 단계 조건부인 집계로 대체하는 것이 수렴 속도와 최종 생성 품질 모두를 크게 향상시킨다는 것을 입증합니다.

이 연구는 교차 계층 라우팅을 계승된 관례에서 명시적인 설계 차원으로 격상시킵니다. 표준 잔여의 경직성이 확산 모델의 제한 요인임을 시사하며, DAR 와 같은 적응형 라우팅 메커니즘이 깊은 DiT 아키텍처의 잠재력을 완전히 unlocking 하는 데 필수적임을, 특히 더 깊은 스택과 더 복잡한 사후 학습 정세로 확장됨에 따라 강조합니다.