[논문 리뷰] MLPruning: A Multilevel Structured Pruning Framework for Transformer-based Models

Zhewei Yao, Linjian Ma외 3명|arXiv (Cornell University)|2021. 05. 30.

Topic Modeling참고 문헌 39인용 수 7

한 줄 요약

MLPruning는 Transformer 기반 모델을 위한 다수준 구조적 프루닝 프레임워크를 제안하며, 유연한 정규화를 위한 학습 가능한 Top-k 임계값을 사용하여 헤드, 행, 블록 단위의 희소 프루닝을 통합한다. 이는 GLUE 벤치마크에서 원본 모델의 약 20%의 가중치를 유지하면서 최대 3.69배의 추론 속도 향상을 달성하고 정확도 저하를 최소화한다.

ABSTRACT

Pruning is an effective method to reduce the memory footprint and computational cost associated with large natural language processing models. However, current approaches either only explore head pruning, which has a limited pruning ratio, or only focus on unstructured pruning, which has negligible effects on the real inference time and/or power consumption. To address these challenges, we develop a novel MultiLevel structured Pruning (MLPruning) framework, which uses three different levels of structured pruning: head pruning, row pruning, and block-wise sparse pruning. We propose using a learnable Top-k threshold, which employs an adaptive regularization to adjust the regularization magnitude adaptively, to select appropriate pruning ratios for different weight matrices. We also propose a two-step pipeline to combine block-wise pruning with head/row pruning to achieve high structured pruning ratios with minimum accuracy degradation. Our empirical results show that for \bertbase, with extapprox20\% of remaining weights, \OURS can achieve an accuracy that is comparable to the full model on QQP/MNLI/\squad, with up to extapprox3.69x speedup. Our framework has been open sourced~\cite{codebase}.

연구 동기 및 목표

• 기존의 프루닝 방법이 헤드 프루닝에만 집중하거나 비정형 프루닝을 통해 실질적인 효율성 향상이 미치지 못하는 한계를 해결하기 위해.
• 높은 압축 비율을 달성하면서도 모델 정확도를 유지할 수 있는 구조적 프루닝 프레임워크를 개발하기 위해.
• 헤드, 행, 블록 단위의 프루닝을 통합하여 Transformer 모델에 대해 통합적이고 효율적인 파이프라인을 설계하기 위해.
• 각 가중치 행렬의 중요도에 따라 조정 가능한 학습 가능한 Top-k 임계값을 활용한 적응형 정규화 메커니즘을 설계하기 위해.

제안 방법

• 프레임워크는 세 가지 수준의 구조적 프루닝을 적용한다: 헤드 프루닝을 통해 어텐션 헤드를 제거하고, 행 프루닝을 통해 피드포워드 레이어의 전체 행을 제거하며, 블록 단위의 희소 프루닝을 통해 가중치 행렬에 구조적인 희소성을 강제로 적용한다.
• 학습 가능한 Top-k 임계값을 도입하여 훈련 중 정규화 강도를 적응적으로 제어함으로써, 각 가중치 행렬에 맞는 최적의 프루닝 비율을 동적으로 선택할 수 있도록 한다.
• 두 단계 프로세스를 설계하여 먼저 헤드 및 행 프루닝을 적용한 후, 이를 블록 단위 프루닝과 통합함으로써 정확도 저하를 최소화하면서도 최대한의 희소성을 달성한다.
• 각 가중치 행렬의 중요도에 따라 스케일링되는 적응형 정규화를 사용하여 안정적인 훈련과 효과적인 프루닝을 보장한다.

실험 결과

연구 질문

• RQ1 헤드 전용 또는 비정형 프루닝 방법보다 높은 프루닝 비율을 달성하면서도 모델 정확도를 유지할 수 있는 다수준 구조적 프루닝 프레임워크는 가능한가?
• RQ2 학습 가능한 Top-k 임계값은 다양한 가중치 행렬에 맞춰 정규화를 적응적으로 조정하여 최적의 프루닝 성능을 달성하는 데 얼마나 효과적인가?
• RQ3 두 단계 파이프라인에서 헤드, 행, 블록 단위의 프루닝을 통합할 경우 추론 속도와 정확도에 어떤 영향을 미치는가?
• RQ4 구조적 프루닝은 자연어 처리 벤치마크에서 정확도 저하 없이 모델 크기와 추론 시간을 얼마나 줄일 수 있는가?

주요 결과

• MLPruning는 BERT-base에서 원본 모델의 약 20%의 파rameter만 유지하면서도 최대 3.69배의 추론 속도 향상을 달성한다.
• QQP, MNLI, SQuAD 벤치마크에서 프루닝된 모델는 상당한 파라미터 감소에도 불구하고 전체 모델과 유사한 정확도를 유지한다.
• 학습 가능한 Top-k 임계값의 사용은 적응형 정규화를 가능하게 하여 프루닝 효율성과 훈련 중 모델 안정성을 향상시킨다.
• 두 단계 파이프라인은 헤드/행 프루닝과 블록 단위 희소성을 효과적으로 통합하여 높은 구조적 프루닝 비율을 달성하면서도 정확도 저하를 최소화한다.