[논문리뷰] WiT: Waypoint Diffusion Transformers via Trajectory Conflict Navigation

링크: 논문 PDF로 바로 열기

저자: Hainuo Wang, Mingjia Li, et al.

1. Key Terms & Definitions

• Flow Matching (FM) : 연속적인 확산 프로세스를 최적 운송 프레임워크로 통합하여 단순한 베이스 분포에서 타겟 분포로의 매핑을 통해 곧은 확률 흐름 ODE(Ordinary Differential Equation) 궤적을 생성하는 모델링 방법론입니다.
• Trajectory Conflict : 픽셀 공간 생성 모델에서 시각적으로 유사하지만 의미론적으로 구별되는 최종 지점들의 최적 운송 경로가 노이즈 공간 내에서 중첩되어, 네트워크가 평균화된 속도 필드를 예측하게 되어 의미론적 블리딩(semantic bleeding)과 느린 수렴을 야기하는 현상입니다.
• Waypoint Diffusion Transformers (WiT) : 픽셀 공간 Flow Matching의 심각한 Trajectory Conflict 문제를 해결하기 위해 식별 가능한 중간 Semantic Waypoint를 도입한 새로운 생성 패러다임입니다.
• Just-Pixel AdaLN : WiT 아키텍처 내에서 동적으로 예측된 Semantic Waypoint를 활용하여 Transformer Feature에 공간적으로 변화하는 변조(modulation)를 제공함으로써, 의미론적 내비게이션(semantic navigation)과 고해상도 픽셀 생성 간의 결합을 명시적으로 해제하는 메커니즘입니다.
• x-prediction : 확산 모델에서 네트워크가 노이즈 ε 또는 속도 v 대신 노이즈 상태 zt로부터 깨끗한 이미지 x를 직접 예측하도록 학습하는 전략입니다.

2. Motivation & Problem Statement

최근 Flow Matching 모델은 Latent Autoencoder의 재구성 병목 현상을 피하기 위해 픽셀 공간에서 직접 작동합니다. 그러나 픽셀 매니폴드(manifold)의 의미론적 연속성이 부족하여 최적 운송 경로가 심하게 얽히게 됩니다. 이는 교차점 근처에서 심각한 Trajectory Conflict 를 유발하여 차선책 솔루션으로 이어집니다. 기존 연구들은 정보 손실이 있는 Latent Representation을 통해 이 문제를 우회하려 했으나, 이는 고주파 텍스처 디테일을 손상시키고 시각적 아티팩트를 생성하는 정보 병목 현상을 초래합니다. 픽셀 도메인은 학습 가능한 Latent Space와 달리 인위적으로 재구성하여 의미를 분리할 수 없으므로, 표준 픽셀 공간 Flow Matching 은 이 문제에 직접적으로 직면하게 됩니다. 결과적으로, 네트워크는 중첩된 경로에 대한 회귀 손실을 최소화하기 위해 평균화된 속도 필드를 예측하게 되어 Semantic Bleeding과 느린 수렴이 발생합니다.

3. Method & Key Results

저자들은 픽셀 공간 Trajectory Conflict 문제를 해결하기 위해 Waypoint Diffusion Transformers (WiT) 를 제안합니다. WiT 는 사전 학습된 비전 모델에서 투영된 중간 Semantic Waypoint를 통해 연속적인 Vector Field를 인수분해합니다. 이를 통해 최적 운송 경로를 "노이즈에서 Waypoint로"와 "Waypoint에서 픽셀로" 두 부분으로 분할하여 생성 궤적을 효과적으로 분리합니다.

