Squeeze, Recover and Relabel: Dataset Condensation at ImageNet Scale From A New Perspective 논문 요약

본 논문은 데이터 압축(Dataset Condensation) 작업을 위한 새로운 프레임워크인 Squeeze, Recover and Relabel (SRe2L)을 제안한다. 기존 데이터 압축 방법론들은 모델 훈련과 합성 데이터 업데이트 간의 바이레벨 최적화(bilevel optimization)로 인해 대규모 데이터셋, 다양한 모델 아키텍처 및 높은 이미지 해상도를 처리하는 데 제약이 있었다. SRe2L은 이 최적화 과정을 디커플링(decoupling)하여 이러한 한계를 극복한다. 제안된 방법은 합성 이미지의 임의 해상도 지원, 저렴한 훈련 비용 및 메모리 소비, 그리고 다양한 평가 네트워크 아키텍처로의 확장성을 특징으로 한다.

1. Introduction

데이터 압축은 대규모 데이터셋을 대표성 있는 작은 서브셋으로 줄여 훈련 시간과 저장 공간을 절약하면서도 원본 데이터의 핵심 정보를 보존하는 것을 목표로 한다. Meta-Model Matching, Gradient Matching, Distribution Matching, Trajectory Matching 등의 기존 방법들은 MNIST, CIFAR, Tiny-ImageNet과 같은 소규모 데이터셋이나 해상도가 낮은 ImageNet 서브셋에 주로 효과적이었다. 이는 바이레벨 최적화 과정에서 발생하는 엄청난 계산 비용과 메모리 오버헤드(inner loop의 unrolling) 때문이다.

그림 1.
왼쪽 그림은 ImageNet-1K에서 클래스당 10장(10 IPC: Images Per Class)만 쓰는 조건에서, 데이터 합성(synthesis) 시간과 정확도(accuracy)의 관계를 보여줘요.
사용한 모델은 ConvNet-4, ResNet-18/50/101이고, † 표시는 파라미터 1,000만(10M)짜리 ViT 모델을 뜻해요(참고문헌 [2]에 있는 ViT).

오른쪽 그림은 기존에 널리 쓰이는 바이레벨 최적화(bilevel optimization) 방식과, 저자들이 제안한 디커플드(분리된) 학습/훈련 방식(decoupled training scheme)을 비교한 결과를 보여준다는 뜻이에요.

그림 2.
우리 프레임워크의 개요입니다. 전체 과정은 3단계로 이루어져 있습니다.

1. 첫 번째 단계에서는 원본 데이터셋에 들어 있는 핵심 정보를 대부분 담아낼 수 있도록 모델을 처음부터(scratch) 학습합니다.

2. 두 번째 단계에서는 가우시안 노이즈(Gaussian noise)로부터 목표 데이터를 만들어내기 위해 복원(recovery) 과정을 수행하여 데이터를 합성(synthesize)합니다.

3. 세 번째 단계에서는 합성된 데이터의 라벨을 crop(잘라낸 조각) 단위로 다시 붙이는(crop-level relabeling) 방식을 적용해, 그 데이터의 실제(진짜) 라벨을 더 잘 반영하도록 합니다.

Table 1 (표 1)
Tiny-ImageNet(해상도 64×64)에서 모든 방법을 동일하게 RTX-4090 GPU 1장으로 돌렸을 때의 합성 시간(Synthesis Time)과 메모리 사용량(Memory Consumption)을 비교한 내용입니다.

• Time Cost는 합성 데이터에 대해 한 번의 반복(iteration) 업데이트를 하면서 이미지 1장을 생성할 때 드는 시간(밀리초, ms)을 의미합니다.
• GPU 메모리 사용량의 피크(최대) 값은 배치 크기 200 기준으로 측정했거나, 그 기준으로 환산(converted)한 값입니다.
• 괄호의 1 IPC는 “클래스당 이미지 수(Images Per Class)”를 뜻하고, Tiny-ImageNet은 클래스가 200개라서 배치 200을 쓴 설정과 연결된 설명이에요.

