[논문 리뷰] Accelerating BERT Inference for Sequence Labeling via Early-Exit

Xiaonan Li, Yunfan Shao외 4명|arXiv (Cornell University)|2021. 05. 28.

Topic Modeling참고 문헌 39인용 수 4

한 줄 요약

이 논문은 문장 분류 작업을 위한 BERT 추론을 가속화하기 위해 SENTEE와 TOKEE라는 두 가지 초기 종료 메커니즘을 제안한다. TOKEE는 윈도우 기반의 불확실성 기준과 자체 샘플링 미세조정을 활용하여 토큰 수준의 초기 종료를 가능하게 하며, 성능 저하를 최소화하면서 최대 75%의 추론 비용 절감을 달성한다. 이는 동일한 속도 향상 비율에서 DistilBERT를 능가한다.

ABSTRACT

Both performance and efficiency are crucial factors for sequence labeling tasks in many real-world scenarios. Although the pre-trained models (PTMs) have significantly improved the performance of various sequence labeling tasks, their computational cost is expensive. To alleviate this problem, we extend the recent successful early-exit mechanism to accelerate the inference of PTMs for sequence labeling tasks. However, existing early-exit mechanisms are specifically designed for sequence-level tasks, rather than sequence labeling. In this paper, we first propose a simple extension of sentence-level early-exit for sequence labeling tasks. To further reduce the computational cost, we also propose a token-level early-exit mechanism that allows partial tokens to exit early at different layers. Considering the local dependency inherent in sequence labeling, we employed a window-based criterion to decide for a token whether or not to exit. The token-level early-exit brings the gap between training and inference, so we introduce an extra self-sampling fine-tuning stage to alleviate it. The extensive experiments on three popular sequence labeling tasks show that our approach can save up to 66%-75% inference cost with minimal performance degradation. Compared with competitive compressed models such as DistilBERT, our approach can achieve better performance under the same speed-up ratios of 2X, 3X, and 4X.

연구 동기 및 목표

• 실시간 또는 모바일 환경에서 특히 높은 계산 비용을 유발하는 BERT의 문장 분류 작업에서의 비용을 해결한다.
• 이전에는 문장 수준 작업(예: 텍스트 분류)에 사용된 초기 종료 메커니즘을 토큰 수준의 문장 분류에 확장하여, 각 토큰의 신뢰도가 필요한 상황에 대응한다.
• 로컬 컨텍스트를 기반으로 개별 토큰이 다른 레이어에서 다른 시점에 종료될 수 있도록 가능한 세분화되고 동적인 계산 비용 할당 전략을 설계한다.
• 추론 중 종료된 토큰의 표현을 복사함으로써 발생하는 학습-추론 불일치 문제를 완화한다.
• 성능을 희생시키지 않은 채 효율성을 향상시켜, 동일한 속도 향상 비율에서 DistilBERT와 같은 정제된 모델을 능가한다.

제안 방법

• 모든 토큰이 동시에 종료될 수 있도록 허용하는 문장 수준의 초기 종료 메커니즘인 SENTEE를 제안한다. 이는 최대 토큰 불확실성이 임계값 이하일 경우에 적용된다.
• 각 토큰이 자신만의 신뢰도에 따라 다른 레이어에서 종료될 수 있도록 허용하는 토큰 수준의 초기 종료 메커니즘인 TOKEE를 도입한다. 이는 계산 효율성을 향상시킨다.
• 지역적 컨텍스트를 고려하여 토큰이 종료될 수 있는지 결정하기 위해 윈도우 기반의 불확실성 기준을 적용한다. 이는 순차적 의존성을 반영한다.
• 학습 중에 초기 종료를 시뮬레이션하기 위해 불확실성 기반으로 종료 레이어를 샘플링하는 자체 샘플링 미세조정 단계를 구현한다. 이는 학습과 추론 간 격차를 줄인다.
• 종료된 토큰의 표현을 재계산 없이 고려하여 상위 레이어로 전파하는 '정지 및 복사' 메커니즘을 사용한다. 이는 모델 구조를 유지한다.
• 표준 BERT 인코더-디코더 프레임워크에 초기 종료 메커니즘을 통합하여, 표준 문장 분류 헤드와의 호환성을 유지한다.

실험 결과

연구 질문

• RQ1 문장 수준 작업에 설계된 초기 종료 메커니즘이, 예측이 토큰 수준이며 토큰 간 의존성이 존재하는 문장 분류 작업에 효과적으로 적용될 수 있는가?
• RQ2 장문의 시퀀스에서 윈도우 기반의 불확실성 기준을 사용한 토큰 수준의 초기 종료가 문장 수준의 초기 종료보다 더 높은 계산 효율성을 달성할 수 있는가?
• RQ3 정지 및 복사 표현으로 인해 발생하는 학습-추론 불일치 문제를 미세조정 과정에서 효과적으로 완화할 수 있는가?
• RQ4 다양한 속도 향상 비율에서 TOKEE의 성능-비용 트레이드오프는 DistilBERT와 같은 정제된 모델과 비교해 어떻게 되는가?
• RQ5 TOKEE를 통한 세분화되고 동적인 계산 자원 할당 전략은 높은 정확도를 유지하면서 추론 비용을 크게 절감할 수 있는가?

주요 결과

• TOKEE는 문장 분류 작업에서 최대 75%의 추론 비용 절감을 달성하며, 성능 저하를 최소화한다. 이는 동일한 속도 향상 비율에서 DistilBERT를 능가한다.
• SENTEE는 최대 66%의 추론 비용 절감을 달성하지만, 모든 토큰이 동시에 종료되어야 하므로 시퀀스 길이가 증가할수록 속도 향상 비율이 감소한다.
• 고속도 향상 비율(예: 4배)에서 TOKEE는 SENTEE보다 높은 성능을 유지한다. 4배 속도 향상에서 SENTEE는 평균적으로 레이어 수가 2.5개로 줄어든다.
• 자체 샘플링 미세조정은 고속도 향상 비율에서 성능 저하를 크게 줄여주며, 무작위 샘플링 및 미세조정 없음 대비 뛰어난 성능을 보인다.
• TOKEE의 레이어 분포를 보면, 속도 향상 비율이 증가할수록 더 많은 토큰이 빨리 종료되며, 일부는 여전히 깊은 레이어를 사용한다. 이는 세분화된 비용 할당을 가능하게 한다.
• 윈도우 기반의 불확실성 기준은 지역적 컨텍스트를 고려함으로써 강건성을 향상시켜, 개별 토큰의 신뢰도에만 의존하는 것보다 더 신뢰할 수 있는 초기 종료 결정을 가능하게 한다.