[논문 리뷰] Learned Token Pruning for Transformers

Sehoon Kim, Sheng Shen외 5명|arXiv (Cornell University)|2021. 07. 02.

Advanced Neural Network Applications인용 수 24

한 줄 요약

이 논문은 Learned Token Pruning (LTP)을 도입하는데, 이는 트랜스포머의 토큰 프루닝을 계층별로 학습 가능한 임계값을 사용하여 적응적으로 수행하는 임계값 기반 토큰 프루닝 방법으로, 최소한의 정확도 손실과 입력 길이 변화에 대한 강건성 향상을 달성하며 FLOPs를 크게 감소시킨다.

ABSTRACT

Deploying transformer models in practice is challenging due to their inference cost, which scales quadratically with input sequence length. To address this, we present a novel Learned Token Pruning (LTP) method which adaptively removes unimportant tokens as an input sequence passes through transformer layers. In particular, LTP prunes tokens with an attention score below a threshold value which is learned for each layer during training. Our threshold-based method allows the length of the pruned sequence to vary adaptively based on the input sequence, and avoids algorithmically expensive operations such as top-k token selection. We extensively test the performance of LTP on GLUE tasks and show that our method outperforms the prior state-of-the-art token pruning methods by up to ~2.5% higher accuracy with the same amount of FLOPs. In particular, LTP achieves up to 2.1x FLOPs reduction with less than 1% accuracy drop, which results in up to 1.9x and 2.0x throughput improvement on Intel Haswell CPUs and NVIDIA V100 GPUs, respectively. Furthermore, we demonstrate that LTP is more robust than prior methods to variations on input sentence lengths. Our code has been developed in PyTorch and has been open-sourced.

연구 동기 및 목표

• 레이어별로 토큰을 적응적으로 프루닝하여 트랜스포머 모델의 추론 비용을 줄인다.
• 상위 k 토큰 선택을 피하고 미분 가능하며 임계값 기반의 프루닝 메커니즘을 개발한다.
• 훈련 중에 레이어별 프루닝 임계값을 학습하고 추론 시 하드 프루닝을 적용한다.
• GLUE와 SQuAD 데이터셋 전반에서 효율성 향상과 강건성을 입증한다.

제안 방법

• 헤드와 토큰 간의 평균 어텐션 확률을 통해 토큰의 중요도를 정의한다.
• 상위-k 프루닝을 시그모이드 기반 임계값을 이용한 학습 가능하고 미분 가능한 소프트 마스크로 대체하여 임계값으로의 그래디언트 흐름을 가능하게 한다.
• 소프트 프루닝 단계에서 모델 매개변수와 함께 임계값을 학습한 후 이진화하고 미세 조정한다(세 단계 과정).
• 프루닝을 촉진하고 임계값 학습을 안정시키기 위해 L1 정규화 항을 도입한다.
• 양자화 및 지식 증류와의 호환성을 보여 모델을 더 압축한다.

실험 결과

연구 질문

• RQ1 임계값 기반 토큰 프루닝이 NLP 작업 전반에서 유의미한 FLOPs 감소와 비슷한 정확도를 달성할 수 있는가?
• RQ2 학습 가능한 레이어별 임계값이 고정 프루닝 구성에 비해 입력 시퀀스 길이의 변화에 대한 강건성을 제공하는가?
• RQ3 LTP는 정확도 대 FLOPs 측면에서 이전 토큰 프루닝 방법(SpAtten, LAT)과 어떻게 비교되는가?
• RQ4 이 접근법이 양자화, 지식 증류와 같은 다른 압축 기법과의 호환성이 있는가?

주요 결과

작업	RoBERTa_base_Accuracy	LTP_Accuracy	RoBERTa_base_GFLOPs	LTP_GFLOPs	속도향상
MNLI-m	87.53	86.53	6.83	3.64	1.88×
MNLI-mm	87.36	86.37	7.15	3.63	1.97×
QQP	90.39	89.69	5.31	2.53	2.10×
QNLI	92.86	91.98	8.94	4.77	1.87×
SST-2	94.27	93.46	4.45	2.13	2.09×
STS-B	90.89	90.03	5.53	2.84	1.95×
MRPC	92.14	91.59	9.33	4.44	2.10×
RTE	77.98	77.98	11.38	6.30	1.81×
SQuAD 2.0	83.04	82.25	32.12	16.99	1.89×

• LTP는 GLUE/SQuAD에서 최대 2.10× FLOPs 감소를 달성하고 정확도 손실은 1% 미만이며 CPU/GPU에서 최대 1.93–1.97× 처리량 이득을 제공합니다.
• LTP는 동일한 FLOPs에서 SpAtten 및 LAT를 consistently 능가하며 각 작업에서 약 최대 ~2.5% 더 높은 정확도를 달성합니다.
• LTP는 입력 길이 변화에 대해 강한 강건성을 보이며, 가변 문장 길이를 가진 작업에서 LAT보다 큰 차이로 우수한 성능을 보입니다.
• 직접 하드웨어 처리량 시연에서 배치 크기에 따라 최대 약 1.9×–2.0× 이득이 나타납니다.
• 정량화 및 지식 증류는 LTP와 결합될 때 정확도 손실이 작으면서 BOPs를 최대 10×까지 더 줄일 수 있습니다.