터미널 속도 매칭을 통한 초고속 이미지 샘플링
초록
본 논문은 확산 모델의 샘플링 단계에서 초기 속도가 아닌 최종(터미널) 속도를 기준으로 한 “터미널 속도 매칭(TVM)” 기법을 제안한다. TVM은 단일 스텝 모델을 터미널 속도에 맞추어 가이드함으로써 1‑NFE(단일 함수 평가) 샘플링을 가능하게 하고, ImageNet‑256/512에서 기존 방법 대비 품질·속도 모두에서 우수한 성능을 보인다.
상세 분석
확산 기반 생성 모델은 고품질 이미지를 생성하기 위해 수백에서 수천 단계의 역확산 과정을 거친다. 최근 연구들은 단계 수를 크게 줄이면서도 성능을 유지하려는 “one‑step” 혹은 “few‑step” 샘플링에 집중하고 있다. 이러한 시도는 주로 초기 속도(초기 노이즈)와 그에 대응하는 스코어 함수(또는 노이즈 예측기)를 이용해 샘플을 추정한다. 그러나 초기 속도에만 의존하면 역확산 경로의 전반적인 동역학을 충분히 반영하지 못해, 특히 고해상도 이미지에서 품질 저하가 두드러진다.
본 논문은 이 한계를 “터미널 속도”라는 새로운 기준으로 전환한다. 터미널 속도는 역확산 과정의 최종 단계에서 기대되는 노이즈 레벨(또는 이미지와 가장 가까운 상태)의 속도를 의미한다. 저자들은 확산 방정식의 시간 역전성을 활용해, 최종 속도를 미리 예측하고 이를 역방향 샘플링에 직접 삽입한다. 구체적으로, 사전 학습된 one‑step 모델에 대해 손실 함수를 “터미널 속도와의 차이”로 재정의하고, 이를 최소화하도록 파라미터를 튜닝한다. 이 과정에서 사용되는 핵심 아이디어는 다음과 같다.
- 속도 매칭 목표 정의: 기존 방법은 초기 노이즈와 목표 이미지 사이의 거리만 최소화했지만, TVM은 최종 속도와 목표 이미지 사이의 거리까지 동시에 최소화한다. 이는 역확산 경로 전체를 보다 균형 있게 설계하게 만든다.
- 단일 함수 평가(1‑NFE) 구현: 터미널 속도를 직접 제공함으로써, 샘플링 과정에서 중간 단계의 스코어를 반복적으로 평가할 필요가 없어진다. 따라서 한 번의 네트워크 호출만으로 최종 이미지를 생성할 수 있다.
- 수치 안정성 확보: 터미널 속도는 일반적으로 낮은 노이즈 레벨에 해당하므로, 기존의 고노이즈 초기값 대비 수치적 불안정성이 크게 감소한다. 저자들은 이를 증명하기 위해 다양한 학습 스케줄과 정규화 기법을 실험하였다.
- 이미지 품질 향상: 실험 결과, TVM은 FID(Frechet Inception Distance)와 IS(Inception Score) 모두에서 기존 one‑step 방법보다 현저히 낮은(좋은) 값을 기록한다. 특히 512×512 해상도에서 1‑NFE 샘플링에도 불구하고, 기존 4‑NFE 방법과 거의 동등한 품질을 달성한다.
이러한 설계는 이론적 근거와 실험적 검증을 동시에 제공한다. 저자들은 확산 과정의 연속적인 미분 방정식 해석을 통해 터미널 속도가 최적의 “경로 종점”임을 보였으며, 이를 기반으로 손실 함수와 학습 전략을 설계하였다. 또한, 다양한 데이터셋(CIFAR‑10, LSUN‑Bedroom, ImageNet)에서 일반화 능력을 검증했으며, 특히 고해상도 ImageNet‑512에서 1‑NFE 샘플링이 가능한 최초의 사례 중 하나로 평가받는다.
한계점으로는 터미널 속도 예측을 위한 별도 사전 학습 단계가 필요하다는 점과, 매우 낮은 노이즈 레벨에서 발생할 수 있는 “모드 붕괴” 현상을 완전히 방지하지 못한다는 점을 들 수 있다. 향후 연구에서는 다중 터미널 속도(다중 목표 레벨)를 동시에 매칭하거나, 조건부 생성(텍스트‑투‑이미지 등) 상황에서 TVM을 확장하는 방안을 제시한다.
전반적으로 TVM은 “속도”라는 물리적 메타포를 활용해 확산 모델의 샘플링 효율성을 크게 향상시키는 혁신적인 접근법이며, 특히 실시간 이미지 생성이나 모바일 디바이스와 같이 연산 자원이 제한된 환경에서 큰 파급 효과를 기대할 수 있다.