2. Approach (SRe2L)

SRe2L은 데이터 압축을 위한 3단계 학습 패러다임으로 구성되어 있으며, 합성 데이터 생성 단계를 실제 데이터 입력으로부터 분리하고, 내부 루프(inner-loop)와 외부 루프(outer-loop) 최적화를 분리하여 효율성을 극대화한다.

2.1. Stage-1: Squeeze ($\phi$)
이 단계의 목표는 원본 데이터셋의 핵심 정보를 딥 뉴럴 네트워크에 압축(extract)하고 캡슐화(encapsulate)하는 것이다. 이는 DeepDream [17] 및 Inverting Image [18, 19]의 아이디어에서 영감을 받았다. 이 단계는 단순히 원본 데이터셋에 대한 일반적인 모델 훈련 과정으로 볼 수 있다.

$\theta_T = \arg \min_\theta \mathcal{L}_T(\theta)$

여기서 $\mathcal{L}_T(\theta)$는 일반적으로 cross-entropy loss인 $E_{(x,y) \in T} [y \log (p(x))]$를 사용한다. 중요한 점은 과도한 데이터 증강(data augmentation)이 반드시 더 나은 복구(recovered) 모델로 이어지지 않는다는 것이다.
또한, SRe2L은 Batch Normalization (BN) 통계(statistics)를 사용하여 매칭(matching)을 수행하며, 이는 전체 데이터셋에 걸쳐 계산되므로 개별 배치(batch)에서 수행되는 Feature Matching보다 포괄적이고 대표적인 정렬(alignment)을 제공한다. ConvNet은 BN 레이어를 흔히 사용하지만, ViT(Vision Transformer)에는 BN이 없다. SRe2L은 ConvNet과 ViT 모두에 적용 가능하도록, ViT의 모든 LayerNorm을 BN 레이어로 대체하고 FFN(Feed-Forward Network)의 두 선형 레이어 사이에 추가 BN 레이어를 삽입하는 BN-ViT를 엔지니어링하였다.

2.2. Stage-2: Recover ($\psi$)
이 단계에서는 Squeeze 단계에서 보존된 정보를 이미지 공간으로 재구성(reconstruct)한다. 합성 데이터 $\mathbf{ex}_{syn}$는 고정된(frozen) 사전 훈련 모델 $\phi_{\theta_T}$ (Squeeze 단계에서 얻은 모델)을 사용하여 최적화된다. 이 과정은 class labels, regularization terms, 그리고 BN trajectory alignment를 활용한다. BN 통계와 예측 확률 분포의 쌍은 최적화 과정을 단일 레벨(singular level)로 제한하여 확장성을 크게 향상시킨다.
학습 목표는 다음과 같다.

$\arg \min_{C_{syn}, |C|} \mathcal{L}(\phi_{\theta_T}(\mathbf{ex}_{syn}), \mathbf{y}) + R_{reg}$

여기서 $R_{reg}$는 정규화 항이다. 저자들은 $\ell_2$ 및 Total Variation (TV)과 같은 이미지 사전 정규화(image prior regularizers) ($R_{prior}(\mathbf{ex}_{syn}) = \alpha_{tv}R_{TV}(\mathbf{ex}_{syn}) + \alpha_{\ell_2} R_{\ell_2}(\mathbf{ex}_{syn})$)의 효과를 평가했으나, 이미지의 시각적 품질보다는 의미론적 정보 회복에 중점을 두므로 이들을 생략하였다.
핵심 정규화 항은 BN 일관성(Consistency)이며, DeepInversion [20]에서와 유사하게 Feature Distribution Regularization의 속성을 활용한다.

