Geometry-Aware Dataset Condensation for Diffusion Model Training

본 논문은 실질적인 서브셋 선택 문제를 일방향 부분 최적 운송(one-sided partial optimal transport)과 의미론적 정규화(semantic regularization)를 사용하여 기하학적 구조를 인지하는 분포 정렬 문제로 재정의함으로써, 효과적인 확산 모델 학습에 필요한 기하학적 구조와 분포적 충실도를 보존하는 압축된 데이터셋을 구축하는 방법론인 기하학 인지 데이터셋 응축(Geometry-Aware Dataset Condensation, GADC)을 제안한다.

핵심

당신이 마스터 셰프(하나의 확산 모델(Diffusion Model))에게 완벽한 요리를 하는 법을 가르치고 싶다고 상정해 봅시다. 전통적인 방식이라면 수백만 개의 레시피와 식재료가 담긴 거대한 도서관(전체 데이터셋)을 주었을 것입니다. 이 방식은 효과적이긴 하지만, 모든 것을 읽어내는 데 시간이 너무 오래 걸리고, 저장 공간 비용이 엄청나며, 이를 처리하기 위한 거대한 주방이 필요합니다.

**데이터셋 응축(Dataset Condensation)**은 단 몇십 개의 완벽한 레시피로 구성된 아주 작은 "치트 시트(요약본)"를 만들어, 불필요한 군더더기 없이 셰프에게 필요한 모든 것을 가르쳐주는 아이디어입니다.

하지만 기존의 치트 시트를 만드는 시도들은 두 가지 큰 문제점이 있었습니다:

1. 합성된 치트(Synthetic Cheats): 어떤 방식은 아예 새로운 레시 recipe를 처음부터 발명하려고 했습니다. 그 결과? 레시피는 기괴해 보였고 맛은 가짜처럼 느껴져 셰프를 혼란스럽게 만들었습니다.
2. 잘못된 선택(Bad Selection): 다른 방식은 기존 레시피 중 "최고"인 것들을 단순히 골라내는 것이었습니다. 하지만 이들은 "이 레시피가 얼마나 어려운가?"와 같은 단일하고 단순한 점수에 기반해 선택했습니다. 이는 전체적인 그림을 놓치게 만들었고, 중요한 변형들을 누락시켜 편향된 메뉴를 만들었습니다.

이 논문은 그 치트 시트를 만드는 새로운 방법인 **기하학적 인지 데이터셋 응축(Geometry-Aware Dataset Condensation)**을 제안합니다. 이 방법이 어떻게 작동하는지 쉬운 비유를 통해 설명하겠습니다.

1. 문제점: "일차원적" 실수

전체 데이터셋이 동네, 공원, 마천루가 있는 거대하고 복잡한 도시라고 상상해 보세요(데이터 분포).

• 기존 방식 (D2C): 이 방식은 "난이도"라는 하나의 직선 위에 집들을 나열하여 가장 좋은 집을 고르려 했습니다. 이는 3D 도시에서 집을 고를 때 오직 '높이'만 보고 결정하는 것과 같습니다. 당신은 높은 마천루 하나는 고를 수 있겠지만, 필수적인 아늑한 오두막집은 놓칠 수 있습니다. 즉, 도시의 모양을 잃어버리는 것입니다.
• 목표: 우리는 도시의 전체적인 모양을 완벽하게 나타낼 수 있도록, 공원, 거리, 동네의 구조를 그대로 유지하면서 작은 집 그룹을 선택해야 합니다.

2. 해결책: "일방향 부분 운송(One-Sided Partial Transport)"

저자들은 **최적 운송(Optimal Transport)**이라는 수학적 도구를 사용합니다. 이는 물류 회사가 창고(전체 데이터셋)에서 더 작은 새 창고(응축된 서브셋)로 화물을 옮기는 것과 같습니다.

• 기존 방식 (균형 운송): 기존의 규칙은 "큰 창고에 있는 모든 모래알을 무게를 완벽히 맞춰서 작은 창고로 옮겨야 한다"라고 말했습니다.
  • 결함: 작은 창고는 매우 작기 때문에, 무게 할당량을 채우기 위해 물류 회사는 도시의 가장자리(저밀도 구역)에 있는 쓸모없는 무거운 모래까지 끌고 와야 합니다. 이는 지도를 왜곡시킵니다.
