Dataset Distillation: A Comprehensive Review 논문 요약

이 논문은 최근 연구가 활발한 데이터셋 증류(Dataset Distillation, DD) 또는 데이터셋 응축(Dataset Condensation, DC) 에 대한 포괄적인 리뷰와 요약을 제공합니다. DD의 목표는 원본 대규모 데이터셋으로부터 합성 샘플을 포함하는 훨씬 작은 데이터셋을 도출하여, 이 합성 데이터셋으로 훈련된 모델이 원본 데이터셋으로 훈련된 모델과 유사한 성능을 달성하도록 하는 것입니다. 이는 대규모 데이터의 저장 및 전송 부담, 훈련 과정의 복잡성, 그리고 프라이버시 및 저작권 문제를 완화하는 데 기여합니다.

그림 1.
이것은 데이터셋 증류(Dataset Distillation, DD)에 대한 개요입니다. 데이터셋 증류는 원본의 대규모 데이터셋으로부터 훨씬 적은 수의 합성 샘플로 구성된 작은 데이터셋을 생성하는 것을 목표로 합니다. 이렇게 생성된 작은 데이터셋으로 학습된 모델은 원본의 대규모 데이터셋으로 학습된 모델과 유사한 테스트 성능을 보여야 합니다.

1. 관련 연구 (Related Works)

지식 증류 (Knowledge Distillation, KD): KD는 큰 teacher 네트워크의 지식을 작은 student 네트워크로 전이하여 모델 압축을 목표로 합니다. DD는 데이터 압축을 목표로 한다는 점에서 KD와 개념적으로 유사하지만 기술적으로는 다릅니다. Data-free KD는 합성 샘플 생성을 포함하나, teacher 모델을 입력으로 받고 student 모델의 성능 향상이 주 목적이며 더 많은 합성 샘플을 생성합니다.

코어셋 또는 인스턴스 선택 (Core-set or Instance Selection): 원본 데이터셋에서 가장 대표적인 샘플을 선택하여 크기를 줄이는 방법입니다. DD는 새로운 데이터를 합성한다는 점에서 차이가 있으며, 이는 더 나은 프라이버시 보호와 최적화된 성능을 제공할 수 있습니다.

생성 모델 (Generative Models): DD의 한 최적화 목표인 분포 매칭은 생성 모델과 유사합니다. 그러나 생성 모델은 사실적인 데이터 생성이 주 목적이지만, DD는 훈련 효율성 및 저장 비용 절감을 위한 정보 함축적인 데이터 생성을 목표로 합니다. 일부 DD 방법은 생성 모델을 합성 데이터 파라미터화에 활용합니다.

하이퍼파라미터 최적화 (Hyperparameter Optimization): DD는 각 합성 샘플을 고차원 하이퍼파라미터로 간주할 경우 하이퍼파라미터 최적화 문제로 전환될 수 있으며, Wang et al. [18]의 방법은 gradient-based 하이퍼파라미터 최적화와 아이디어를 공유합니다.

2. 데이터셋 증류 방법론 (Dataset Distillation Methods)

2.1 정의 및 일반적인 워크플로우

DD는 원본 데이터셋 T로부터 작은 합성 데이터셋 $S=(X_S, Y_S)$ 를 학습하는 문제로 정의됩니다: $S = \arg \min_S \mathcal{L}(S, T)$, 여기서 $\mathcal{L}$ 은 DD 목적 함수입니다.

일반적인 워크플로우는 다음과 같습니다 (Algorithm 1):

1. $S$ 초기화: Gaussian noise 또는 원본 데이터셋 $T$의 랜덤 샘플로 초기화합니다.

2. 반복 루프:
   - 네트워크 $θ$ 를 얻습니다 (랜덤 초기화 또는 캐시된 체크포인트에서 로드).
   - 필요시 $θ$를 $S$ 또는 $T$ 를 사용하여 업데이트하고 캐시합니다.
   - $\mathcal{L}(S, T)$ 를 통해 $S$ 를 업데이트합니다.
   
   

