Improving Bi-encoder Neural Ranking Models using Knowledge Distillation and Lightweight Fine-tuning

Abstract

In recent studies, pre-trained language models, especially bidirectional encoder representations from transformers (BERT) have been essential in enhancing the performance of neural ranking models (NRMs). Various BERT-based NRMs have been proposed, and many have achieved state-of-the-art performance. BERT-based NRMs can be classified according to how the query and document are encoded through BERTs self-attention layers: bi-encoder versus cross-encoder. Bi-encoder models are highly efficient because all the documents can be pre-processed before the query time, but their performance is inferior compared to cross-encoder models. Because of their efficiency, bi-encoder models are much more deployable in real search engines and tend to receive more attention from industrial practitioners. However, their performance does not reach that of cross-encoder models. Therefore, improving the performance of bi-encoder models is a promising research direction. This thesis explores the methods to improve bi-encoder NRMs using knowledge distillation and lightweight fine-tuning. We consider a method that transfers the knowledge of a teacher cross-encoder model to a student bi-encoder model using knowledge distillation. Knowledge distillation enables a bi-encoder student to imitate the representation of a cross-encoder teacher and have the advantages of both types of models. The resulting student bi-encoder achieves an improved performance by simultaneously learning from a cross-encoder teacher and a bi-encoder teacher. We also investigate lightweight fine-tuning to improve bi-encoder NRMs. Lightweight fine-tuning is a method of fine-tuning only a small portion of the model weights, and is known to have a regularization effect. We demonstrate two approaches for improving the performance of BERT-based bi-encoders using lightweight fine-tuning. The first approach is to replace the full fine-tuning step with lightweight fine-tuning. The second is to develop semi-Siamese models in which queries and documents are handled with a limited amount of difference. The limited difference is realized by learning two lightweight fine-tuning modules, where the main language model of BERT is kept common for both query and document. We provide extensive experimental results, which confirm that both lightweight fine-tuning and semi-Siamese models are considerably helpful for improving BERT-based bi-encoders. Finally, we present a model that uses these two methods simultaneously. Using knowledge distillation and lightweight fine-tuning methods together, a model can gain the effects of both methods, resulting in further performance improvement over the individual methods. We anticipate that these techniques will be broadly applicable to industrial domains.

최근 연구에서 다양한 BERT기반의 신경망 랭킹 모델이 제안되었고, 이 모델들은 최고의 성능을 보여주고 있다. BRET기반 랭킹 모델은 쿼리와 문서간의 관계가 BERT의 셀프 어텐션을 통해서 계산되는가의 여부에 따라 크로스 인코더와 바이 인코더로 구분된다. 크로스 인코더 모델은 높은 성능을 가지고 있지만 효율이 좋지 못한 단점이 있다. 반면, 바이 인코더 모델은 크로스 인코더에 비해 성능은 떨어지지만, 모든 문서의 벡터 표현형을 미리 구해놓을 수 있기 때문에 높은 효율성을 가지고 있다. 바이 인코더 모델은 효율적이기 때문에 실제 검색 엔진에 배포가 가능하다. 이런 이유로 바이 인코더 모델은 검색 업계로부터 더 많은 관심을 받는다. 그러나 앞에서 언급했듯이, 바이 인코더 모델의 성능이 크로스 인코더 모델에 도달하지 못한다는 문제가 있다. 따라서 바이 인코더 모델의 성능을 향상시키는 것은 랭킹 모델을 실제로 이용하려고 하는 영역에서는 매력적인 문제이다. 이 연구에서는 지식 증류와 경량 파인튜닝을 이용하여 바이 인코더 모델을 개선하는 방법을 탐구한다. 우리는 지식 증류를 사용하여 크로스 인코더 모델의 지식을 바이 인코더 모델로 전달하는 방법을 연구한다. 지식 증류를 통해 만들어진 바이 인코더 모델은 크로스 인코더로부터 배운 지식을 이용하기 때문에 성능이 향상된다. 우리는 또한 바이 인코더 모델을 개선하기 위한 경량 파인튜닝 방법을 이용한다. 경량 파인튜닝은 모델 가중치의 일부만 미세하게 학습하는 방법으로, 모델의 정규화 효과가 있는 것으로 알려져 있다. 경량 파인튜닝을 사용하여 BERT기반 바이 인코더 모델의 성능을 개선하기 위한, 두 가지 접근 방식을 이용한다. 첫 번째 접근 방식은 파인튜닝을 경량 파인튜닝으로 대체하는 것이다. 두 번째 접근 방식은 쿼리와 문서를 서로 다르게 처리하는 세미 샴 모델을 이용하는 것이다. 우리는 다양한 실험을 통하여 경량 파인튜닝 방법과 세미 샴 모델이 바이 인코더 모델을 개선하는 데 상당히 도움이 됨을 확인하였다. 마지막으로 지식증류와 경량 파인튜닝 방법을 동시에 사용하는 모델을 제시한다. 두 방법을 모두 사용한 모델이 두 방법을 사용한 각각의 방법보다 성능이 더 좋음을 실험으로 확인하였다. 우리가 제안한 방법이 검색 업계에 도움이 될 것으로 기대한다.