$R_{BN}(\mathbf{ex}) = \sum_l \| \mu_l(\mathbf{ex}) - E (\mu_l | T ) \|^2 + \sum_l \| \sigma^2_l (\mathbf{ex}) - E (\sigma^2_l | T ) \|^2$

$\approx \sum_l \| \mu_l(\mathbf{ex}) - BNRM_l \|^2 + \sum_l \| \sigma^2_l (\mathbf{ex}) - BNRV_l \|^2$

여기서 $l$은 BN 레이어의 인덱스, $\mu_l(\mathbf{ex})$ 및 $\sigma^2_l (\mathbf{ex})$는 각각 평균과 분산, $BNRM_l$ 및 $BNRV_l$는 사전 훈련된 모델의 $l$번째 레이어에 있는 Running Mean 및 Running Variance이다. 또한, Multi-crop Optimization 전략을 도입하여 합성 데이터의 정보 함량을 향상시킨다. 이는 RandomResizedCrop과 유사하게 전체 이미지에서 무작위로 크롭(crop)하고, 이 크롭된 영역만 매 반복(iteration)마다 업데이트하는 방식이다.

2.3. Stage-3: Relabel ($\rho$)
이 단계는 합성 데이터와 실제 레이블의 정렬(alignment)을 목표로 한다. Recover 단계에서 생성된 합성 데이터에 대해 사전 훈련된 모델 $\phi_{\theta_T}$를 사용하여 소프트 레이블(soft label)을 생성한다 (FKD [22]와 유사).

$\mathbf{ey}_i = \phi_{\theta_T}(\mathbf{ex}_{R_i})$

여기서 $\mathbf{ex}_{R_i}$는 합성 이미지의 $i$번째 크롭(crop)이고 $\mathbf{ey}_i$는 해당 소프트 레이블이다. 마지막으로, 모델 $\phi_{C_{syn}}$을 이 소프트 레이블이 적용된 합성 데이터로 훈련시킨다.

$\mathcal{L}_{syn} = - \sum_i \mathbf{ey}_i \log \phi_{C_{syn}}(\mathbf{ex}_{R_i})$

이 단계는 합성 데이터와 레이블을 더욱 정렬하고 훈련된 모델의 성능을 크게 향상시키는 데 중요하다.

2.4. 계산 및 메모리 비용 절감:
SRe2L은 바이레벨 최적화를 두 단계(Squeezing 및 Recovering)로 분리하여 실제 데이터와 합성 데이터를 훈련 단계에서 디커플링한다. 이를 통해 각 훈련 세션에서 최소한의 메모리만 필요하게 된다. 기존 방법들은 실제 데이터와 합성 데이터를 동시에 GPU에 로드해야 하여 상당한 계산 및 메모리 오버헤드를 발생시켰다.

3. Experiments

SRe2L은 Tiny-ImageNet 및 ImageNet-1K 데이터셋에서 평가되었다. 백본 네트워크로는 ResNet-{18, 50, 101}, ViT-Tiny 및 BN-ViT-Tiny가 사용되었다.

3.1. Squeezing 분석:
Squeezing 예산(훈련 반복 횟수)이 증가하고 데이터 증강(Mixup, CutMix)을 적용할수록 복구된 모델의 성능이 감소하는 경향을 보였다. 이는 더 많이 훈련된 모델로부터 데이터를 복구하는 것이 더 어렵다는 것을 시사한다. Tiny-ImageNet의 경우 50 반복이 기본 설정으로 채택되었다.

3.2. Recovering 분석:

• Image Prior Regularization: $\ell_2$ 및 TV 정규화는 의미론적 정보 복구에 기여하지 않거나 오히려 방해가 될 수 있어 SRe2L에서는 생략되었다.
• Multi-crop Optimization: Multi-crop Optimization 전략은 validation 정확도를 크게 향상시켰다. 이는 합성 이미지의 크롭된 영역이 목표 카테고리와 더 밀접하게 연관되어 있고 소프트 레이블 모델 훈련에 더 유익함을 보여준다.
• Recover Budget: 더 긴 Recover 예산(반복 횟수)이 합성 데이터의 품질을 향상시켰다. ImageNet-1K의 경우 ResNet-50과 같은 대규모 모델에서 데이터 복구가 더 도전적이며, 비교 가능한 성능을 위해 더 많은 반복이 필요했다.