2.2 최적화 목적 함수 (Optimization Objectives)

• 성능 매칭 (Performance Matching): 합성 데이터셋 $S$로 훈련된 모델의 성능이 원본 데이터셋 $T$에 대해 가장 낮은 손실을 보이도록 $S$를 최적화합니다.

  • Meta-Learning 기반: Wang et al. [18]의 시초 방법입니다. 내재 루프(inner loop)에서 $S$를 사용하여 모델 $θ$를 훈련하고, 외재 루프(outer loop)에서 $T$에 대한 $θ$의 손실을 통해 $S$를 역전파합니다. 이는 Bi-level optimization 문제를 형성하며, 재귀적 계산 그래프를 통한 역전파로 인해 계산 비용과 GPU 메모리 요구량이 큽니다. 모멘텀($β$)을 추가한 $θ$ 업데이트:
    $\quad m^{(0)} = 0,$ $\quad m^{(t)} = \beta m^{(t-1)} + \nabla \ell(S; \theta^{(t-1)}),$ $\quad \theta^{(t)} = \theta^{(t-1)} - \eta m^{(t)}.$

  • Kernel Ridge Regression (KRR) 기반: KIP [20, 21], FRePo [19], RFAD [85]. 내재 루프 훈련을 피하고 KRR의 닫힌 형태 해를 사용하여 계산 효율성을 높입니다. 선형 모델에 적합하며, 신경망 $f_\theta$를 통한 고차원 특징 공간 투영을 활용합니다. 목적 함수는:
    $\quad \mathcal{L}(S, T) = E_{\theta\sim\Theta}[\|Y_t - K^\theta_{X_t X_S} (K^\theta_{X_S X_S} + \lambda I)^{-1} Y_S\|^2].$ FRePo는 특징 추출기($f_\theta$)는 훈련하고 선형 분류기의 최적 파라미터는 KRR로 얻는 분해된 접근법을 사용합니다.

• 파라미터 매칭 (Parameter Matching): 합성 데이터셋과 원본 데이터셋으로 훈련된 신경망의 파라미터가 일관성을 유지하도록 합니다.

  • Single-Step Parameter Matching (Gradient Matching): DC [22], DSA [23], IDC [82]. 네트워크를 $S$와 $T$로 각각 한 스텝 업데이트한 후, 결과 그라디언트 간의 거리를 최소화합니다. $\mathcal{L}(S, T)= E_{\theta^{(0)}\sim\Theta}[D(S, T; \theta^{(t)})]$ (여기서 $D$는 그라디언트 거리 측정 함수). Zhao et al. [22]은 각 클래스별로 레이어 및 채널별 코사인 유사도를 사용했습니다. 과적합 및 그라디언트 소실 문제를 완화하기 위해 IDC [82]는 네트워크를 $T$로 훈련합니다.

  • Multi-Step Parameter Matching (Training Trajectory Matching, MTT): MTT [81], DDPP [94], TESLA [95], FTD [96], HaBa [99]. $\theta^{(0)}$에서 시작하여 $S$와 $T$로 각각 $T_S$와 $T_t$ 스텝을 훈련한 후, 두 훈련 궤적의 끝점 간 거리를 최소화합니다. $\quad \mathcal{L}(S, T)= E_{\theta^{(0)}\sim\Theta}[D(\theta^{(T_S)}_S,\theta^{(T_t)}_T)],$ $\quad \theta^{(t)}_S = \theta^{(t-1)}_S - \eta\nabla\ell(S; \theta^{(t-1)}_S),$ $\quad \theta^{(t)}_T = \theta^{(t-1)}_T - \eta\nabla\ell(T; \theta^{(t-1)}_T).$ 긴 훈련 궤적에 대한 역전파는 메모리 비효율성을 초래하며, TESLA [95]는 이를 상수 메모리 복잡도로 해결합니다. DDPP [94]는 매칭하기 어려운 파라미터를 가지치기하여 강건성을 향상합니다.

