기술 요약: WaiT (Wait for the Signal)

문제 정의

이미지 생성 모델이 고해상도(512×512, 1024×1024 및 그 이상)로 확장됨에 따라, 전역적 일관성(global coherence), 국소적 디테일(local detail), 그리고 질감 충실도(texture fidelity)를 동시에 달성하는 것이 점점 더 어려워지고 있습니다. 표준 플로우 매칭(flow matching) 모델은 모든 공간 주파수를 균등하게 취급하며, 이미지 신호의 자연스러운 계층 구조를 무시합니다. 경험적 분석에 따르면, 고주파(HF) 대역은 저주파(LF)의 거친 구조보다 훨씬 더 이른 디노이징 궤적 단계에서 순수 노이즈와 구별할 수 없는 상태가 됩니다. 결과적으로, 기존 모델들은 생성 초기 단계에서 학습 가능한 신호가 없는 고주파 대역을 디노이징하는 데 상당한 모델링 용량을 낭비하며, 이 시기에는 여전히 전역적 구조가 구축되는 중입니다. 또한, FID와 같은 표준 평가 지표는 공격적인 다운샘플링(299×299로)에 의존하여, 고충실도 생성을 정의하는 바로 그 고주파 질감과 국소적 디테일을 버리게 되며, 이는 네이티브 해상도에서의 품질을 불충분하게 평가하는 결과를 초래합니다.

방법론: WaiT (Wavelet-aware image Transformer)

저자들은 손실 없는 이산 웨이브릿 변환(Discrete Wavelet Transform, DWT)을 사용하여 이미지의 자연스러운 주파수 구조를 생성 과정에 직접 내장하는 픽셀 공간 접근 방식인 WaiT를 제안합니다. 핵심 혁신은 고주파 대역이 저주파 신호가 나타날 때까지 지연되는 "신호를 기다리는(wait for the signal)" 전략입니다.

1. 주파수별 노이즈 스케줄 (Frequency-Specific Noise Schedules)

균등한 노이즈 스케줄 대신, WaiT는 단일 전환 임계값 $t^*$를 사용하여 생성 과정을 두 단계로 분리합니다:

• Phase 시점 0 (거친 단계, $0 \le t < t^*$): 모델은 저주파(LF) 대역(근사 계수)에 대해서만 작동합니다. 고주파(HF) 대역은 순수한 단위 분산 노이즈($\epsilon_{HF}$) 상태로 유지됩니다. 이를 통해 모델은 고주파 노이즈의 방해 없이 전역적 구조적 일관성을 구축하는 데 역량을 집중할 수 있습니다.
• Phase 1 (정밀 단계, $t^* \le t \le 1$): 교차 지점 $t^*$에서, 부분적으로 디노이징된 LF 대역은 (학습 분포와 일치하도록 주입된) 새로운 HF 노이즈와 결합됩니다. 그 후 모델은 픽셀 공간에서 두 대역을 공동으로 디노이징합니다.

노이즈 스케줄은 다음과 같이 정의됩니다:

• $t_{LF} = t$ ( $[0, 1]$ 구간에 대한 표준 선형 스케줄).
• $t_{HF} = \max(0, \frac{t - t^*}{1 - t^*})$ (지연된 스케줄, $t^*$부터 시작).

이 공식화는 전환 지점 $t^*$에서 완벽한 분포 일치를 보장하여, 계층형(cascaded) 또는 불연속적 접근 방식에서 흔히 발생하는 학습-테스트 간의 불일치 문제를 해결합니다.

2. 아키텍처 구현

WaiT는 최소한의 수정만 거친 JiT (Just Image Transformer) 아키텍처 상에 구현되었습니다:

• 손실 없는 분해 (Lossless Decomposition): 단일 레벨의 Haar DWT를 사용하여 이미지를 하나의 LF 대역($D/4$ 토큰)과 세 개의 HF 대역($3D/4$ 토큰)으로 분리합니다.
• 대역 정규화 (Band Normalization): LF 대역은 노이즈 스케줄의 동적 범위를 맞추기 위해 스칼라 $S_{LF}$(절대 계수의 95 백분위수)에 의해 정규화되는 반면, HF 대역은 자연스러운 희소성을 보존하기 위해 정규화되지 않은 상태로 둡니다.
• 해상도 컨디셔닝 (Resolution Conditioning): 모델이 거친(저해상도) 단계와 정밀한(전체 해상도) 단계를 구분할 수 있도록 시간 및 클래스 컨디셔닝에 해상도 임베딩을 추가합니다.
• 학습 목적 함수 (Training Objective): 모델은 각 대역의 노이즈 레벨에 따른 대역별 재가중치를 적용하여 $v$-loss를 사용하는 $x$-예측(x-prediction)을 사용합니다.