3.3. Relabeling 분석:

• Relabeling Model: Relabeling 모델이 Recover 모델 아키텍처와 동일하거나 유사할 때 가장 좋은 정확도를 얻었다. Relabeling 모델과 Recover 모델 간의 차이가 클수록 Top-1 에러가 증가하는 경향을 보였다.
• Temperature on Soft Label: 레이블 소프트맥스(softmax)를 위한 온도(temperature) 설정은 초기에는 Top-1 정확도를 빠르게 증가시키다가 10 이상에서는 안정화되었다. 최적의 온도는 20으로 설정되었다.
• Model Training: SRe2L로 압축된 데이터는 다양한 아키텍처(ResNet-{18, 50, 101})에 걸쳐 효과적인 훈련 확장성을 보여주었으며, 오버피팅 문제가 발생하지 않았다.

3.4. 압축된 데이터셋 평가:

• Tiny-ImageNet: 50 IPC(Images Per Class)에서 SRe2L은 42.5%의 정확도를 달성하여 기존 SOTA (MTT의 28.0%)를 크게 능가했다.
• ImageNet-1K: 10 IPC에서 SRe2L은 21.3%의 정확도를 달성하여 TESLA의 7.7%를 크게 능가했다. 50 IPC에서 SRe2L은 60.8%의 정확도를 달성하여 모든 기존 방법을 큰 차이로 앞섰다. SRe2L은 대규모 아키텍처에서 성능 향상을 보였다.

3.5. 교차 아키텍처 일반화 (Cross-Architecture Generalization):
SRe2L로 압축된 데이터셋은 ResNet-{18, 50, 101} 및 ViT-T와 같은 새로운 아키텍처에 효과적으로 일반화(generalize)되는 능력을 입증했다.

3.6. 합성 이미지 시각화:
SRe2L이 생성한 합성 이미지는 MTT [1]의 이미지보다 의미론적으로 더 선명하고 대상 클래스의 속성과 윤곽을 효과적으로 담고 있었다. 이는 SRe2L이 더 높은 품질의 이미지를 생성하며, 이는 더 많은 의미론적 정보를 포함하여 validation 정확도를 높이는 데 기여한다.

3.7. 응용: Continual Learning:
SRe2L은 Tiny-ImageNet의 class-incremental learning에서 기존 baseline보다 우수한 성능을 보이며, 대규모 데이터 처리에서의 우수성을 입증했다.

Table 2 (표 2)
Tiny-ImageNet에서 squeezing budget(데이터를 “압축/추려내는” 데 쓰는 예산·제약, 예: 합성 이미지 수/업데이트 횟수/연산량 같은 리소스 한도)과 데이터 증강(data augmentation)이 성능에 미치는 영향을 요소별로 떼어(ablations) 검증한 실험 결과라는 뜻이에요.

Table 3 (표 3)
ResNet-{18, 50}을 복원(recovery) 모델로 사용하고 업데이트 반복(iteration) 횟수를 달리했을 때의 ablation(요소 제거/변경 실험)에서의 Top-1 검증 정확도이며, 학생(student) 모델로는 ResNet-18을 사용했다. “Time”은 단일 NVIDIA 4090 GPU에서 반복 1회(iteration)마다 이미지 1장을 학습할 때 소요되는 시간(ms)을 나타낸다.

그림 3.
다양한 정규화 항과 크롭 증강 설정 하에서 ImageNet-1K에 대한 증류된 예시들의 시각화.
선택된 클래스는 {Volcano, Hammerhead Shark, Bee, Valley}이다.

