데이터셋 증류를 최적 양자화로 풀다: 잠재공간 클러스터링 기반 새로운 프레임워크

초록

본 논문은 기존의 분리형(디스엔탱글드) 데이터셋 증류 방법을 잠재공간에서의 최적 양자화 문제로 재구성한다. 저자들은 양자화 이론과 Wasserstein 거리·바리센터 개념을 이용해 증류된 데이터가 원본 분포에 점근적으로 수렴함을 증명하고, CLVQ(경쟁 학습 벡터 양자화) 기반의 클러스터링 알고리즘을 제안한다(D​DOQ). 실험에서는 ImageNet‑1K에서 기존 SOTA인 D⁴M을 능가하는 정확도와 모델 간 일반화 능력을 보이며, 강력한 Diffusion Transformer와 결합해 최고 수준의 증류 성능을 달성한다.

상세 분석

이 논문은 데이터셋 증류를 두 단계로 나누는 기존의 바이레벨 최적화 접근법이 계산량과 모델 의존성에서 한계를 가진다는 점을 지적한다. 대신, 인코더‑디코더 구조를 이용해 고차원 이미지 데이터를 저차원 잠재공간으로 매핑하고, 그 공간에서 K개의 대표점과 가중치를 찾는 최적 양자화 문제로 전환한다. 양자화는 ‘Quadratic distortion’ G를 최소화하는 점 집합 {x₁,…,x_K}와 각 Voronoi 셀의 질량 µ(C_i) 을 가중치로 하는 확률 측도 ν_K = ∑_i µ(C_i)δ(x_i) 로 정의된다. 저자들은 Proposition 1을 통해 G의 최소화가 Wasserstein‑2 거리에서 원본 분포와 가장 가까운 K‑점 지원 측도를 찾는 것과 동등함을 보인다.

Theorem 1에서는 최적 양자화가 인구 위험의 그래디언트 ∇R(θ) 에 대한 기대값을 O(K^{-1/d}) 속도로 근사한다는 수렴율을 제시한다. 여기서 d 는 잠재공간 차원, K 는 양자화 점 수이며, 차원을 낮출수록 같은 K에 대해 더 빠른 수렴이 가능함을 의미한다. 이는 기존 디스엔탱글드 방법이 저차원 잠재공간에서 클러스터링을 수행하는 경험적 성공을 이론적으로 뒷받침한다.

알고리즘적 구현은 CLVQ(Competitive Learning Vector Quantization) 혹은 미니배치 k‑means와 동일한 업데이트 규칙을 사용한다. 샘플 x 를 뽑아 가장 가까운 중심 c_i 를 찾고, 학습률 γ_i (보통 1/|C_i|) 를 곱해 중심을 이동시킨다. 동시에 각 중심에 대한 Voronoi 셀 질량을 온라인으로 추정해 가중치 w_i = µ(C_i) 를 얻는다. 이렇게 얻은 (점, 가중치) 쌍을 디코더에 입력하면, 가중치를 반영한 합성 이미지가 생성된다.

다음으로 저자들은 확산 모델(Score‑based diffusion)과의 연결을 분석한다. 잠재공간에서의 양자화 오차가 이미지 공간으로 전달되는 정도를 SDE dx = f(x,t)dt + g(t)dW와 그 역방향식으로 모델링한다. Theorem 2(논문에 명시되지 않았지만 암시)에서는 Wasserstein‑2 수렴이 일정 시간 δ>0 후에도 유지된다는 점을 보이며, 이는 증류된 초기화가 실제 데이터와 동일한 훈련 역학을 유도함을 의미한다.

실험에서는 동일한 Latent Diffusion Model(LDM) 백본을 사용해 D⁴M과 DDOQ를 비교한다. ImageNet‑1K에서 IPC(Images per Class) = 1, 5, 10 조건에서 DDOQ가 평균 +2.3% 정확도 향상을 보였으며, 특히 높은 IPC(>50)에서는 기존 방법을 크게 앞선다. 또한, 증류된 노이즈 초기화를 Diffusion Transformer(DiT)와 결합했을 때, SOTA인 D4M‑Guided와 Diffusion‑Guidance 기반 방법들을 모두 능가하는 Top‑1 ≈ 78% 성능을 기록했다. 교차‑아키텍처 실험에서도 ResNet‑50, ViT‑B/16 등 다양한 모델에 대해 일반화 격차가 1% 이하로 감소했다.

결론적으로, 이 논문은 디스엔탱글드 데이터셋 증류를 최적 양자화라는 수학적으로 엄밀한 프레임워크에 매핑함으로써, 이론적 수렴 보장과 실용적 성능 향상을 동시에 달성한다는 점에서 중요한 기여를 한다.