3. 평가 프로토콜

표준 FID가 고해상도 디테일을 포착하지 못한다는 점을 인식하여, 저자들은 3축 평가 프로토콜을 도입합니다:

1. FID: 전역적 일관성을 측정합니다 (표준 다운샘플링 사용).
2. 5-crop FID (5cFID): 다운샘플링 없이 다섯 개의 네이티브 해상도 크롭(모서리 4개 + 중앙 1개)에 대해 FID를 계산하여 국소적 구조 디테일을 측정합니다.
3. 고주파 FWD (hFWD): 저주파 DC 성분을 명시적으로 제외하고 프레셰 웨이브릿 거리(Fréchet Wavelet Distance)를 계산하여 순수 질감 충실도를 측정합니다.

주요 결과

ImageNet 512×512

• 성능: WaiT는 픽셀 공간 FID 1.43(WaiT-H/16) 및 1.30(WaiT-G/16)을 달성하여, 픽셀 공간 모델의 새로운 SOTA(State-of-the-Art)를 기록했습니다.
• 효율성: JiT 베이스라인과 비교하여 샘플링 연산량(GFLOPs)을 최대 50% 줄이면서도 더 나은 혹은 대등한 품질을 달 achieves 합니다.
• 질감 충실도: WaiT는 hFWD(WaiT-H/16 기준 0.67)에서 최첨단 픽셀 공간 및 잠재 공간 모델을 모두 능가하며, 우수한 질감 생성을 입증했습니다.
• 파레토 최적성 (Pareto Optimality): WaiT는 세 가지 지표(FID, 5cFID, hFWD)와 연산 비용 모두에서 파레토 프런트를 지배합니다.

OpenImages (512×512 및 1024×1024)

• WaiT는 JiT 베이스라인을 일관되게 능가하며, 해상도가 높아질수록(1024×1024) 성능 격차가 더 커집니다.
• 1024 해상도에서 WaiT는 거의 모든 지표에서 JiT와 대등하거나 이를 능가하면서도 최대 **3배 높은 처리량(throughput)**을 달성합니다.

비디오 생성

• 이 프레임워크는 시공간 웨이브릿(높이, 너비, 시간에 대해 작동)을 사용하여 비디오로 매끄럽게 확장됩니다.
• Kinetics-600에서 WaiT는 JiT 베이스라인 대비 연산량을 30% 감소시키면서 SOTA FVD 0.84(WaiT-XL/8)를 달성했습니다.
• Taichi-HD에서 WaiT는 FVD를 개선하고 연산량을 약 22% 줄였습니다.

텍스트-투-이미지 (Text-to-Image)

• 1024×1024 텍스트-투-이미지 합성에 적용되었을 때, WaiT는 더 높은 처리량을 유지하면서 거의 모든 지표(FID, 5cFID, hFWD, CLIPScore, GenEval)에서 JiT 베이스라인을 개선했습니다.

의의 및 주장

본 논문은 WaiT가 이미지의 자연스러운 주파수 계층 구조와 생성 과정을 정렬시키는 **원칙적인 유도 편향(principled inductive bias)**을 나타낸다고 주장합니다. "신호를 기다림으로써"(LF 구조가 확립될 때까지 HF 생성을 지연함으로써), 이 방법은:

1. 품질 향상: 전역적 일관성을 희생하지 않으면서 질감 충실도와 국소적 디테일을 크게 향상시킵니다.
2. 비용 절감: 초기 거친 생성 단계 동안 토큰 수를 줄임으로써 상당한 계산 절감(샘플링 시 최대 50%)을 실현합니다.
3. 일반화 가능성: 이 접근 방식은 모달리티에 구애받지 않으며, 이미지와 비디오 모두에 매끄럽게 작동하고, 픽셀 공간 및 잠재 공간 아키텍처(DDT를 통해 입증됨) 모두와 호환됩니다.
4. 한계 노출: 제안된 3축 평가 프로토콜은 표준 FID의 약점을 드러내며, 현재의 잠재 공간 모델들이 올바르게 평가될 때 픽셀 공간 접근 방식에 비해 고주파 질감 충실도 측면에서 어려움을 겪는 경우가 많음을 강조합니다.

저자들은 이러한 개선이 복잡한 아키텍처 피라미드나 학습된 오토인코더가 아니라, 노이즈 스케줄에 대한 간단한 수정과 손실 없는 파라미터 프리(parameter-free) 웨이브릿 변환을 통해 달성되었다는 점을 강조합니다.