[논문 리뷰] InfoBatch: Lossless Training Speed Up by Unbiased Dynamic Data Pruning

Ziheng Qin, Kai Wang외 9명|arXiv (Cornell University)|2023. 03. 08.

Advanced Neural Network Applications인용 수 7

한 줄 요약

InfoBatch는 손실이 낮은(잘 학습된) 샘플을 무작위로 제거하고 기울기 재스케일링을 통해 원래 기울기 기대값을 유지함으로써 성능 손실 없이 딥러닝 학습을 가속화하는 새로운 비편향 동적 데이터 프루닝 프레임워크를 제안한다. 이는 이미지 분류, 디퓨전 모델, LLaMA 지시 fine-tuning에서 각각 최대 40%, 27%, 20%의 손실 없는 학습 속도 향상을 달성하며, 최소한의 오버헤드와 시각 및 자연어 처리 작업 간 호환성을 보인다.

ABSTRACT

Data pruning aims to obtain lossless performances with less overall cost. A common approach is to filter out samples that make less contribution to the training. This could lead to gradient expectation bias compared to the original data. To solve this problem, we propose extbf{InfoBatch}, a novel framework aiming to achieve lossless training acceleration by unbiased dynamic data pruning. Specifically, InfoBatch randomly prunes a portion of less informative samples based on the loss distribution and rescales the gradients of the remaining samples to approximate the original gradient. As a plug-and-play and architecture-agnostic framework, InfoBatch consistently obtains lossless training results on classification, semantic segmentation, vision pertaining, and instruction fine-tuning tasks. On CIFAR10/100, ImageNet-1K, and ADE20K, InfoBatch losslessly saves 40\% overall cost. For pertaining MAE and diffusion model, InfoBatch can respectively save 24.8\% and 27\% cost. For LLaMA instruction fine-tuning, InfoBatch is also able to save 20\% cost and is compatible with coreset selection methods. The code is publicly available at \href{https://github.com/henryqin1997/InfoBatch}{github.com/NUS-HPC-AI-Lab/InfoBatch}.

연구 동기 및 목표

• 대규모 데이터셋에서 대규모 딥러닝 모델을 학습하는 데 드는 높은 계산 비용을 해결하기 위해.
• 정적 및 동적 데이터 프루닝 방법이 샘플을 제거할 때 보정 없이 기울기 기대값 편향을 유발하는 문제를 해결하기 위해.
• 손실 없는 학습 속도 향상을 가능하게 하는 즉시 사용 가능한, 아키텍처에 관계없는 프레임워크를 개발하기 위해.
• 기존 데이터 압축 기법(예: 코어셋 선택 및 데이터셋 양자화)과의 호환성을 확보하기 위해.

제안 방법

• InfoBatch는 순환적으로 각 샘플의 손실 점수를 포워드 프로파게이션 동안 유지하여 손실이 낮은(잘 학습된) 샘플을 식별한다.
• 각 학습 에포크에서, 가장 낮은 손실을 가진 샘플의 고정 비율을 무작위로 프루닝하며, 정렬 또는 히우리스틱 점수 계산 오버헤드를 피한다.
• 남아 있는 샘플의 기울기를 재스케일링하여 총 기울기 업데이트의 기대값을 유지함으로써 편향 없는 최적화를 보장한다.
• 기존 데이터셋에서 전체 배치 학습과 동일한 기울기 기대값 방향을 유지하기 위해 기대값 재스케일링을 사용한다.
• 학습 중 실시간 손실 값에 기반해 동적으로 프루닝을 조정함으로써 온라인, 저오버헤드 적응이 가능하다.
• 최종 몇 에포크 동안은 전체 데이터셋 학습을 복원하여 수렴 안정성 향상과 분산 감소를 달성한다.

실험 결과

연구 질문

• RQ1 동적 데이터 프루닝이 기울기 기대값 편향을 유발하지 않고도 손실 없는 학습 가속화를 달성할 수 있는가?
• RQ2 간단하고 저오버헤드 프루닝 전략이 기존 히우리스틱 또는 반복적 방법보다 속도 향상과 정확도 측면에서 뛰어나게 작용할 수 있는가?
• RQ3 샘플 제거에도 불구하고 기울기 재스케일링이 원래 최적화 궤적을 효과적으로 유지하는가?
• RQ4 InfoBatch는 아키텍처 수정 없이도 다양한 딥러닝 작업(시각 및 자연어 처리 포함)에 일반적으로 적용 가능한가?
• RQ5 InfoBatch는 학습 비용 절감 및 계산 오버헤드 측면에서 최신 기술과 비교해 어떻게 성능을 내는가?

주요 결과

• InfoBatch는 CIFAR-10/100, ImageNet-1K, ADE20K 데이터셋에서 전체 배치 학습과 동일한 정확도를 유지하면서 전체 학습 비용을 40% 감소시켰다.
• 마스킹 자동에코딩(MAE) 사전학습에서는 비용을 24.8% 절감했으며, 이미지 복원 품질에 하락이 없었다.
• 디퓨전 모델 사전학습에서는 생성된 이미지 품질에 영향을 주지 않고 학습 비용을 27% 절감했다.
• LLaMA-7B의 지시 파인튜닝에서 InfoBatch는 학습 비용을 20% 줄였으며, BBH, DROP, MMLU, HumanEval 벤치마크에서 최신 기술 수준의 성능을 유지했다.
• EL2N과 같은 최신 기술(샘플 점수 추정에 O(MNT) 시간이 소요됨)에 비해 샘플 점수 계산 오버헤드를 최소 10배 이상 감소시켰다.
• 이 프레임워크는 코어셋 선택 및 데이터셋 양자화와 완전히 호환되어 기존 압축 기법과 조합하여 추가 비용 절감이 가능하다.