• 새로운 방식 (일방향 부분 운송): 저자들은 이렇게 말합니다. "우리는 중요한 화물만 옮기면 된다. 도시 외곽의 빈 공간에 있는 모래까지 옮길 필요는 없다."
  • 이점: 이를 통해 작은 창고는 도시의 "핵심", 즉 번화한 거리와 인기 있는 동네에 온전히 집중할 수 있습니다. 이는 작은 서브셋이 노이즈에 휘둘리지 않고 원본 데이터의 진정한 기하학적 구조(모양과 구조)를 포착하도록 보장합니다*합니다.

3. 안전장치: "통계적 정규화(Statistical Regularization)"

화물을 옮기는 것만으로는 충분하지 않습니다. 우리는 새 창고가 여전히 원래의 도시와 같은 느낌을 주는지 확인해야 합니다. 저자들은 두 가지 "안전장치"를 추가했습니다.

• 평균-분산 체크: 그들은 작은 창고의 건물의 평균 "높이"와 "퍼짐 정도"가 큰 도시와 일치하는지 확인합니다. 큰 도시에 높고 낮은 건물이 섞여 있다면, 작은 도시에도 동일한 혼합 비율이 있어야 합니다.
• 신뢰도 체크: 그들은 선택된 집들이 명확하게 인식 가능한지 확인합니다. 만약 어떤 집이 헛간인지 차고인지 구분하기 힘든 흐릿한 덩어리처럼 보인다면, 그 집은 탈락시킵니다. 이는 "셰프"가 모호한 예시 때문에 혼란을 겪지 않도록 하기 위함입니다.

4. 전략: "탐욕적 구축 + 교체 정교화(Greedy Construction + Swap Refinement)"

그렇다면 이 구체적인 집들을 어떻게 고를까요? 모든 가능한 조합을 다 확인할 수는 없습니다(조합이 너무 많기 때문입니다!). 그래서 그들은 2단계 전략을 사용합니다.

1. 탐욕적 구축 (Greedy Building): 빈 부지에서 시작하여 한 번에 집을 하나씩 추가하며, 현재 시점에서 지도를 가장 많이 개선하는 집을 항상 선택합니다. 이는 퍼즐 조각을 하나씩 맞춰가는 것과 같습니다.
2. 교체 (The Swap): 퍼즐이 완성되면, 실수를 찾아냅니다. "이 구석에 있는 집은 잘 안 맞네. 저 바깥에 있는 집이랑 바꾸자." 그들은 지도가 완벽해질 때까지 계속해서 교체 작업을 수행합니다.

결과

연구진이 AI 이미지 생성기를 훈련시키기 위해 ImageNet(140만 개의 이미지가 담긴 거대한 데이터베이스)에 이 방법을 테스트했을 때의 결과입니다:

• 더 나은 품질: AI가 생성한 이미지는 이전 방식들에 비해 훨씬 더 선명하고 다양했습니다(더 낮은 FID 점수).
• 효율성: 이들은 원래 데이터의 단 **0.8%**만을 사용하여(140만 개 대신 10,000개의 이미지) 전체 데이터의 무작위 조각을 사용하는 것보다 더 나은 결과를 얻었습니다.
• 속ness: 이 10,000개의 이미지를 선택하는 과정은 이전 방식들보다 훨씬 빨랐습니다.

요약하자면:
이 논문은 강력한 AI를 작은 데이터셋으로 훈련시키려면, 단순히 "가장 어렵거나" "가장 쉬운" 예시를 고르는 것이 아니라, 노이즈가 섞인 빈 가장자리를 무시하면서도 원본의 방대한 데이터가 가진 모양, 구조, 다양성을 완벽하게 보존하는 아주 작은 이미지 그룹을 수학적으로 선택해야 한다는 것을 가르쳐줍니다. 이는 마치 미술사의 전체 컬렉션이 가진 영혼을 단 하나의 전시실에 담아내는 박물관 큐레이팅과 같습니다.

기술 요약

기술 요약: 확산 모델 학습을 위한 기하학 인지적 데이터셋 응축 (Geometry-Aware Dataset Condensation)

1. 문제 정의

데이터셋 응축(Dataset condensation)은 딥러닝 모델 학습에 수반되는 저장 및 계산 비용을 줄이기 위해, 합성 또는 선택을 통해 실제 데이터로부터 압축된 데이터셋을 구축하는 것을 목표로 합니다. 기존의 응축 방법들은 판별 작업(예: 분류)에서는 성공적인 성과를 보여왔으나, 확산 모델(Diffusion models)을 학습시키기에는 부적합합니다.

