한 줄 요약
SiT는 DiT 기반으로 설계된 생성 모델로, 다양한 디퓨전 계수와 연속 시간 학습을 통해 ImageNet에서 FID-50K 2.06을 달성한다.
핵심 기여도
- SiT는 DiT와 동일한 모델 구조, 파라미터 수, GFLOPs를 사용하면서 ImageNet 256x256에서 FID-50K 2.06을 달성함.
- 연속 시간 학습, 벨로시티 예측, 선형 인터폴랜트를 도입하여 디퓨전 모델의 성능을 일관되게 개선함.
- 디퓨전 계수는 학습과 별도로 조정 가능하며, 이는 샘플링 과정에서 유연성을 제공함.
- SDE 샘플러에서 벨로시티 모델만으로도 성능을 향상시킬 수 있음을 보여줌.
핵심 아이디어
SiT는 기존 디퓨전 모델에서 제한적인 노이징-데이터 연결 방식을 벗어나, 더 유연한 인터폴랜트 프레임워크를 도입함. 이는 두 분포를 연결하는 방식을 모듈식으로 연구할 수 있게 하며, 시간의 연속/이산성, 예측 대상(스코어 vs. 벨로시티), 인터폴랜트 유형, 샘플링 방식(결정적 vs. 확률적)을 독립적으로 탐색할 수 있도록 함. 특히, 벨로시티 모델 $ \mathbf{v}_\theta(\mathbf{x}, t) $를 학습하고 이를 기반으로 스코어 $ \mathbf{s}(\mathbf{x}, t) $를 표현함으로써, 샘플링 과정에서 더 정확한 경로를 추적할 수 있음. 이는 기존 디퓨전 모델에서의 노이즈 예측 방식보다 더 직관적이고 수학적으로 안정적인 접근법임.
기술적 접근법
- **모델 구조**: DiT-B 기반으로 구축, 파라미터 수와 GFLOPs는 DiT와 동일.
- **학습 방식**: 연속 시간 학습을 도입, 벨로시티 $ \mathbf{v}(\mathbf{x}, t) $를 예측.
- **인터폴랜트**: 선형 인터폴랜트를 사용, VP( Variance Preserving) 인터폴랜트 대비 더 유연한 연결.
- **샘플링**: ODE는 Heun 적분기, SDE는 Euler-Maruyama 적분기 사용, 샘플링 단계 수는 250 NFE로 제한.
- **디퓨전 계수**: $ w_t = \sigma_t $로 설정하여 수치적 특이점을 제거.
주요 결과
- ImageNet 256x256 데이터셋에서 SiT는 FID-50K 2.06을 달성함 (DiT 대비 개선).
- 512x512 이미지에서도 FID-50K 2.62를 달성함.
- 벨로시티 모델 학습 시, 가중치 스코어 모델 대비 FID-50K에서 0.5 이상 개선됨 (Table 5 참조).
의의 및 한계
SiT는 디퓨전 모델의 핵심 구성 요소를 모듈식으로 분리하여, 각 요소가 성능에 미치는 영향을 체계적으로 분석할 수 있는 프레임워크를 제공함. 특히, 벨로시티 예측과 연속 시간 학습은 기존 디퓨전 모델의 한계를 극복하고, 샘플링 과정의 유연성을 높이는 데 기여함. 그러나 본 연구는 주로 이미지 생성에 초점을 맞춘 것이며, 다른 생성 작업(예: 텍스트, 음성)에 대한 적용 가능성은 아직 검증되지 않았음. 또한, 수치적 특이점은 분석적으로는 해소되지만, 실제 구현에서는 여전히 문제가 될 수 있음.
실용적 활용
SiT는 고해상도 이미지 생성, 생성 모델의 샘플링 효율성 향상, 디퓨전 모델의 학습-샘플링 분리 설계 등에 활용 가능함. 특히, 디퓨전 계수를 학습과 독립적으로 조정할 수 있는 점은 생성 모델의 실시간 응용(예: VR, 게임 콘텐츠 생성)에 유리함.