노이즈 없는 디퓨전 특징 추출, CleanDIFT

초록

CleanDIFT는 사전 학습된 대규모 디퓨전 모델을 약 30분의 가벼운 미세조정만으로 깨끗한 이미지에서 타임스텝에 독립적인 고품질 특징을 추출한다. 기존 방식이 필요로 하던 노이즈 추가와 타임스텝 튜닝을 없애고, 다양한 비전 과제에서 기존 디퓨전 특징 및 앙상블 방법을 크게 능가한다.

상세 분석

본 논문은 디퓨전 모델이 학습 과정에서 다양한 수준의 가우시안 노이즈를 입력받아 그 노이즈를 제거하는 과제를 수행한다는 점에 착안한다. 이때 모델은 타임스텝 t에 따라 서로 다른 “목표”를 학습하게 되며, 높은 노이즈 단계에서는 이미지의 전반적인 구조만을, 낮은 노이즈 단계에서는 세부 텍스처와 색상 정보를 복원한다. 이러한 다중 목표 학습은 타임스텝마다 서로 다른 의미적 특성을 가진 내부 특징 맵을 생성한다는 것을 의미한다. 기존 연구들은 이러한 타임스텝‑특정 특징을 활용하기 위해 이미지에 인위적으로 노이즈를 추가하고, 최적의 타임스텝을 과제별로 탐색한다. 그러나 노이즈를 추가하면 원본 이미지의 정보가 손실되고, 노이즈 자체가 특징에 큰 비중을 차지한다는 실험적 증거(그림 3)를 제시한다. 즉, “노이즈 ε”가 특징 벡터의 상당 부분을 설명한다는 사실은 기존 방식이 본질적으로 정보를 파괴하고 있음을 보여준다.

CleanDIFT는 이 문제를 두 단계로 해결한다. 첫째, 사전 학습된 디퓨전 U‑Net을 복제한 가중 가능한 복사본을 만든다. 이 복사본은 깨끗한 이미지 x₀만을 입력받으며, 원본 디퓨전 모델은 동일 이미지에 노이즈 ε와 타임스텝 t를 적용해 입력한다. 둘째, 각 디코더 레이어 k에 대해 타임스텝‑조건부 선형(또는 비선형) 프로젝션 헤드 projₖ(·, t)를 도입한다. 학습 중에는 projₖ가 복제 모델의 특징 featₖᶜ(x₀)를 원본 디퓨전 모델의 타임스텝‑특정 특징 featₖ(x_t, t)와 정렬하도록, 부정 코사인 유사도 −sim(·,·)를 최소화한다. 타임스텝을 균등하게 샘플링해 여러 t에 대해 동시에 정렬을 강제함으로써, 복제 모델은 “모든 타임스텝의 정보를 하나의 타임스텝‑독립 특징 featᶜ(x₀)”에 압축한다.

핵심적인 설계 선택은 다음과 같다. (1) 프로젝션 헤드는 학습 단계에서만 사용되며, 추론 시에는 제거하고 복제 모델의 내부 표현을 바로 특징으로 활용한다. (2) 미세조정은 이미지‑캡션 쌍을 이용해 400 스텝(≈30 분, A100 1대)만 수행해도 충분히 수렴한다. (3) 전체 파이프라인은 기존 디퓨전 특징 추출 대비 8배 가량 빠르며, 추가적인 메모리 비용도 거의 없다.

실험에서는 Stable Diffusion 1.5와 2.1을 백본으로 사용해 네 가지 대표 과제(무감독 의미 매칭, 단일 이미지 깊이 추정, 의미 분할, 이미지 분류)를 평가한다. 모든 과제에서 CleanDIFT는 기존 노이즈‑기반 디퓨전 특징을 능가하고, 특히 의미 매칭에서는 이전 최첨단(DIFT)보다 큰 AP 상승을 기록한다. 또한 DINOv2와 같은 다른 자기지도 학습 특징과 결합했을 때도 시너지 효과가 확인되었다. 앙상블(여러 타임스텝 t를 동시에 사용) 대비 동일하거나 더 높은 정확도를 보이며, 연산량은 크게 감소한다. 이는 “노이즈‑의존성”이 제거된 특징이 보다 풍부하고 일관된 의미 정보를 담고 있음을 의미한다.

이 논문은 디퓨전 모델을 “특징 추출기”로 전환하는 새로운 패러다임을 제시한다. 기존의 “노이즈 → 복원” 흐름을 유지하면서도, 학습된 정렬 메커니즘을 통해 노이즈를 완전히 배제한 특징을 얻는다. 결과적으로 디퓨전 기반 특징은 이제 전통적인 비전 백본(ResNet, ViT 등)과 동등하거나 그 이상으로 활용 가능하며, 타임스텝 튜닝이라는 번거로운 하이퍼파라미터 관리 없이 다양한 다운스트림 작업에 바로 적용할 수 있다. 향후 연구는 더 큰 스케일의 디퓨전 모델(예: Stable Diffusion XL)이나 비디오/오디오 디퓨전 모델에 동일한 정렬 기법을 확장하는 방향으로 진행될 수 있다.