[논문 리뷰] Scaling Up Dataset Distillation to ImageNet-1K with Constant Memory

Justin Cui, Ruochen Wang외 2명|arXiv (Cornell University)|2022. 11. 19.

Multimodal Machine Learning Applications인용 수 14

한 줄 요약

이 논문은 TESLA를 제안하며, 경량 메모리 트레이젝터리 매칭 방법을 통해 데이터셋 정련에서 GPU 메모리 복잡도를 O(T)에서 O(1)로 감소시켜 ImageNet-1K에까지 확장 가능하게 한다. 이는 클래스당 최대 50장의 이미지로 가능하다. 또한, 교사 모델 출력을 사용해 소프트 레이블 할당(SLA)을 도입하여 대규모 클래스 데이터셋에서 수렴성을 향상시켰으며, 50 IPC에서 ImageNet-1K에서 27.9%의 top-1 정확도를 달성했다. 이는 이전 최고 성능(SOTA) 대비 18.2%의 절대적 향상이며, 원본 데이터의 4.2%만을 사용한다.

ABSTRACT

Dataset Distillation is a newly emerging area that aims to distill large datasets into much smaller and highly informative synthetic ones to accelerate training and reduce storage. Among various dataset distillation methods, trajectory-matching-based methods (MTT) have achieved SOTA performance in many tasks, e.g., on CIFAR-10/100. However, due to exorbitant memory consumption when unrolling optimization through SGD steps, MTT fails to scale to large-scale datasets such as ImageNet-1K. Can we scale this SOTA method to ImageNet-1K and does its effectiveness on CIFAR transfer to ImageNet-1K? To answer these questions, we first propose a procedure to exactly compute the unrolled gradient with constant memory complexity, which allows us to scale MTT to ImageNet-1K seamlessly with ~6x reduction in memory footprint. We further discover that it is challenging for MTT to handle datasets with a large number of classes, and propose a novel soft label assignment that drastically improves its convergence. The resulting algorithm sets new SOTA on ImageNet-1K: we can scale up to 50 IPCs (Image Per Class) on ImageNet-1K on a single GPU (all previous methods can only scale to 2 IPCs on ImageNet-1K), leading to the best accuracy (only 5.9% accuracy drop against full dataset training) while utilizing only 4.2% of the number of data points - an 18.2% absolute gain over prior SOTA. Our code is available at https://github.com/justincui03/tesla

연구 동기 및 목표

• ImageNet-1K과 같은 대규모 데이터셋으로 확장할 때 트레이젝터리 매칭 기반 데이터셋 정련(MTT)의 금전적인 메모리 비용을 극복하기 위해.
• MTT가 CIFAR-10/100에서 뛰어난 성능을 보이는 것이 ImageNet-1K와 같이 훨씬 더 많은 클래스를 가진 대규모 데이터셋으로 일반화되는지 조사하기 위해.
• 합성 샘플에 하드 레이블을 사용할 경우 MTT가 ImageNet-1K에서 수렴하지 못하는 문제를 해결하기 위해.
• 합성 샘플에 정보성 레이블 신호를 주입하기 위한 초모수 없는, 학습이 필요 없는 방법을 개발하기 위해.
• 기존 방법보다 훨씬 높은 압축률을 달성하면서 ImageNet-1K에서 최고 성능을 내는 것을 목표로 한다.

제안 방법

• T는 편향된 SGD 단계 수일 때, 메모리 복잡도를 O(T)에서 O(1)로 감소시키기 위해 트레이젝터리 매칭 기울기의 캐싱 및 재배열을 수행하는 새로운 기울기 계산 절차를 제안한다.
• 소프트 레이블 할당(SLA)을 도입하여, 사전 훈련된 교사 모델의 로짓을 사용해 합성 이미지에 대해 소프트 레이블을 생성하며, 레이블을 함께 학습하지 않는다.
• 다양한 훈련 에포크에서 얻은 교사 모델 앙상블을 활용해, 합성 이미지 최적화 중 동적이고 고품질의 소프트 레이블을 제공한다.
• 일정 메모리 MTT의 구현과 SLA를 결합하여, 단일 GPU에서 최대 50 IPC로 ImageNet-1K에서의 엔드 투 엔드 훈련을 가능하게 한다.
• SLA에서 참조 모델 가중치 외에 추가 초모수를 사용하지 않아 기존 MTT 프레임워크와 호환 가능하다.
• 학습이 필요 없는, 추론 기반의 레이블 주입 전략을 사용하여 공동 레이블 최적화의 불안정성을 피한다.

실험 결과

연구 질문

• RQ1 일괄 최적화로 인해 높은 메모리 비용이 발생하는 트레이젝터리 매칭 기반 데이터셋 정련(MTT)이 ImageNet-1K로 확장 가능한가?
• RQ2 CIFAR-10/100과 같은 소규모 데이터셋에서 MTT의 뛰어난 성능이 ImageNet-1K와 같은 대규모 데이터셋으로 일반화되는가?
• RQ3 합성 샘플에 하드 레이블을 사용할 경우 MTT가 왜 ImageNet-1K에서 효과적으로 수렴하지 못하는가?
• RQ4 교사 모델에서 유도된 소프트 레이블 할당이 추가 초모수 없이도 대규모 클래스 데이터셋에서 MTT의 성능을 크게 향상시킬 수 있는가?
• RQ5 단일 GPU와 일정 메모리로 50 IPC에서 ImageNet-1K에서 SOTA 성능을 달성하는 것이 가능한가?

주요 결과

• TESLA는 계산 오버헤드를 거의 느끼지 않으면서 MTT의 메모리 복잡도를 O(T)에서 O(1)로 감소시켜 ImageNet-1K로의 확장 가능성을 열었다.
• TESLA는 50 IPC에서 ImageNet-1K에서 27.9%의 top-1 정확도를 달성했으며, 이는 전체 데이터셋 훈련 대비 5.9%의 감소에 불과하지만, 원본 데이터의 4.2%만을 사용한다.
• TESLA는 ImageNet-1K에서 50 IPC로 MTT를 확장한 최초의 방법으로, 이전 SOTA 대비 18.2%의 절대 정확도 향상을 달성했다.
• 1 IPC에서 소프트 레이블 할당(SLA)은 하드 레이블 대비 성능을 7.0%포인트 향상시켰으며, 추가 초모수 없이도 가능했다.
• TESLA에서 정련된 데이터셋은 아키텍처 간에 잘 일반화되며, CIFAR-10, CIFAR-100, ImageNet-1K에서 훈련했을 때 기존 방법보다 ResNet18과 ViT에서 모두 뛰어난 성능을 보였다.
• 절단 연구 결과, SLA만으로도 1 IPC에서 하드 레이블일 경우 0.7%에서 소프트 레이블일 경우 7.7%로 성능 향상됨을 확인했으며, 최종 에포크보다는 초기 훈련 에포크에서 유도된 소프트 레이블이 더 효과적임을 밝혔다.