Improving Diversity in Black-box Few-shot Knowledge Distillation

Tri-Nhan Vo 외의 논문 "Improving Diversity in Black-box Few-shot Knowledge Distillation"에 대한 상세한 기술 요약입니다.

1. 문제 정의

본 논문은 **블랙박스 퓨샷 지식 증류 (Black-box Few-shot Knowledge Distillation, KD)**의 과제를 다룹니다.

• 배경: 표준 KD 는 대규모 "교사" 모델에서 소형 "학생" 모델로 지식을 전이합니다. 그러나 전통적인 KD 는 대규모 학습 데이터셋과 교사 모델의 내부 파라미터에 대한 완전한 접근 (화이트박스) 을 필요로 합니다.
• 제약 조건: 실제 시나리오 (예: Google 의 FaceNet 또는 ChatGPT 와 같은 독점 모델을 증류하는 경우) 에서는 학생 모델 개발자가 다음과 같은 제한을 겪습니다:
  1. 블랙박스 접근: 교사의 최종 예측 확률만 이용 가능하며, 내부 기울기, 로그 (logits), 또는 파라미터는 숨겨져 있습니다.
  2. 퓨샷 데이터: 학생 모델 학습을 위해 사용 가능한 레이블이 지정된 이미지가 매우 적어 ($N$), 실제 데이터 매니폴드를 대표하기에 부족합니다.
• 기존 해결책의 한계: 현재 퓨샷 KD 방법들은 대부분 화이트박스 교사에 의존하거나 MixUp 과 같은 단순한 데이터 증기를 사용하는데, 이는 다양성이 부족하거나 현실적인 의미론을 결여한 합성 이미지를 생성합니다. 변이형 오토인코더 (CVAE) 를 사용하는 다른 방법들은 흐림 아티팩트 (blurring artifacts) 로 고통받으며 다양성을 명시적으로 촉진하지 못합니다.

2. 방법론: DivBFKD

저자들은 교사의 내부에 접근하지 않고도 교사의 원래 학습 분포를 모방하는 다양하고 고품질의 합성 이미지를 생성하도록 설계된 2 단계 프레임워크인 DivBFKD(Diverse Black-box Few-Shot Knowledge Distillation) 를 제안합니다.

1 단계: 생성 (높은 신뢰도 선택을 통한 WGAN)

핵심 혁신은 다양하고 고품질의 합성 이미지를 생성하도록 **워asserstein 생성 적대 신경망 (WGAN)**을 학습시키는 데 있습니다.

• 높은 신뢰도 이미지: 모든 생성된 이미지를 사용하는 대신, 이 방법은 교사의 신뢰도를 기반으로 합성 이미지를 필터링합니다. 교사의 최대 예측 확률 $c_{\tilde{x}}$가 임계값 $\tau$를 초과할 때 이미지 $\tilde{x}$를 "높은 신뢰도"로 간주합니다.
• 적응형 임계값 ($\tau$): 고정된 임계값은 교사가 클래스별 편향 (예: "셔츠"보다 "바지"를 더 높은 신뢰도로 예측하는 등) 을 보이기 때문에 비효율적입니다. 저자들은 퓨샷 실제 이미지의 신뢰도 점수에서 $q$-분위수 (예: $q=0.1$) 를 사용하여 클래스별로 계산된 적응형 임계값을 제안합니다. 이는 임계값이 각 특정 클래스에 대한 교사의 행동에 상대적이도록 보장합니다.
• 온더플라이 적대적 학습:
  • WGAN 생성기 ($G$) 는 합성 이미지를 생성합니다.
  • 판별기 ($D$) 는 실제 이미지(퓨샷 세트 $D$에서) 와 높은 신뢰도 합성 이미지(이전 반복에서 선택됨) 를 구별하도록 학습됩니다.
  • 피드백 루프: 판별기에 대해 높은 신뢰도 합성 이미지를 "실제"로 취급함으로써, 생성기는 교사가 신뢰도 있게 분류하는 더 다양한 이미지를 생성하도록 동기를 부여받습니다. 이는 생성된 세트의 다양성을 확장하는 긍정적 피드백 루프를 만듭니다.

2 단계: 증류

• 데이터셋 구성: 증류 데이터셋 ($D_{KD}$) 은 원래 퓨샷 실제 이미지 ($N$) 와 학습된 WGAN 이 생성한 합성 이미지 예산 ($M$) 을 결합하여 구성됩니다.
• 클래스 균형: 편향을 방지하기 위해 합성 이미지는 거부 샘플링을 통한 클래스 균형을 유지하며 생성됩니다.
• 학습: 학생 네트워크는 교사의 예측 (소프트 레이블) 에 대한 표준 교차 엔트로피 손실을 사용하여 $D_{KD}$에서 학습됩니다.

