Diffusion Transformers에서의 Cross-Layer 정보 라우팅 재고

왜 중요한가

Diffusion Transformers의 잔차 흐름은 여전히 고정된 합산으로 구성되어 깊은 층에서의 정보 전달이 비효율적이다. DAR은 timestep에 따라 가중치를 학습하는 라우팅을 도입해 PreNorm dilution과 gradient 흐름의 비대칭, 블록 간 중복 문제를 완화한다. 이로써 학습 속도와 최종 품질이 함께 개선될 수 있다.

핵심 기여

Diffusion-Adaptive Routing (DAR) 도입

DiT의 각 서브레이어에서 과거 출력들을 softmax 가중합으로 선택적으로 통합하는 학습 가능한 라우팅을 도입한다.

timestep-aware 쿼리 파라미터화

정적(static), 동적(dynamic), 명시적 인젝션(explicit) 등의 방식으로 timestep t를 per-layer query에 주입해 라우팅이 timestep에 적응하도록 한다.

Chunked aggregation

L개의 서브레이어를 N개의 chunk로 나누고 각 chunk의 요약을 사용해 softmax를 수행하되, 소스 수를 |S|+N으로 제한해 계산/메모리 비용을 줄인다.

REPA와의 상호 보완성

DAR은 REPRESENTATION-ALIGNMENT Objective와 독립적으로 작동하며 REPA와 함께 사용 시 초기 학습 속도 및 품질 개선 효과가 가속된다.

실험적 성과

ImageNet-256×256에서 DAR은 SiT-XL/2 대비 2.11의 FID 개선(7.56 vs 9.67) 및 8.75× 적은 학습 이터레이션으로 수렴하는 품질을 달성한다. REPA와 결합 시 초기 단계에서 2× 학습 가속을 달성한다.

핵심 아이디어 이해하기

단락 1: DiT의 Residual Routing은 깊이에 따라 잔차에 의한 정보 축적이 비선형적으로 증가하는 문제를 야기한다. 이로 인해 PreNorm dilution이 나타나고, 깊은 층으로 갈수록 gradient가 약화되며, 이웃한 블록은 높은 유사성을 보여 중복이 심화된다. 단락 2: 본 논문은 Cross-Layer 정보를 timesteps(t)와 depth에 걸쳐 다루는 진단적 분석을 수행하고, timesteps에 따라 가중치를 조정하는 필요한 라우팅이 존재함을 발견한다. 단락 3: DAR은 softmax로 이전 서브레이어 출력들을 가중 합산하는 방식으로 라우팅을 학습 가능하게 만들고, ql(t)를 Static/Dynamic/Explicit Injection으로 설계해 timestep에 따른 적응을 확보한다. 단락 4: Chunked aggregation은 깊이가 깊어질 때 소스 수를 줄이면서도 정보의 손실을 최소화하도록 설계되었고, 최적 chunk_size(S)로는 S=4가 일반적으로 최적임을 이론적으로 예측하고 실험으로 확인한다. 단락 5: DAR은 REPA 등 다른 개선점과 orthogonal하게 작동하며, 전반적 학습 속도와 품질 개선에 기여한다.

방법론

전체 접근: cross-layer routing을 DiT의 핵심 디자인 축으로 간주하고, residual routing을 DAR로 대체한다. hl+1 = fl(hl; t) 대신 hl = sum_{i<l} α_{i→l}(t) vi 형태의 가중 합으로 구성한다. α_{i→l}(t) = exp(ql(t)⊤ ki / sqrt(d)) / sum_{j<l} exp(ql(t)⊤ kj / sqrt(d)) 이다. ki = RMSNorm(vi)이며, Sl = {v0,...,v(l-1)}에서 softmax를 수행한다. 쿼리 파라미터화: ql(t) = static: wl, dynamic: W(l)_q vl-1, explicit: wl + e(t). 타임스텝 주입 방식에는 3가지 변형이 있으며, 실험에서 timestep awareness가 성능의 핵심 요인임을 확인한다. Chunked aggregation: L 서브레이어를 S 크기의 chunk로 묶고, 각 chunk는 cn := v(nS)S의 요약으로 표현되며, 소스 집합 Sl은 |Sl| ≤ S+N로 축약된다. 최종 Aggregator: 마지막 Chunk의 raw 서브레이어 출력까지 접근 가능하도록 설계하여 상세 정보를 보존한다( REPA와의 호환성 유지 시 예외적으로 마지막 chunk의 aggregator는 독립적으로 작동한다). 구현: 고성능 커널(Triton)을 사용해 forward/backward 연산에서 소스들을 스트리밍하며 필요 중간 값을 register에 유지해 속도/메모리 효율을 크게 향상시키고, 28층 기준 N=57까지 확장 가능한 구조로 구현한다.

주요 결과

