[논문 리뷰] Only Train Once: A One-Shot Neural Network Training And Pruning Framework

Tianyi Chen, Bo Ji외 8명|arXiv (Cornell University)|2021. 07. 15.

Advanced Neural Network Applications참고 문헌 81인용 수 43

한 줄 요약

OTO는 제로-불변 그룹을 사용하고 새로운 Half-Space Stochastic Projected Gradient를 통해 미세조정 없이 한 번의 패스로 처음부터 전체 신경망을 학습하고 압축합니다.

ABSTRACT

Structured pruning is a commonly used technique in deploying deep neural networks (DNNs) onto resource-constrained devices. However, the existing pruning methods are usually heuristic, task-specified, and require an extra fine-tuning procedure. To overcome these limitations, we propose a framework that compresses DNNs into slimmer architectures with competitive performances and significant FLOPs reductions by Only-Train-Once (OTO). OTO contains two keys: (i) we partition the parameters of DNNs into zero-invariant groups, enabling us to prune zero groups without affecting the output; and (ii) to promote zero groups, we then formulate a structured-sparsity optimization problem and propose a novel optimization algorithm, Half-Space Stochastic Projected Gradient (HSPG), to solve it, which outperforms the standard proximal methods on group sparsity exploration and maintains comparable convergence. To demonstrate the effectiveness of OTO, we train and compress full models simultaneously from scratch without fine-tuning for inference speedup and parameter reduction, and achieve state-of-the-art results on VGG16 for CIFAR10, ResNet50 for CIFAR10 and Bert for SQuAD and competitive result on ResNet50 for ImageNet. The source code is available at https://github.com/tianyic/only_train_once.

연구 동기 및 목표

• pruning 워크플로에서 다단계 미세조정 없이 효율적인 모델 압축의 필요성에 대해 동기를 부여한다.
• 처음부터 전체 모델을 동시에 학습하고 가지치기할 수 있는 프레임워크를 소개한다.
• 안전한 출력 보존 가지치를 가능하게 하는 제로-불변 그룹(ZIGs)을 정의한다.
• 구조화된 희소성 정규화를 통해 그룹 희소성을 유도하기 위한 새로운 최적화 알고리즘 Half-Space Stochastic Projected Gradient(HSPG)을 개발한다.
• CNN 및 NLP 모델 전반에서 표준 벤치마크에서의 최첨단 또는 경쟁적인 압축 및 정확도를 입증한다.

제안 방법

• 모델 파라미터를 제로-불변 그룹(ZIGs)으로 분할하여 제로 그룹이 출력에 영향을 주지 않고 가지치기될 수 있도록 한다.
• ZIGs에 대한 혼합 l1/l2 노름으로 구조적 희소성 규제를 형성하여 그룹 희소성을 유도한다.
• 초비연속 규제 문제를 Half-Space Stochastic Projected Gradient(HSPG)로 해결하는데, 이는 초기화 단계(SGD)와 그룹-희소성 단계(Half-Space Projection)로 구성된다.
• 하강 조건이 충족되면 전체 그룹을 0으로 만드는 Half-Space Projection 연산자를 도입하여 공격적인 희소성을 가능하게 한다.
• 추가 미세조정 없이 제로 ZIG를 제거하여 더 얇은 아키텍처로 전체 모델을 가지치기한다.
• 일반적인 아키텍처(Conv-BN, 잔차 블록, Fully Connected 및 Multi-Head Attention)에 대한 자동 또는 사전 지정을 위한 ZIG 분할 가이드를 제공한다.

실험 결과

연구 질문

• RQ1 처음부터 한 번의 학습 실행이 미세조정 없이도 압축된 고성능 모델을 낼 수 있는가?
• RQ2 파라미터를 ZIG로 분할하여 출력 변경 없이 제로 값인 그룹을 가지치기할 수 있는가?
• RQ3 HSPG를 이용한 구조화된 희소성 접근법이 표준 인접 방법보다 제로 그룹을 더 효과적으로 촉진하는가?
• RQ4 OTO를 적용할 때 표준 벤치마크에서 FLOPs, 파라미터 수 및 정확도에 어떤 실질적 이득이 있는가?
• RQ5 이 접근법이 아키텍처에 독립적이며 CNN과 BERT와 같은 NLP 모델에 확장 가능한가?

주요 결과

• OTO는 벤치마크 전반에서 경쟁력 있는 정확도를 유지하면서 FLOPs와 파라미터 감소를 크게 달성한다.
• CIFAR-10의 VGG16 및 VGG16-BN에서 Top-1 정확도를 유지하면서 FLOPs 및 파라미터 감소가 크게 나타난다.
• CIFAR-10 및 ImageNet에서 ResNet-50은 상당한 FLOPs 감소와 파라미터 감소를 달성하되 정확도 손실은 최소화된다.
• SQuAD에서 BERT는 원샷 가지치기 프레임워크 하에서 경쟁력 있는 결과를 보여준다.
• HSPG는 그룹 희소성을 촉진하는 데 있어 표준 프로ximal 방법보다 우수하며 수렴 특성을 유지한다.
• 가지치기 과정은 전체 모델과 동일한 입력에 대해 동일한 출력을 계산하는 압축된 모델을 생성하며 미세조정이 필요 없다.