3. 주요 기여

1. WGAN 을 위한 새로운 학습 방식: 높은 신뢰도 이미지를 적대적 학습 루프에 온더플라이로 주입하는 메커니즘을 도입합니다. 이는 교사를 지도자로 활용하여 생성기를 교사의 알려지지 않은 학습 분포로 안내하는 효과적인 방법입니다.
2. 적응형 임계값: 교사의 편향을 완화하고 비정상적인 이미지 선택을 방지하기 위해 분위수 기반 임계값을 사용하여 "높은 신뢰도"를 정의하는 클래스별, 데이터셋 무관한 방법을 제안합니다.
3. 다양한 블랙박스 퓨샷 KD (DivBFKD): 퓨샷 KD 의 중요한 병목 현상인 합성 데이터의 다양성 부족을 해결하는 완전한 프레임워크입니다.
4. 최첨단 성능: 다양한 데이터셋과 아키텍처에 걸쳐 기존 퓨샷 KD 방법보다 상당한 개선을 입증했습니다.

4. 실험 결과

이 방법은 다양한 아키텍처 (LeNet, AlexNet, ResNet, VGG) 를 사용하여 7 개의 데이터셋(MNIST, FMNIST, SVHN, CIFAR10, CIFAR100, Tiny-ImageNet, Imagenette) 에서 평가되었습니다.

• 성능 향상:
  • 단순 데이터셋: DivBFKD 는 FMNIST 와 CIFAR10 에서 두 번째로 좋은 블랙박스 방법 (FS-BBT) 보다 약 2% 더 우수한 성능을 보였습니다. CIFAR10 에서 표준 KD 보다 17% 향상되었습니다.
  • 복잡한 데이터셋: CIFAR100 에서 DivBFKD 는 **59.91%**의 정확도를 달성하여 FS-BBT(56.28%) 를 약 3.6% 상회했습니다. Imagenette(고해상도) 에서는 **86.48%**에 도달하여 "Student-Full" 상한선과의 격차를 크게 좁혔습니다.
  • 크로스 아키텍처: 이 방법은 교사와 학생이 서로 다른 아키텍처를 가질 때 (예: ResNet 교사에서 VGG 학생) 도 견고하게 유지되었으며, 표준 KD 대비 정확도를 14~22% 향상시켰습니다.
• 다양성 및 품질 지표:
  • 커버리지: 합성 이미지는 교사의 학습 이미지 이웃의 **42%**를 커버한 반면, 퓨샷 이미지는 18%, 표준 WGAN 출력은 29% 만 커버했습니다.
  • FID/IS: 베이스라인 중 가장 우수한 프리쳇 시작 거리 (13.96) 와 시작 점수 (4.42) 를 달성하여 높은 다양성과 품질을 나타냈습니다.
• 애블레이션 연구: **적응형 임계값 (AT)**과 **클래스 균형 (CB)**이 모두 필수적인 구성 요소임을 확인했으며, AT 가 다양성 향상에 가장 큰 기여를 했습니다.
• 데이터 프리 KD 와의 비교: DivBFKD 는 데이터 프리 (제로샷) KD 방법보다 현저히 우수한 성능을 보여주어, 소량의 실제 데이터 (2K 이미지) 와 다양한 합성 데이터를 결합하는 것이 실제 데이터가 전혀 없는 경우보다 우월함을 입증했습니다.

5. 의의

본 논문은 모델 압축 및 개인정보 보호 AI 분야에서 중요한 기여를 합니다:

• 실용적 적용 가능성: 조직이 개인정보 또는 지식재산권 문제로 인해 대규모 데이터셋이나 모델 가중치를 공유할 수 없는 실제 세계의 중요한 문제를 해결합니다.
• 다양성 중심: 단순히 어떤 합성 데이터를 생성하는 것에서 다양한 그리고 교사 정렬된 데이터를 생성하는 것으로 초점을 이동시켜, 데이터가 부족한 환경에서 효과적인 증류의 핵심을 다룹니다.
• 효율성: 제안된 방법은 단순한 임계값 설정과 온더플라이 선택을 통해 계산 오버헤드를 거의 추가하지 않으면서 최첨단 결과를 제공합니다.
• 일반화 가능성: 이 접근법은 아키텍처 무관하며 단순하고 복잡한 이미지 분류 작업 전반에 걸쳐 작동하므로, 제한된 환경에서 경량 모델을 배포하기 위한 다재다능한 도구가 됩니다.