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

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

단락 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이 커져도 효율성을 유지하도록 설계된다.