• 분포 매칭 (Distribution Matching): 합성 데이터의 분포가 실제 데이터의 분포를 근사하도록 최적화합니다. DM [24], CAFE [97].

  • 신경망을 임베딩 함수로 사용하여 특징 공간에서 분포를 비교합니다. DM은 각 클래스별 출력 임베딩의 평균 벡터($\mu$) 간 거리를 최소화합니다:
    $\quad D(S, T; \theta) = \sum_{c=0}^{C-1}\|\mu_{\theta,s,c}\mu_{\theta,t,c}\|^2.$ CAFE는 각 네트워크 계층의 특징 통계량 일관성을 강제합니다.

• 목적 함수 간의 연결 (Connections between Objectives):

  • Proposition 1: KRR 모델의 성능 매칭 목표는 최적 파라미터 매칭 목표(무한 스텝 파라미터 매칭)와 동등합니다.
  • Proposition 2: 랜덤 특징 추출기를 따르는 KRR 모델에서 1차 분포 매칭 목표는 각 클래스에 대한 그라디언트 매칭과 근사적으로 동일합니다.
  • Proposition 3: KRR 모델에서 2차 분포 매칭 목표는 그라디언트 매칭 목표의 상한을 최적화합니다.
    
    

2.3 합성 데이터 파라미터화 (Synthetic Data Parameterization)

합성 데이터 $S$를 직접 픽셀 공간에서 최적화하는 대신, 더 압축적이고 정보 함축적인 방식으로 표현하는 기술입니다. $X_S = \{g_\phi(z_j)\}$ 형태로, $z_j$는 잠재 코드이고 $g_\phi$ 는 생성 함수입니다.

• Differentiable Siamese Augmentation (DSA): 데이터 증강 정책을 $g_\phi$로 사용하여 합성 데이터의 일반화 능력을 향상합니다.
• Upsampling: IDC [82]. 원본 이미지를 다운샘플링하여 $z$로 저장하고, $g_\phi$는 업샘플링 함수로 사용하여 동일 저장 예산 내에서 더 많은 합성 이미지를 얻습니다.
• Generators and Latent Vectors: 잠재 벡터 $z$로부터 데이터를 생성하는 생성 모델을 $g_\phi$로 활용합니다.
  • GAN: IT-GAN [67]은 사전 훈련된 GAN 생성기를 사용하여, GAN 역전(inversion)으로 얻은 잠재 벡터를 최적화합니다.
  • Addressing Matrices: Deng et al. [84]은 학습 가능한 주소 지정 행렬(addressing matrices)을 통해 공유 메모리 표현으로부터 합성 데이터를 구성합니다.
  • Decoders: HaBa [99]와 KFS [70]는 비선형 디코더를 사용하여 잠재 코드와 디코더의 조합으로 합성 샘플을 생성합니다.
    
    

2.4 라벨 증류 방법 (Label Distillation Methods)

대부분의 방법은 클래스 라벨 $Y_S$를 one-hot 형태로 고정하지만, 일부 연구는 라벨을 학습 가능하게 하거나 soft 라벨을 사용하면 성능이 향상될 수 있음을 보여줍니다.

3. 응용 (Applications)

DD는 높은 압축률 덕분에 다양한 분야에서 혁신적인 응용을 이끌어냈습니다.