그림 4.
IPC 50 조건에서, 서로 다른 라벨 종류와 temperature 설정(보통 소프트 라벨/지식증류에서 쓰는 온도 파라미터)을 바꿔가며 학습한 모델들의 Top-1 검증 정확도를 나타낸다.
여기서 T와 S는 각각 리라벨링(relabeling)을 위한 기준(reference) 모델과 학습 대상(target) 모델을 의미한다.
또한 R18, R50, R101은 각각 ResNet-18, ResNet-50, ResNet-101의 약어이다.

Table 4 (표 4)
IPC 50 조건에서, 서로 다른 라벨 종류와 temperature 설정(보통 소프트 라벨/지식증류에서 쓰는 온도 파라미터)을 바꿔가며 학습한 모델들의 Top-1 검증 정확도를 나타낸다.
여기서 T와 S는 각각 리라벨링(relabeling)을 위한 기준(reference) 모델과 학습 대상(target) 모델을 의미한다.
또한 R18, R50, R101은 각각 ResNet-18, ResNet-50, ResNet-101의 약어이다.

Table 5 (표 5)
ImageNet-1K에서 교차 아키텍처 일반화(Cross-Architecture Generalization) 성능을 평가한 Top-1 정확도 결과이다.
두 가지 복원/압축(recovery/squeezed) 모델을 사용했으며: DeiT-Tiny-BN과 ResNet-18이다.
평가에는 네 가지 평가 모델을 사용했는데: DeiT-Tiny, ResNet-18, ResNet-50, ResNet-101이다.

ImageNet-1K
Table 4의 두 번째 그룹에서 볼 수 있듯이, IPC 10에서 ResNet-18과 동일한 모델 아키텍처를 사용할 때, 우리의 접근 방식은 TESLA의 성능을 7.7%의 기준선에서 21.3%로 향상시킵니다. 성능이 더 큰 모델 아키텍처에서 저하되는 TESLA와 달리, 우리가 제안하는 접근 방식은 더 큰 아키텍처를 활용하여 주목할 만한 비례적인 성능 향상을 보여줍니다. 이는 대규모 모델의 현대 시대에 상당한 가능성을 시사합니다. IPC 50, 100, 200에서는 우리의 방법이 일관된 정확도 향상을 달성합니다.

그림 5.
MTT [1] 및 SRe2L 시각화.
윗 두 행은 합성 Tiny-ImageNet이며, 아래 두 행은 합성 ImageNet-1K입니다. (첫 번째 행은 MTT, 두 번째 행은 저희의 결과입니다).

그림 6.
Tiny-ImageNet에서의 5단계 및 10단계 클래스 증분 학습

4. Related Work

데이터 압축 분야는 1) Meta-Model Matching, 2) Gradient Matching, 3) Distribution Matching, 4) Trajectory Matching 의 네 가지 주요 범주로 나뉜다. SRe2L의 BN-matching 복구 절차는 Distribution Matching 방식의 특별한 형태로 볼 수 있다.

5. Conclusion

SRe2L은 Squeezing, Recovering, Relabeling의 3단계 접근 방식을 통해 대규모 ImageNet-1K 데이터셋을 효율적이고 효과적으로 압축하는 새로운 방법을 제시한다. 이 방법은 기존 SOTA 압축 방법을 큰 폭으로 능가하며, 계산 효율성과 모델 성능 모두에서 향상을 보여준다.

제한 사항 및 향후 연구:
현재 압축된 데이터셋과 원본 전체 데이터셋 간에는 여전히 성능 차이가 존재하며, 전체 데이터의 완전한 대체는 아직 불가능하다. 또한 소프트 레이블에 대한 추가 저장 공간이 필요하다. 향후 연구는 ImageNet-21K와 같은 더 큰 데이터셋, 그리고 언어 및 음성과 같은 다른 데이터 양식으로의 확장에 중점을 둘 것이다.