방법론의 핵심은 다음과 같습니다. 먼저, 타겟 이미지 x로부터 DINOv3 Feature φ(x) 를 추출한 후 PCA 를 적용하여 저차원(예: d=64 ) Semantic Waypoint s0 를 구성합니다. 이 Waypoint는 클래스 분리 가능성을 최적화하여 구조적인 감독 신호 역할을 합니다. 다음으로, 경량의 Waypoints Generator (Wψ) ( 21M 파라미터)가 현재 노이즈 상태 zt 로부터 깨끗한 Semantic Waypoint ŝ0 를 동적으로 추론합니다. Wψ 는 Semantic Space에서 병렬 확률 흐름 ODE를 매칭하는 Lsem Loss를 통해 학습됩니다. 마지막으로, 주된 Pixel Space Generator (Gθ) 는 예측된 Semantic Waypoint ŝ0 를 공간 조건으로 활용하는 제안된 Just-Pixel AdaLN 메커니즘을 통해 픽셀 생성을 수행합니다. Just-Pixel AdaLN 은 글로벌 컨디셔닝과 로컬 Semantic Map을 결합하여 Transformer Feature에 공간적으로 변화하는 변조를 제공하며, 이는 Semantic 내비게이션과 고해상도 픽셀 생성 간의 명시적인 분리를 보장합니다. WiT 는 Wψ 를 먼저 학습시킨 후 고정하고 Gθ 를 학습시키는 두 단계 훈련 패러다임을 따릅니다 [Algorithm 1, Figure 2, Figure 3(a)].

주요 실험 결과는 다음과 같습니다:

• ImageNet 256x256 벤치마크에서 WiT-B/16 은 200 Epoch만으로 FID 3.34 를 달성하여, 600 Epoch 학습된 바닐라 JiT-B/16 의 FID 3.66 을 넘어섭니다. 이는 JiT [22] 대비 2.2배 의 훈련 가속을 보여줍니다 [Figure 1(e), Table 2].
• WiT-L/16 은 265 Epoch 훈련 시 FID 2.36 및 IS 293.7 을 달성하여, 600 Epoch 훈련된 JiT-L/16 의 성능과 일치하며 2.27배 의 훈련 속도 향상을 입증했습니다. 600 Epoch까지 훈련된 WiT-L/16 은 FID 2.22 및 IS 303.3 를 기록하여, Latent Space 벤치마크인 DiT-XL/2 의 FID 2.27 까지 능가합니다

Table 2: Comprehensive comparison of class-conditional ImageNet 256 × 256.

.

• Just-Pixel AdaLN 메커니즘은 Semantic Waypoint 주입 방법 중 가장 우수한 성능(FID 3.34 , IS 270.73 )을 보였으며, 이는 Channel Concat (FID 3.93 ) 및 In-context Concat (FID 3.63 )보다 현저히 높습니다 [Table 3].
• Trajectory Conflict 에 대한 정량적 분석 결과, WiT 는 통합 중간 지점(t ≈ 0.5)에서 JiT [22] 대비 Pairwise Conflict가 1.55배 더 안정적이었으며, 최대 Conflict Peak 지점에서는 1.62배 더 높은 안정성을 보였습니다 [Table 4].

4. Conclusion & Impact

본 연구는 픽셀 공간 Flow Matching 에 내재된 심각한 Trajectory Conflict 를 해결하기 위한 새로운 생성 패러다임인 Waypoint Diffusion Transformers (WiT) 를 제시합니다. WiT 는 원시 픽셀 매니폴드가 본질적으로 얽혀 있고 직접적인 정규화에 저항한다는 점을 인식하여, 생성 프로세스를 Semantic 내비게이션과 고해상도 텍스처 합성으로 명시적으로 분리합니다. 사전 학습된 비전 모델의 식별 가능한 특징 공간을 Compact한 Semantic Waypoint로 투영함으로써, WiT 는 복잡한 노이즈-투-픽셀 최적 운송 경로를 성공적으로 인수분해합니다.

통합 과정 동안 경량의 Waypoints Generator 는 이러한 구조적 앵커를 동적으로 추론하며, 이는 제안된 Just-Pixel AdaLN 메커니즘을 통해 주된 확산 Transformer에 공간적으로 변화하는 가이던스를 제공합니다. ImageNet 256x256에 대한 광범위한 실험은 WiT 가 순수 픽셀 공간 아키텍처 중 최첨단 성능을 달성하며, 심지어 무거운 VAE 압축 Latent 모델까지 능가한다는 것을 입증했습니다. 또한 JiT [22] 대비 2.2배 의 놀라운 훈련 가속을 제공합니다. 이 연구는 Latent 압축이나 무차별적인 파라미터 스케일링에 의존하지 않고도 고해상도 픽셀 공간 생성 모델의 성능을 크게 향상시키는 중요한 시사점을 제공합니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.