[논문 리뷰] TR-BERT: Dynamic Token Reduction for Accelerating BERT Inference

Deming Ye, Yankai Lin외 2명|arXiv (Cornell University)|2021. 05. 25.

Topic Modeling참고 문헌 48인용 수 7

한 줄 요약

TR-BERT는 강화학습을 통해 각 레이어에서 작업의 관련성이 높은 토큰을 유지할지 여부를 학습함으로써, BERT 추론을 가속화하기 위한 동적 토큰 감소 방법을 제안한다. 정확도 저하를 최소화하면서 2–5배의 속도 향상을 달성하며, 특히 자기주의성 주의의 복잡도를 선택적 토큰 전파를 통해 제곱형으로 감소시켜 장문의 시퀀스 작업에 유리하다.

ABSTRACT

Existing pre-trained language models (PLMs) are often computationally expensive in inference, making them impractical in various resource-limited real-world applications. To address this issue, we propose a dynamic token reduction approach to accelerate PLMs' inference, named TR-BERT, which could flexibly adapt the layer number of each token in inference to avoid redundant calculation. Specially, TR-BERT formulates the token reduction process as a multi-step token selection problem and automatically learns the selection strategy via reinforcement learning. The experimental results on several downstream NLP tasks show that TR-BERT is able to speed up BERT by 2-5 times to satisfy various performance demands. Moreover, TR-BERT can also achieve better performance with less computation in a suite of long-text tasks since its token-level layer number adaption greatly accelerates the self-attention operation in PLMs. The source code and experiment details of this paper can be obtained from https://github.com/thunlp/TR-BERT.

연구 동기 및 목표

• 자원 제한 환경에서 사전 학습된 언어 모델(PLMs)의 높은 추론 지연을 해결하기 위해.
• 모든 레이어를 균일하게 줄이는 고정 레이어 프루닝의 한계를 극복하기 위해, 복잡한 작업에서 성능 저하를 야기하지 않도록 하기 위해.
• 각 토큰에 대해 동적으로 어떤 토큰을 더 깊은 레이어로 전파할지 선택함으로써, 토큰 단위의 적응형 계산을 가능하게 하기 위해.
• 시퀀스 길이에 따라 자기주의성 주의의 복잡도가 제곱형으로 증가하는 장문의 텍스트 환경에서의 효율성을 향상시키기 위해.
• 실제 배포 환경에서 유연하고 요구에 따라 속도-정확도 트레이드오프를 제공하기 위해.

제안 방법

• TR-BERT는 각 레이어에서 각 토큰이 선택되거나 스킵될지를 결정하는 다단계 선택 문제로 토큰 감소를 공식화한다.
• 작업별 중요도에 기반해 동적 토큰 선택 전략을 학습하기 위해 강화학습 프레임워크 내 정책 네트워크를 사용한다.
• 예측 신뢰도와 선택된 토큰 수에 대한 페널티를 조합한 보상을 수신함으로써, 속도와 정확도의 균형을 이루는 RL 에이전트를 운영한다.
• 낮은 중요도로 간주되는 토큰들은 조기에 프루닝되어 이후 계산을 줄이고 불필요한 자기주의성 주의 연산을 피한다.
• 레이어 수준의 프루닝과 달리 토큰 수준에서 운영되어 입력 토큰당 계산을 세밀하게 제어할 수 있다.
• 기본 BERT 모델의 재학습이 필요 없고 기존의 PLMs와 호환되는 플러그 앤 플레이 프레임워크를 제공한다.

실험 결과

연구 질문

• RQ1 동적이고 토큰 수준의 프루닝은 모델 성능을 저하시키지 않고 추론 속도를 향상시킬 수 있는가?
• RQ2 작업 관련성에 기반한 동적 토큰 선택 전략이 고정 레이어 또는 레이어 단위의 프루닝보다 어떻게 다른가?
• RQ3 인간의 주석 없이 강화학습이 토큰 중요도 신호를 효과적으로 학습할 수 있는가?
• RQ4 장문의 자연어 처리 작업에서 토큰 수준의 프루닝이 얼마나 효율성을 향상시키는가?
• RQ5 이 방법은 다양한 최종 작업에 걸쳐 다양한 속도-정확도 트레이드오프를 달성할 수 있는가?

주요 결과

• TR-BERT는 11개의 자연어 처리 벤치마크에서 원하는 정확도 수준에 따라 2–5배의 더 빠른 추론 속도를 달성한다.
• 질문 응답 작업에서 TR-BERT는 특히 고정 프루닝이 실패하는 복잡한 추론 시나리오에서 베이스라인 방법보다 뛰어난 성능을 보인다.
• 장문의 텍스트 작업(예: 1,024개의 입력 토큰)에서는 TR-BERT가 순수 BERT보다 훨씬 빠른 추론 시간을 확보하면서도 높은 성능을 유지한다.
• 프루닝된 시퀀스에서도 높은 예측 신뢰도를 유지함으로써 의미적으로 유의미한 토큰을 효과적으로 선택하고 있음을 시사한다.
• 인간의 주석 없이도 강화학습 기반의 토큰 선택 정책이 질문어나 감정 지표와 같은 작업 관련 토큰을 유지하는 데 성공한다.
• 텍스트 분류, 자연어 추론, 질문 응답과 같은 다양한 자연어 처리 작업에서 이 방법이 뛰어난 강건성을 보였다.