[논문 리뷰] Prospect Pruning: Finding Trainable Weights at Initialization using Meta-Gradients

Milad Alizadeh, Shyam A. Tailor외 5명|arXiv (Cornell University)|2022. 02. 16.

Advanced Neural Network Applications인용 수 11

한 줄 요약

Prospect Pruning (ProsPr)는 초기화 시점에서의 프루닝을 위한 새로운 방법으로, 첫 번째 학습 단계 수개에 걸친 메타기울기를 사용하여 프루닝 이후에도 학습 가능한 하위망을 식별한다. 이는 프루닝 이후의 학습 가능성에 명시적으로 고려함으로써, 단일 스텝 프루닝과 더 적은 데이터 요구량으로도 최신 기술 수준의 정확도를 달성한다. 이는 최적화 동역학을 중요도 기준에 통합하여 기존 방법보다 뛰어난 성능을 발휘한다.

ABSTRACT

Pruning neural networks at initialization would enable us to find sparse models that retain the accuracy of the original network while consuming fewer computational resources for training and inference. However, current methods are insufficient to enable this optimization and lead to a large degradation in model performance. In this paper, we identify a fundamental limitation in the formulation of current methods, namely that their saliency criteria look at a single step at the start of training without taking into account the trainability of the network. While pruning iteratively and gradually has been shown to improve pruning performance, explicit consideration of the training stage that will immediately follow pruning has so far been absent from the computation of the saliency criterion. To overcome the short-sightedness of existing methods, we propose Prospect Pruning (ProsPr), which uses meta-gradients through the first few steps of optimization to determine which weights to prune. ProsPr combines an estimate of the higher-order effects of pruning on the loss and the optimization trajectory to identify the trainable sub-network. Our method achieves state-of-the-art pruning performance on a variety of vision classification tasks, with less data and in a single shot compared to existing pruning-at-initialization methods.

연구 동기 및 목표

• 기존의 초기화 시점 프루닝 방법들이 프루닝 이후의 학습 가능성에 고려하지 못하는 문제를 해결하기 위해.
• 현재의 중요도 기준이 초기화 시점에서 단일 단계만 평가할 뿐 이후 학습을 고려하지 않는 단기적 사고 방식을 극복하기 위해.
• 프루닝이 첫 번째 학습 단계 동안 최적화 궤적과 손실에 미치는 영향을 명시적으로 모델링하는 방법을 개발하기 위해.
• 반복적 재학습이나 점진적 프루닝 없이도 초기화 시점에서 고정밀도의 단일 스텝 프루닝을 가능하게 하기 위해.
• 학습 가능한 하위망을 초기 단계부터 조밀하고도 학습 가능한 상태로 식별함으로써 신경망 프루닝의 실용성과 효율성을 향상시키기 위해.

제안 방법

• ProsPr는 초기 가중치에 대한 손실의 기울기인 메타기울기를 기반으로 중요도 점수를 계산한다. 이 메타기울기는 수개의 최적화 단계 이후에 계산된다.
• 확률적 경사 하강법의 첫 번째 수개 단계를 통해 역전파하여, 특정 가중치를 프루닝할 경우 초기 학습 단계 동안 손실에 어떤 영향을 미치는지 추정한다.
• 계산 효율성과 기울기 안정성을 향상시키기 위해 메타기울기의 일阶 근사(approximation)를 사용한다.
• 모든 가중치에 대해 중요도 점수를 계산하고, 이 점수에 기반해 가장 중요도가 낮은 가중치들을 프루닝한다.
• 프루닝 마스크는 초기 가중치에 직접 적용되며, 나머지 하위망은 처음부터 재학습한다.
• 이 방법은 반복적 프루닝과 재학습 루프를 피할 수 있도록 설계되어 있어, 이전 방법에서 흔히 볼 수 있는 단계적 프로세스를 피한다.

실험 결과

연구 질문

• RQ1 프루닝 이후 최적화 궤적을 고려하는 초기화 시점 프루닝 방법이 기존 방법보다 더 뛰어난 성능을 낼 수 있는가?
• RQ2 초기 학습 단계에 걸친 메타기울기를 통합함으로써 정적 중요도 기준에 비해 학습 가능한 하위망을 얼마나 더 잘 식별할 수 있는가?
• RQ3 ProsPr는 단일 프루닝 단계와 더 적은 데이터 사용으로 최신 기술 수준의 정확도를 얼마나 잘 달성할 수 있는가?
• RQ4 메타기울기를 통한 학습 가능성 모델링이 반복적 또는 점진적 프루닝 전략의 필요성을 줄이는가?
• RQ5 메타기울기가 손실과 최적화 동역학에 대한 프루닝의 고차원 효과를 효과적으로 포착할 수 있는가?

주요 결과

• ProsPr는 다양한 비전 분류 벤치마크에서 최신 기술 수준의 정확도를 달성하며, 기존의 초기화 시점 프루닝 방법을 능가한다.
• 이 방법은 반복적 재학습이나 점진적 프루닝 없이도 단일 스텝 프루닝 과정을 통해 높은 정확도를 달성한다.
• 고밀도 프루닝 수준에서도 이전 방법 대비 성능 저하가 적게 나타나며, 성능 저하를 줄였다.
• 메타기울기의 사용으로 프루닝이 초기 학습 동역학에 미치는 영향을 모델링함으로써 학습 가능한 하위망을 더 잘 식별할 수 있었다.
• 이전 방법보다 더 적은 데이터로도 효과적으로 작동하여, 하위망 탐색에서 더 뛰어난 샘플 효율성을 보였다.
• 세 단계를 초과해 메타기울기를 계산할 경우 기울기 안정성 문제가 발생하여, 이는 메서드를 소수의 초기 단계로 제한한다.