• 지속 학습 (Continual Learning): 제한된 메모리 버퍼에서 이전 태스크의 지식을 보존하는 데 활용됩니다.
• 연합 학습 (Federated Learning): 클라이언트가 모델 업데이트 대신 합성 데이터셋을 전송하여 통신 비용을 줄이고 프라이버시를 보호합니다.
• 신경망 아키텍처 탐색 (Neural Architecture Search, NAS): DD로 생성된 작은 데이터셋을 NAS의 프록시 셋으로 사용하여 모델 평가를 가속화합니다.
• 프라이버시, 보안 및 강건성 (Privacy, Security and Robustness): 합성 데이터가 실제 데이터를 재현하기 어려워 데이터 프라이버시를 제공하고, 강건성 최적화에 기여하며, 백도어 공격에 대한 취약성도 연구됩니다.
• 그래프 신경망 (Graph Neural Network, GNN): 그래프 구조 데이터의 증류를 통해 GNN 훈련 효율성을 향상합니다.
• 추천 시스템 (Recommender System): 대규모 사용자-아이템 상호작용 데이터를 요약하는 데 사용됩니다.
• 텍스트 분류 (Text Classification): 텍스트 데이터의 이산적인 특성으로 인해 어렵지만, 임베딩 공간에서 텍스트를 증류하는 연구가 진행됩니다.
• 지식 증류 (Knowledge Distillation): 지식 응축을 통해 불필요한 지식 전이를 줄이고 효율성을 높입니다.
• 의료 (Medical): 의료 데이터의 프라이버시 보호, 전송 및 저장 비용 절감을 위해 증류된 익명 데이터셋을 활용합니다.
• 패션, 예술 및 디자인 (Fashion, Art and Design): 시각적으로 흥미로운 타일형 증류 텍스처 생성 및 의상 호환성 모델링에 응용됩니다.

그림 2.
이중(bi-level) 학습 프레임워크는 합성 데이터셋 S로 메타 학습된 모델의, 실제 데이터셋 T에서의 메타 테스트 손실을 최적화하는 것을 목표로 합니다.

그림 3.
a) 단일 단계 파라미터 매칭. (b) 다중 단계 파라미터 매칭. 이들은 실제 데이터와 합성(가짜) 데이터를 사용해 학습된 모델 파라미터(매개변수)의 일관성(consistency)을 최적화한다. 단일 단계 파라미터 매칭의 경우, 이는 그래디언트(기울기)를 매칭하는 것과 동등하다. 다중 단계 파라미터 매칭의 경우, 이는 훈련 궤적(training trajectories)을 매칭하는 것으로도 알려져 있다.

그림 4.
분포 매칭(Distribution matching). 합성(synthetic) 데이터셋과 실제(real) 데이터셋에 대해, 어떤 네트워크(들)에서 추출된 특징(feature)들의 통계(statistics)가 서로 같아지도록 맞춘다.

그림 5.
기존 데이터셋 증류(DD) 방법들에 대한 분류 체계
Optimization objective (최적화 목표) : 어떤 기준을 바탕으로 합성 데이터셋을 최적화하는지
Fashion of updating networks (네트워크 업데이트 방식) : 데이터셋 증류 과정에서 신경망을 업데이트하는 방식
Synthetic data parameterization (합성 데이터 매개변수화) : 합성 데이터를 어떤 방식으로 표현하고 생성하는지
Fashion of learning labels (레이블 학습 방식) : 합성 데이터에 대한 레이블을 어떻게 학습하는지

그림 6.
합성 데이터 파라미터화(synthetic data parameterization). 이는 DD(데이터셋 증류)에서의 한 하위 분야이며, 데이터셋 증류의 목표들과는 직교(별개)한다. 왜냐하면 $S$를 생성하는 과정이 미분 가능하고, $S$를 업데이트하기 위한 그래디언트가 합성 파라미터까지 더 역전파될 수 있기 때문이다.

4. 실험 (Experiments)

MNIST, Fashion-MNIST, CIFAR-10, CIFAR-100, Tiny-ImageNet에서 ConvNet, AlexNet, ResNet, VGG 등의 아키텍처를 사용하여 성능과 훈련 비용을 평가했습니다.

