WiT: 궤적 충돌 내비게이션을 통한 경로점 확산 트랜스포머

기존 픽셀 기반 생성 모델은 이미지의 복잡한 구조 때문에 학습 경로가 엉키는 '궤적 충돌' 문제를 겪어 학습이 느리고 화질이 떨어지는 한계가 있었다. WiT는 시각 모델의 의미 정보를 '경로점'으로 활용해 생성 과정을 단계별로 안내함으로써 학습 효율을 2.2배 높이고 고해상도 텍스처를 완벽하게 보존한다.

Diffusion 모델이나 Flow Matching이 픽셀 공간에서 직접 이미지를 생성할 때, 서로 다른 의미를 가진 이미지들이 노이즈 공간에서 비슷한 경로를 공유하게 된다. 이를 '궤적 충돌(Trajectory Conflict)'이라 하며, 모델이 여러 경로의 평균값을 학습하게 만들어 경계가 흐릿해지거나 학습이 매우 느려지는 원인이 된다.

WiT는 이 문제를 해결하기 위해 생성 경로 중간에 '의미론적 경로점(Semantic Waypoint)'을 도입한다. 이는 마치 복잡한 미로에서 목적지까지 한 번에 가는 대신, 중간 기착지를 설정해 경로를 명확히 하는 것과 같다. 사전 학습된 시각 모델(DINOv3)에서 추출한 특징을 PCA로 압축해 핵심 구조 정보만 남긴 경로점을 생성 가이드로 삼아 궤적을 분리한다.

이러한 접근은 모델이 '무엇을 그릴지(의미론적 내비게이션)'와 '어떻게 그릴지(픽셀 텍스처 생성)'를 단계적으로 처리하게 만든다. 결과적으로 모델은 엉킨 경로에서 벗어나 더 직선적이고 명확한 최적 운송 경로를 따라 이미지를 생성할 수 있게 되며, 이는 곧 학습 속도 향상과 고화질 결과물로 이어진다.

Flow Matching 프레임워크를 기반으로 하며, 픽셀 공간의 속도 필드를 직접 예측하는 대신 x-prediction 방식을 채택하여 안정성을 높였다. 궤적 충돌을 수학적으로 정의하기 위해 조건부 분산 Var(x|zt)을 사용하며, 이를 경로점 s0에 대한 조건부 분산으로 분해하여 최적화 부담을 획기적으로 줄였다.

의미 경로점은 DINOv3의 패치별 토큰 phi(x)를 입력으로 받아 데이터셋 평균 mu를 빼고 PCA 투영 행렬 Ud를 곱해 d=64 차원으로 압축하여 생성한다. [N x D 차원의 토큰 입력 → 평균 차감 및 행렬 곱셈 연산 → N x 64 차원의 경로점 출력 → 구조적 핵심 정보만 포함된 저차원 가이드 생성].

Just-Pixel AdaLN은 예측된 경로점 s0_hat을 선형 투영하여 공간 컨디션 cs를 만든 후, 이를 통해 6개의 변조 파라미터를 산출한다. [공간 컨디션 입력 → 선형 레이어 연산 → 픽셀별 스케일 및 시프트 값 출력 → 각 트랜스포머 블록의 정규화 및 잔차 연결 제어]. 이 메커니즘은 전역적인 클래스 정보와 국소적인 의미 구조를 동시에 주입한다.

ImageNet 256x256 벤치마크에서 WiT-L/16 모델은 600 에포크 학습 후 FID 2.22, IS 303.3을 기록했다. 이는 기존 픽셀 기반 SOTA 모델인 JiT-L/16(FID 2.36)을 능가하는 수치이며, 심지어 강력한 잠재 공간 모델인 DiT-XL/2(FID 2.27)와 대등하거나 앞서는 성능이다.

학습 효율성 측면에서 WiT는 JiT 대비 약 2.2배의 수렴 가속도를 보였다. WiT-L/16은 단 265 에포크 만에 JiT-L/16의 600 에포크 성능과 대등한 FID 2.36을 달성하여, 의미론적 경로점이 학습 난이도를 크게 낮춤을 입증했다.

Ablation study 결과, Just-Pixel AdaLN 방식이 단순 채널 결합(FID 3.93)이나 인컨텍스트 결합(FID 3.63)보다 우수한 성능(FID 3.34)을 보였다. PCA 차원은 64일 때 최적의 균형을 유지했으며, 차원이 너무 높으면 차원의 저주로 인해 오히려 성능이 저하되는 현상이 관찰됐다.

WiT 아키텍처는 21M 파라미터의 경량 Waypoints Generator와 주 Pixel Space Generator로 구성된 이단계 디커플링 구조를 가진다. Waypoints Generator는 ViT-S/16 설정을 따르며 노이즈 상태 zt에서 깨끗한 경로점 s0를 추론하도록 학습된다. 이 과정에서 독립적인 가우시안 노이즈를 사용하는 병렬 확률 흐름 ODE를 구축하여 의미 공간에서의 매핑을 감독한다.

이론적으로 궤적 충돌은 베이즈 리스크(Bayes-risk) 관점에서 분석된다. 표준 픽셀 공간 예측의 오차는 조건부 분산 E[Var(x|zt)]로 표현되는데, WiT는 경로점 s0를 도입하여 이를 E[Var(x|zt, s0)] 수준으로 낮춘다. 총 분산의 법칙에 따라 전체 불확실성에서 경로점에 의해 설명되는 부분을 제외함으로써 주 생성 모델의 검색 공간을 효과적으로 축소한다.

Just-Pixel AdaLN은 AdaLN-Zero 포뮬레이션을 확장하여 공간적으로 변화하는 맵을 통해 특징을 변조한다. 이는 전역적인 클래스 정보뿐만 아니라 국소적인 의미 구조를 트랜스포머의 각 레이어에 직접 주입하여 구조적 일관성을 강화한다. 학습 시에는 logit-normal 분포를 사용한 비균일 타임스텝 샘플링을 통해 궤적 충돌이 가장 심한 중간 노이즈 레벨에 학습 용량을 집중시켰다.