• 합성(Synthesis)의 한계: 연속적인 픽셀 최적화를 통해 합성 데이터를 생성하는 방식은 낮은 충실도(low-fidelity)의 샘플을 생성하는 경우가 많습니다. 확산 모델은 노이즈와 구조적 왜곡에 매우 민감하므로, 이러한 합성 샘플은 실제 모델링에 적합하지 않습니다.
• 선택(Selection)의 한계: 기존의 실제 서브셋 선택 방식은 고정된 기준이나 휴리스틱 순위(예: 1차원 축을 따라 확산 난이도에 의해 이미지를 정렬하는 D2C)에 의존합니다. 이러한 접근 방식은 확산 모델의 우도(likelihood) 목적 함수에 필요한 복잡한 분포 기하학을 보존하지 못합니다. 이들은 종 대개 확산 학습과 정렬된 원칙적인 최적화 목적 함수가 부족하여, 전체 데이터 분포와 정렬되지 못하고 필요한 매니폴드 구조를 포착하지 못하는 서브셋을 생성하게 됩니다.

핵적인 과제는 확산 모델의 우도 기반 학습 목적에 최적화되어, 전체 데이터셋의 기하학적 구조와 분포적 충실도를 보존하는 압축된 실제 서브셋을 선택하는 것입니다.

2. 방법론