주요 벤치마크: ImageNet-256×256에서 DAR은 SiT-XL/2 대비 2.11의 FID 개선(7.56 vs 9.67) 및 8.75× 더 적은 학습 이터레이션으로 수렴하는 성능을 보인다. DAR Static(c4) 버전은 600K 이터레이션에서 FID 6.92(SDE)로 도달하며, CFG 없이도 Baseline 대비 개선을 보장한다. DAR Dynamic(c4) 버전은 CFG와 함께 ODE FID 2.05를 달성하였고 SiT baseline보다 우수한 성능을 보인다. REPA를 함께 사용할 때도 100K 이터레이션에서 FID가 7.09로 개선되며, REPA 단독 대비 초기 속도가 2× 증가한다. 타임스텝 인젝션에 따른 ablation에서 timestep-aware 구성(Dynamic/Explicit)이 Static보다 유의하게 우수하며, Explicit Injection과 Static은 초기 학습에서 차이가 작지 않음을 보였다. Chunk 사이즈 실험에서 S=4가 최적의 균형을 보여 U-형 비용 구조를 갖는다고 이론적으로 예측되며, 실제로도 Tab.4에서 S=4가 최적값으로 나타났다. 대규모 T2I 프리트레이닝 및 Distribution Matching Distillation(DMD)에서 DAR은 고주파 디테일 보존에 도움을 준다. DAR은 REPA와 조합 시 초기 속도와 최종 품질을 모두 향상시키며, 두 가지 접근 방식은 상호 보완적으로 작동함이 확인된다.

기술 상세

4.2절에서 제시된 DAR의 수학적 구성은 hl = sum_{i<l} α_{i→l}(t) vi 형태의 누적 집계이며, α_{i→l}(t) = exp(ql(t)⊤ ki / sqrt(d)) / ∑ exp(ql(t)⊤ kj / sqrt(d)) 로 정의된다. ki = RMSNorm(vi). ql(t)는 static/dynamic/explicit injection 중 하나로 설정될 수 있으며 timestep e(t) 삽입도 가능하다. Chunked aggregation은 L개의 서브레이어를 N개의 chunk로 나누고, 각 chunk의 요약 cn을 사용해 Sl를 구성하고 softmax를 수행한다. 최종 Aggregator는 마지막 chunk의 요약뿐 아니라 해당 Chunk의 RAW한 서브레이어 출력을 일부 포함해 더 상세한 정보를 보존한다. DAR-REPA 조합의 경우, REPA의 invariant를 유지하면서 DAR의 Aggregation 파라미터를 그대로 활용하는 방식으로 구현된다. 이때 구현은 Triton 기반의 fused 커널을 사용하여 forward/backward에서 소스 수 N이 커져도 효율성을 유지하도록 설계된다.

실무 활용

DAR은 Diffusion Transformers의 잔차 라우팅을 개선하여 학습 속도와 품질을 동시에 향상시키므로, 대규모 T2I 모델의 전처리/사전학습 및 사후 학습(stepwise distillation)에서 실무적으로 적용 가능하다. REPA와의 조합으로 초기 단계의 학습 가속도도 함께 얻는다.

대형 Diffusion 모델의 프리트레이닝에서 학습 속도 가속 및 품질 향상
Distribution Matching Distillation 같은 고주파 디테일 보존이 중요한 디노이징 기반 distillation에서의 안정성 향상
대규모 T2I 모델의 파인튜닝 및 도메인 특화 모델의 양자화/경량화 과정에서의 성능 유지
REPA와 결합한 포지션에서의 초기 학습 가속 및 안정성 증가

코드 공개 여부: 미확인

키워드

Diffusion Transformers (확산 트랜스포머)cross-layer routing (크로스-레이어 라우팅)Diffusion-Adaptive Routing (DAR)REPAFIDDMD (Distribution Matching Distillation)

용어 해설

크로스-레이어 라우팅(Cross-Layer Routing): — 딥 뉴럴 네트워크에서 서로 다른 층들 간에 정보를 주고받는 방식에 관한 일반적 용어로, DiTs의 잔차 스트림에서 층 간 정보의 흐름을 동적으로 조정하는 문제를 다룬다.
PreNorm dilution: — Residual stream의 Magnitude 증가로 인해 노멀라이제이션(PreNorm) 이후 각 층의 출력이 점점 확산되듯 커지는 현상으로, 깊은 DiT에서 정보 손실과 gradient 흐름 불균형을 초래한다.
블록 단위 중복(Block-wise redundancy): — 연속된 Transformer 블록 간 표현이 유사하게 반복되어 계산 효율이 떨어지는 현상으로, 인접 블록의 출력이 서로 크게 달라지지 않는다.
Diffusion-Adaptive Routing (DAR): — DiT의 잔차 합산을 학습 가능한 softmax 가중합으로 대체하고, timesteps(t)에 따라 동적으로 라우팅을 조정하는 크로스-레이어 라우팅 방법.
REPA: — Representation Alignment for diffusion transformers. 중간 은닉 상태를 pretrained 비주얼 인코더의 표현과 정렬하는 목표로 DiT의 학습 효율을 높이는 기법.

Diffusion Transformers에서의 Cross-Layer 정보 라우팅 재고

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

관련 토론

댓글

관련 피드