[논문 리뷰] Simplified TinyBERT: Knowledge Distillation for Document Retrieval

Xuanang Chen, Ben He외 3명|arXiv (Cornell University)|2020. 09. 16.

Topic Modeling참고 문헌 21인용 수 6

한 줄 요약

이 논문은 지식 정복 프레임워크인 Simplified TinyBERT를 제안하며, 두 단계로 나누어진 TinyBERT의 훈련 과정을 하나로 통합하고 속도 손실를 손실 함수에 통합하여 문서 검색 성능을 향상시킵니다. BERT-Base 대비 15배 빠른 추론 속도를 달성하면서도 MS MARCO 및 TREC 2019 DL 벤치마크에서 BERT-Base와 표준 TinyBERT를 크게 앞서며 뛰어난 성능을 보입니다.

ABSTRACT

Despite the effectiveness of utilizing the BERT model for document ranking, the high computational cost of such approaches limits their uses. To this end, this paper first empirically investigates the effectiveness of two knowledge distillation models on the document ranking task. In addition, on top of the recently proposed TinyBERT model, two simplifications are proposed. Evaluations on two different and widely-used benchmarks demonstrate that Simplified TinyBERT with the proposed simplifications not only boosts TinyBERT, but also significantly outperforms BERT-Base when providing 15$ imes$ speedup.

연구 동기 및 목표

• BERT의 높은 추론 비용으로 인해 배포가 제한되는 문서 랭킹 작업에 대해 지식 정복의 효과성을 조사하기 위해.
• 문서 랭킹 작업에 적합한 성능 향상을 위해 훈련 과정을 단순화하고 정복 효과를 향상시키기 위해 TinyBERT 모델을 개선하기 위해.
• 더 작고 빠른 학생 모델을 구현하여 추론 속도가 15배 빨라지더라도 BERT-Base의 랭킹 성능를 초월하면서도 높은 효율성을 유지하기 위해.

제안 방법

• 두 단계로 나누어진 TinyBERT의 작업별 정복 단계를 하나의 훈련 단계로 통합하기 위해 복합 손실 함수를 사용하였습니다.
• 주의, 은닉 상태, 임bedding 및 소프트 레이블 정복 손실을 결합한 통합 손실 함수를 제안하였습니다: $\mathcal{L} = \mathcal{L}_{attn} + \mathcal{L}_{hidn} + \mathcal{L}_{emb} + \mathcal{L}_{soft}$.
• 관련 문서와 비관련 문서 간의 구분 능력을 향상시키기 위해 하드 레이블을 손실 함수에 통합하였습니다.
• 일괄 정복 단계를 사용하여 학생 모델을 BERT-Base의 첫 번째 k개의 레이어로 초기화함으로써 별도의 일반 정복 단계가 필요 없도록 하였습니다.
• 고정된 크기의 학생 모델에서 온도 $T \in \{1,5,10\}$ 과 지식 정복 가중치 $\alpha \in \{0.2,0.5,0.7\}$ 에 대해 초모수 탐색을 수행하였습니다.
• MRR@10 및 NDCG@10와 같은 표준 랭킹 메트릭을 사용하여 MS MARCO 및 TREC 2019 DL 트랙 벤치마크에서 모델을 평가하였습니다.

실험 결과

연구 질문

• RQ1 고비용 추론이 요구되는 문서 검색 작업에 표준 지식 정복 및 TinyBERT를 효과적으로 적용할 수 있는가?
• RQ2 TinyBERT의 두 단계 작업별 정복을 하나로 통합하면 성능 향상과 훈련 효율성이 향상되는가?
• RQ3 정복 손실에 하드 레이블을 추가하면 관련 문서와 비관련 문서를 구분하는 능력이 향상되는가?
• RQ4 Simplified TinyBERT의 단순화된 훈련 과정은 15배 빠른 추론 속도를 확보하면서도 BERT-Base의 랭킹 성능를 뛰어넘는가?
• RQ5 제안된 단순화 방법은 다양한 모델 크기와 재랭킹 깊이에서 성능을 유지하거나 향상시킬 수 있는가?

주요 결과

• L6_H768 설정을 가진 Simplified TinyBERT는 15배 빠른 속도에서 MS MARCO 개발 세트에서 BERT-Base 및 표준 TinyBERT를 크게 앞서며 뛰어난 성능을 보입니다.
• TREC 2019 DL 테스트 세트에서, L3_H384 학생 모델을 사용할 경우 Simplified TinyBERT는 NDCG@10에서 BERT-Base를 앞서며 성능을 냅니다. 반면 동일한 조건에서 TinyBERT는 BERT-Base보다 뚜렷하게 열등한 성능을 보입니다.
• 모든 재랭킹 깊이(상위 10개에서 상위 100개)에서 Simplified TinyBERT는 BERT-Base 및 TinyBERT를 항상 앞서며, 깊이 100일 때 MS MARCO 개발 세트에서 MRR@10가 3.5% 향상되었습니다.
• 단일 단계 훈련으로 인해 표준 TinyBERT 대비 두 번째 단계 훈련 시간이 42~52% 감소하였으며, L6_H768의 경우 훈련 시간이 14.37시간에서 29.95시간으로 감소하였습니다.
• 제거 실험 결과, 단일 단계 훈련과 하드 레이블 통합이라는 두 가지 개선 조치가 성능 향상에 독립적이고 상호보완적으로 기여하는 것으로 확인되었습니다.
• Simplified TinyBERT는 정복된 모델 중 문서 검색 분야에서 최고 성능을 기록하였으며, L3_H384 버전은 15배의 빠른 속도를 확보하면서도 BERT-Base의 랭킹 효과성보다 뛰어난 성능을 보였습니다.