저자들은 **기하학 인지적 데이터셋 응축(Geometry-Aware Dataset Cond(GADC)**을 제안하며, 이는 실제 서브셋 선택을 기하학 인지적 분포 정렬 문제로 재정의합니다. 이 프레임워크는 세 가지 주요 구성 요소로 이루어져 있습니다.

A. 일방향 부분 최적 운송 (One-Sided Partial Optimal Transport, POT)
작은 서브셋과 전체 데이터셋 사이의 용량 불일치(capacity mismatch)를 해결하기 위해, 저자들은 일방향 부분 최적 운송을 채택합니다.

• 개념: 전체 질량 정렬을 강제하는 균형 잡힌 OT(balanced OT)와 달리, 일방향 POT는 소스 질량(선택된 서브셋)은 완전히 운송되도록 허용하면서 용량 제약 조건 하에 타겟 질량(전체 데이터셋)을 완화합니다.
• 메커니즘: 이를 통해 밀도가 낮은 주변 영역의 일치하지 않는 질량을 허용함으로써, 운송을 데이터 매니폴드의 고밀도, 기하학적으로 안정적인 지배적 영역에 집중시킬 수 있습니다.
• 구현: 이 문제는 "더미 소스(dummy-source)" 재정의와 엔트로피 정규화를 사용하여 효율적으로 해결되며, Sinkhorn 반복법을 통해 최적화됩니다. 이를 통해 서브셋이 대표적인 영역에 집중하면서도 필요한 커버리지를 유지할 수 있는 유연한 질량 할당이 가능해집니다.

B. 통계적 및 의미론적 정규화
기하학적 정렬을 보완하고 분포적 충실도를 보장하기 위해 두 가지 경량 정규화 항이 도입됩니다.

• 평균-분산 정규화 ($L_{sta}$): 선택된 서브셋과 전체 데이터셋 간의 특징 표현(feature representation)의 1차 및 2차 통계량(평균 및 분산)을 정렬하여 전역적인 분포 충실도를 보존합니다.
• 신뢰도 정규화 ($L_{conf}$): 예측된 클래스 확률이 낮은 샘플의 선택을 억제함으로써 의미론적 일관성을 보장합니다. 이는 정렬을 저해할 수 있는 신뢰할 수 없는 기하학적 앵커(anchor)의 포함을 방지합니다.

C. 2단계 이산 최적화 전략
고정된 크기의 서브셋을 선택하는 것은 조합 최적화 문제이므로, 저자들은 효율적인 2단계 솔버를 제안합니다.

1. 1단계 (탐욕적 기하 가이드 선택): 복합 목적 함수(POT 손실 + 정규화 항)를 최소화하는 샘플을 점진적으로 구축하여 광범위한 매니폴드 커버리지를 설정합니다.
2. 2단계 (교환 기반 정제): 선택된 샘플과 선택되지 않은 샘플 간의 쌍별 교환(pairwise swaps)을 수행하여 초기 단계의 근시안적인 선택을 수정하고, 전역적인 기하학적 정렬과 일관성을 더욱 최적화하는 정제 단계입니다.

3. 주요 기여

1. 문제 재정의: 본 논문은 데이터셋 응축을 휴리스틱한 순위 지정이나 스칼라 기반 선택에서 벗어나, 분포 정렬 문제로 재정의합니다. 또한 대표적인 실제 서브셋을 선택하기 위해 통계적 정규화가 포함된 일방향 POT 목적 함수를 도입합니다.
2. 최적화 프레임워크: 고정된 기준이나 순위 기반 샘플링의 한계를 피하면서, 정렬 문제를 효율적으로 해결하는 2단계 이산 최적화 프레임워크(탐욕적 구축 후 교환 기반 정제)를 제안합니다.
3. 실험적 성능: 광범위한 실험을 통해 다양한 확산 변체(DiT, SiT), 서브셋 크기, 이미지 해상도, 학습 라운드에 걸쳐 기존 방식들보다 우수한 충실도와 분포 커버리지를 달s는 것을 입증했습니다.

4. 실험 결과

본 방법은 DiT-L/2 및 SiT-L/2 아키텍처를 사용하여 데이터 예산(budget) 0.8%(10K 이미지)부터 8%(100K 이미지)까지 ImageNet-1K에서 평가되었습니다.

• 성능 지표: 제안된 방법은 Random Sampling, K-Center, Herding, CCS, DQ, D2C를 포함한 베이스라인들과 비교하여 지속적으로 가장 낮은 FID와 가장 높은 Inception Score (IS), Precision, Recall을 달성했습니다.
  • 예시: ImageNet 256×256, 10K 데이터 예산(0.8%) 환경에서, 이 방법은 FID 3.43을 달성하여 D2C(4.20)와 Random(35.86)을 능가했습니다.
  • 예시: 512×512 해상도, 10K 예산 환경에서, 이 방법은 FID 6.17을 달성하여 D2C(14.8)보다 현저히 우수했습니다.
• 강건성(Robustness): 개선된 성능은 서로 다른 확산 변체(DiT 및 SiT), 서로 다른 평가 프로토콜(10K vs 50K 생성 샘플), 그리고 확장된 학습 반복 횟수(최대 300K) 전반에 걸쳐 유지되었습니다.
• 효율성: 선택 과정이 계산적으로 효율적입니다. 단일 GPU에서 이 방법은 ImageNet에서 10K 서브셋을 선택하는 데 5.5시간이 소요되었으며, 이는 D2C(41.9시간) 및 DQ(30.4시간)와 비교됩니다.

절제 연구 (Ablation Studies):

• POT 손실($L_{OT}$)을 제거하면 FID가 저하되고 정밀도/재현율이 감소하여, 기하학적 정렬의 필요성을 확인했습니다.
• 일방향 POT를 균형 잡힌 OT로 대체하면 성능이 저하되어, 주변 영역의 일치하지 않는 질량을 허용하는 것의 이점을 검증했습니다.
• 통계적($L_{sta}$) 또는 신뢰도($L_{conf}$) 정규화를 제거하면 각각 FID 또는 IS가 악화되었으며, 이는 전역 통계와 의미론적 명확성의 중요성을 강조합니다.

5. 의의 및 주장

본 논문은 데이터셋 응축을 기하학 인지적 분포 정렬 문제로 다룸으로써, 데이터 선택와 확산 모델 학습의 특수한 요구 사항 사이의 간극을 성공적으로 메웠다고 주장합니다.

• 데이터 효율성: 이 방법은 압축된 데이터셋을 통한 고품질 생성 모델링을 가능하게 하여, 확산 모델 학습과 관련된 저장 및 계산 비용을 크게 줄여줍니다.
• 충실도: 합성 생성 방식과 달리, 이 접근 방식은 실제 샘플의 고충실도 구조를 유지하면서 서브셋이 전체 데이터 분포의 기하학적 구조와 일치하도록 보장합니다.
• 실용성: 제안된 2단계 최적화 전략은 계산적으로 효율적이고 확장 가능하여, 자원이 제한된 환경에서 실용적인 솔루션이 됩니다.

저자들은 자신들의 방법이 확산 스타일의 생성 학습을 위한 신뢰할 수 있는 데이터 기질(substrate)을 제공하며, 모델 측면의 최적화와 병행하여 효율성을 개선할 수 있는 상호 보완적인 데이터 중심 경로를 제시한다고 결론짓습니다.