• 증류 성능: FRePo가 CIFAR-10, CIFAR-100 등 복잡한 데이터셋에서 SOTA 성능을 보였고, MTT가 그 뒤를 이었습니다. DSA는 성능 향상에 기여했습니다.
• 교차 아키텍처 일반화: DD, DC, DSA, DM, MTT 등은 정규화 계층(normalization layer)에 따라 성능이 크게 저하되었습니다. 이는 합성 데이터셋이 특정 훈련 아키텍처의 유도 편향(inductive bias)을 인코딩함을 시사합니다. FRePo는 ConvNet-BN 기반으로 훈련되어 ResNet-BN, VGG-BN 등에서 비교적 강건한 전이성을 보였습니다.
• 훈련 비용: DD와 MTT는 unrolled gradient computation으로 인해 높은 런타임과 GPU 메모리를 요구했습니다. DM은 가장 낮은 GPU 메모리를 필요로 했습니다. FRePo(JAX 버전)는 IPC가 작을 때 매우 빠르고 효율적입니다.
  
  

5. 과제 및 개선 방향 (Challenges and Possible Improvements)

• 계산 비용: unrolled computational graphs를 통한 역전파는 여전히 큰 비용을 발생시킵니다. KRR 기반 방법이나 그래디언트 계산 흐름을 재구성하는 TESLA 같은 노력이 필요합니다.
• 스케일업 (Scaling Up):
  • 높은 압축률: IPC(images per class)가 증가할수록 기존 DD 방법의 지식 압축 능력이 상대적으로 떨어져, 선택 기반 방법과의 격차가 줄어들거나 역전됩니다.
  • 대규모 원본 데이터셋: ImageNet-1K 같은 대규모 데이터셋에서는 대부분의 DD 방법이 높은 GPU 메모리와 긴 훈련 시간으로 인해 적용이 어렵습니다.
  • 복잡한 네트워크: 대형 신경망에 대한 증류 및 평가 성능이 저조하며, 이는 작은 합성 데이터셋에서의 과적합 문제 때문일 수 있습니다.

• 다양한 아키텍처에 대한 일반화: 합성 데이터셋이 특정 훈련 아키텍처에 강하게 결합되어 전이성이 낮아지는 문제입니다. 네트워크를 합성 데이터셋과 독립적으로 업데이트하거나, 훈련 아키텍처의 다양성을 늘리는 방법을 모색해야 합니다.
• 다른 태스크 및 응용: 현재 DD 연구는 주로 이미지 분류에 집중되어 있습니다. 의미론적 분할, 객체 탐지, 기계 번역 등 다른 컴퓨터 비전 및 자연어 처리 태스크로의 확장이 필요합니다.
• 보안 및 프라이버시: DD의 잠재적 보안 및 프라이버시 문제는 아직 충분히 연구되지 않았습니다. DOORPING과 같은 백도어 공격에 대한 방어 메커니즘 연구가 필요합니다.
  
  

그림 7.
a) 루프당 실행 시간(한 루프에서 합성 데이터와 네트워크를 업데이트하는 데 필요한 시간을 포함). (b) 스텝당 실행 시간(합성 데이터를 업데이트하는 데 필요한 시간만 포함). (c) 최대 GPU 메모리 사용량. 모든 평가는 CIFAR-10에서 전체 배치(full batch) 학습 조건으로 수행되었다.

6. 결론 (Conclusion)

이 논문은 데이터셋 증류의 포괄적인 개요를 제공하며, 기존 DD 방법들을 최적화 목적(성능 매칭, 파라미터 매칭, 분포 매칭)에 따라 체계적으로 분류하고 그 이론적 연결성을 분석했습니다. 또한, 모든 기존 DD 방법이 따르는 공통 알고리즘 프레임워크를 제시하고, 지속 학습, 연합 학습, 프라이버시 보호 등 다양한 응용 분야에서의 활용을 다루었습니다. 실험을 통해 각 접근 방식의 정확도, 시간 효율성, 확장성을 비교 분석하여, 현재 DD 분야가 직면한 계산 비용, 스케일업, 아키텍처 일반화 및 보안 문제와 같은 주요 과제를 강조하고 미래 연구 방향을 제시했습니다.

그림 8.
CIFAR10에서 Random selection, K-Center, DD, DC, DSA, DM, MTT, FRePo-JAX를 사용하여 서로 다른 압축 비율에서의 성능 비교. 결과는 [75] 및 우리의 실험에서 얻어졌다.