Long 스킵 연결로 안정화된 효율적인 Diffusion Transformer

초록

본 논문은 Diffusion Transformer(DiT)의 동적 특징 불안정을 해결하기 위해 U‑Net에서 영감을 얻은 Long‑Skip‑Connection(LSC)을 도입한 Skip‑DiT를 제안한다. LSC는 깊은 블록과 얕은 블록을 직접 연결해 스펙트럴 노름을 제어하고, 특성 전파의 안정성을 확보한다. 이를 통해 학습 속도는 최대 4.4배 가속화되고, 추론 시 1.5‑2배의 속도 향상을 달성하면서도 원본 DiT와 거의 동일한 이미지·비디오 품질을 유지한다.

상세 분석

Skip‑DiT의 핵심 아이디어는 DiT 내부에 Long‑Skip‑Connection(LSC)을 삽입해 깊은 레이어에서 얕은 레이어로 정보를 직접 전달하는 것이다. 기존 DiT는 순차적인 트랜스포머 블록만으로 구성돼 각 블록의 Jacobian 스펙트럴 노름이 제어되지 않아 작은 입력 교란이 층을 거치면서 기하급수적으로 증폭된다. 논문은 이를 “Dynamic Feature Instability(DFI)”라 정의하고, 실험적으로 파라미터와 입력에 대한 작은 섭동(α·δ, β·η)을 가했을 때 코사인 유사도가 급격히 감소함을 보여준다.

스펙트럴 분석에서는 각 트랜스포머 블록 Tₗ이 노이즈 감소 불변성을 만족한다는 가정 하에 σ_max(∂Tₗ/∂h) = γₗ < 1임을 증명한다. LSC를 도입하면 깊은 블록의 출력이 얕은 블록에 직접 더해지면서 전체 경로의 스펙트럴 노름이 γⁿ이 아닌 γ·(1+γ) 형태로 낮아진다. 즉, 특성 전파 과정에서 증폭 효과가 억제되고, 역전파 시 그래디언트 소실 위험도 감소한다.

또한 LSC는 “statistical caching”을 가능하게 한다. 깊은 블록의 출력을 일정 타임스텝에서 캐시하고, 추론 단계에서는 얕은 블록만 재계산하면 된다. 이는 기존 DiT 캐싱 기법(FORA, Faster‑Diff 등)이 얕은 블록까지 모두 재계산해야 하는 문제를 해결한다. 실험 결과, 캐시 간격을 크게 늘려도 Skip‑DiT는 코사인 유사도가 0.98 이상 유지돼 품질 저하가 거의 없으며, 영상 생성 PSNR에서도 동일하거나 더 높은 값을 기록한다.

학습 효율성 측면에서는 ImageNet‑50K FID를 기준으로, Skip‑DiT가 1.6M 스텝에서 FID 8.37을 달성하는 반면, 기존 DiT‑XL은 7M 스텝이 필요했다. 이는 스펙트럴 노름 제어가 학습 수렴을 가속화하고, 초기 단계에서의 불안정성을 감소시켜 최적화 경로를 더 부드럽게 만든 결과이다.

요약하면, LSC는 DiT의 구조적 결함을 보완해 동적 특성 불안정을 근본적으로 해결하고, 스펙트럴 제약을 통해 학습·추론 모두에서 효율성을 크게 향상시킨다.