今天給大家分享下知乎搜尋中文字相關性和知識蒸餾的工作實踐,主要內容包括:
知乎搜尋文字相關性的演進BERT在知乎搜尋的應用和問題知識蒸餾及常見方案知乎搜尋在BERT蒸餾上的實踐01
知乎搜尋文字相關性的演進
1. 文字相關性的演進
我們首先來介紹下知乎搜尋中的文字相關性。在搜尋場景中,文字相關性可以定義為⽤戶搜尋query的意圖與召回 doc 內容的相關程度。我們需要透過不同模型來對這種相關程度進行建模。整體而言,文字的相關性一般可以分為兩個維度,字面匹配和語義相關。知乎搜尋中文字相關性模型的演進也是從這兩個方面出發並有所側重和發展。在知乎搜尋的整個架構中,文字相關性模型主要定位於為二輪精排模型提供更高維/抽象的特徵,同時也兼顧了一部分召回相關的工作。
2. Before NN
知乎搜尋中的文字相關性整體演進可以分為三個階段。在引入深度語義匹配模型前,知乎搜尋的文字相關性主要是基於TF-IDF/BM25的詞袋模型,下圖右邊是BM25的公式。詞袋模型通常來說是一個系統的工程,除了需要人工設計公式外,在統計詞的權重、詞頻的基礎上,還需要覆蓋率、擴充套件同義詞,緊密度等各種模組的協同配合,才能達到一個較好的效果。知乎搜尋相關性的一個比較早期的版本就是在這個基礎上迭代的。右下部分為在基於詞袋模型的基礎上,可以參考使用的一些具體特徵。
3. Before BERT
基於 BM25 的詞袋模型不管如何設計,主要還是隻解決文字相關性中的字面匹配這部分問題。第二階段引入的深度語義匹配模型則聚焦於解決語義相關的問題,主要分為兩部分:雙塔表示模型和底層互動模型。微軟的DSSM(左下)是雙塔模型的典型代表。雙塔模型透過兩個不同的 encoder來分別獲取query和doc的低維語義句向量表示,然後針對兩個語義向量來設計相關性函式(比如cosine)。DSSM擺脫了詞袋模型複雜的特徵工程和子模組設計,但也存在固有的缺陷:query和doc的語義表示是透過兩個完全獨立的 encoder 來獲取的,兩個固定的向量無法動態的擬合doc在不同 query的不同表示。這個反應到最後的精度上,肯定會有部分的損失。
底層互動模型一定程度上解決了這個問題。這個互動主要體現在 query 和 doc term/char 互動矩陣(中)的設計上,互動矩陣使模型能夠在靠近輸入層就能獲取 query 和 doc 的相關資訊。在這個基礎上,後續透過不同的神經網路設計來實現特徵提取得到 query-doc pair 的整體表示,最後透過全連線層來計算最終相關性得分。Match-Pyramid(右下)、KNRM(右上)是互動模型中比較有代表性的設計,我們在這兩個模型的基礎上做了一些探索和改進,相比於傳統的 BM25 詞袋模型取得了很大的提升。
4. BERT
對於互動模型來說,如下左圖,query和doc分別為sentence1和sentence2直接輸入到BERT模型中,透過BERT做一個整體的encoder去得到sentence pair的向量表示,再透過全連線層得到相似性打分,因為每個doc都是依賴query的,每個query-doc pair都需要線上實時計算,對GPU機器資源的消耗非常大,對整體的排序服務效能有比較大的影響。
基於上述原因,我們也做了類似於DSSM形式的表示模型,將BERT作為encoder,訓練資料的中的每個query和doc在輸入層沒有區分,都是做為不同的句子輸入,得到每個句向量表示,之後再對兩個表示向量做點乘,得到得到相關度打分。透過大量的實驗,我們最終採用了 BERT 輸出 token 序列向量的 average 作為句向量的表示。從互動模型到表示模型的妥協本質是空間換時間,因為doc是可以全量離線計算儲存的,線上只需要實時計算比較短的 query ,然後doc直接透過查表,節省了大量的線上計算。相比於互動模型,精度有一部分損失。
02
BERT在知乎搜尋的應用和問題
1. 搜尋業務架構中的BERT
在下圖中我們可以看到,BERT在知乎搜尋業務的召回和排序階段都扮演了比較重要的角色。互動模型的主要服務於二輪精排模型,依賴於線上實時的計算query和doc,為精排模組提供相關性特徵。表示模型又分為線上和離線兩塊,線上表示模型實時的為使用者輸入的query提供句向量表示,離線表示模型為庫中的doc進行批次句向量計算。一方面,doc向量透過TableStore/TiDB 和Redis的兩級儲存設計,為線上排序做查詢服務;另一方面,使用 faiss 對批次doc 向量構建語義索引,在傳統的 term 召回基礎上補充向量語義召回。
2. BERT表示模型語義召回
下面詳細介紹下我們的語義召回模型。首先看個例子,對於「瑪莎拉蒂 ghlib」這個case,使用者真正想搜的是「瑪莎拉蒂 Ghibli」這款車,但使用者一般很難記住完整的名稱,可能會輸錯。在輸錯的情況下,基於傳統的term匹配方式(Google搜尋的例子)只能召回“瑪莎拉蒂”相關的 doc,而無法進行這輛車型的召回,這種場景下就需要進行語義召回。更通用的來說,語義召回可以理解為增加了字面不匹配但是語義相關的 doc 的召回。
語義召回模型整體是BERT 相關性任務中雙塔表示模型的一個應用。BERT做為encoder來對query和doc進行向量的表示,基於faiss對全量 doc 向量構建語義索引,線上實時的用query向量進行召回。這個策略上線後,線上top20 doc中語義召回doc數量佔總召回 doc 數量的比例能到達 5%+。
3. BERT帶來的問題
BEER 模型上線後,為不同的模組都取得了不錯收益的同時,也給整個系統帶來了不少問題。這些問題整體可以歸結為線上實時計算、離線儲存、模型迭代三個方面。具體的見上圖。
4. 蒸餾前的嘗試
針對上述效能或儲存的問題,在對BERT 蒸餾之前,我們也進行了很多不同的嘗試。
BERT 互動模型的部署放棄了使用原生TF serving,而是在cuda 的基礎上用c++ 重寫了模型的載入和serving,加上混合精度的使用。在我們的業務規模上,線上實時效能提高到原來的約 1.5 倍,使BERT互動模型滿足了的最低的可上線要求。在這個基礎上,對線上的 BERT 表示模型增加 cache,減少約 60% 的請求,有效減少了GPU 機器資源的消耗。
另一個思路是嘗試給BERT在橫向和縱向維度上瘦身。橫向上,一方面可以減小serving 時 max_seq_length長度,減少計算量;另一方面可以對錶示向量進行維度壓縮來降低儲存開銷。這兩種嘗試在離線和線上指標上都有不同程度的損失,因此被放棄。縱向上,主要是減少模型的深度,即減少 transformer層數。這對於視訊記憶體和計算量都能得到顯著的最佳化。前期嘗試過直接訓練小模型,以及使用BERT-base若干層在下游的相關性任務上進行fine-tune。這兩種方案,在離線指標上的表現就沒法達到要求,因此也沒有上線。
針對 doc數量過大,儲存開銷過大和語義索引構建慢的問題。在這方面做了一個妥協的方案:透過wilson score 等規則過濾掉大部分低質量的 doc,只對約 1/3 的doc 儲存表示向量和構建語義索引。該方案會導致部分文件的相關性特徵存在缺失。對於表示模型存在的低互動問題,嘗試Poly-encoder(Facebook方案)將固定的 768維表示向量轉為多個head的形式,用多個head做attention的計算,保證效能在部分下降的前提得到部分精度的提升。
03
智知識蒸餾及常見方案
1. 知識蒸餾
下面簡單介紹下知識蒸餾。從下圖中看,我們可以把知識蒸餾的整體形式簡化為:大模型不考慮效能問題儘量學習更多的知識(資料),小模型透過適量的資料去高效地學習大模型的輸出,達到一個知識遷移的效果。實際 serving 使用的是小模型。
知識蒸餾為什麼能有效?關鍵點在於 soft target 和 temperature。soft target對應的是teacher模型的輸出,類似於機率分佈,知識蒸餾從hard target轉為soft target的學習有利於模型更好的去擬合標籤,引入temperature則是為了進一步平滑標籤,讓模型去學習到類別和類別中的知識。這裡需要注意的是,temperature 的選取不宜過大,太大的 temperature 會導致不同類別之間的差異被完全平滑掉。
2. BERT蒸餾方案
對與BERT的蒸餾我們做了大量的調研,並對目前主流的蒸餾方案做了歸納分類。基於任務維度來說,主要對應於現在的pretrain + fine-tune 的兩段式訓練。在預訓練階段和下游任務階段都有不少的方案涉及。技巧層面來分的話,主要包括不同的遷移知識和模型結構的設計兩方面。後面我會選兩個典型的模型簡單介紹一下。
3. 蒸餾-MiniLM
MiniLM是基於預訓練任務的蒸餾,其是一種通用的面向Transformer-based預訓練模型壓縮演算法。主要改進點有三個,一是蒸餾teacher模型最後一層Transformer的自注意力模組,二是在自注意模組中引入 values-values點乘矩陣的知識遷移,三是使⽤了 assistant ⽹絡來輔助蒸餾。
4. 蒸餾-BERT to Simple NN
BERT to Simple NN更多的是做了一些loss形式的設計,使其訓練方式更高效。
04
知乎搜尋再BERT蒸餾上的實踐
1. BERT蒸餾上的實踐和收益
前面的介紹中我有提到,在做 BERT蒸餾前其實已經做了很多嘗試,但是多少都會有精度的損失。因此,我們做蒸餾的第一目標是離線模型對⽐線上 BERT精度⽆損。但對BERT-base 直接進行蒸餾,無論如何都沒辦法避免精度的損失,所以我們嘗試用更大的模型(比如BERT-large/Robert-large/XLNET)來作為 teacher 進行蒸餾。這些多層的模型均在我們知乎全量語料先做pretrain,再做fine-tune,得到微調後的模型再做蒸餾。
2. 蒸餾-Patient KD
我們對互動模型和表示模型都做了蒸餾,主要採用了Patient KD模型的結構設計,Student模型基於BERT-base的若干層運用不同的策略進行引數的初始化,去學習Robert-large大模型的方案。
其中知識遷移主要有三部分:student的預測與真實標籤的交叉熵、student與teacher的預測的交叉熵和中間隱層的向量之間的normalized MSE。
3. BERT互動模型蒸餾
對於我們選的teacher模型Robert-large,單純預訓練模型其nDCG指標為0.914,線上之前使用的BERT-base 是0.907,若對BERT-base的若干6層直接去做fine-tune能達到的最高指標是0.903,對比於BERT-base精度會損失很多。
我們這塊做了一些嘗試,基於Robert-large從24層蒸餾到6層的話能到0.911,能超過線上BERT-base的效果。
4. BERT表示模型蒸餾
在BERT表示模型上,蒸餾時我們希望對向量維度和模型層數同時進行壓縮,但蒸餾後得到的student模型表現不及預期。所以最後上線的方案中,表示模型層數還是維持了12層。在蒸餾時,為了提高精度,選取互動模型作為teacher進行蒸餾。因為互動模型是query和doc之間的打分,互動模型得到的logits與表示模型點乘後的打分在數量值會有較大差值,所以用pairwise形式透過teacher差值擬合來進行loss的計算。
在維度壓縮方面我們做了對比實驗,BERT模型輸出做 average pooling 後接全連線層分別壓縮至8維到768維。如圖所示,128維和64維的表現跟768維差別不大,在上線時選擇維度為64和128進行嘗試,兩者在線上表現沒有太明顯的差異,最終選擇了64維的方案,把模型的維度壓縮了12倍,儲存消耗更低。
5. 蒸餾的收益
蒸餾的收益主要分為線上和離線兩部分。
線上方面:
互動模型的層數從12層壓縮到6層,排序相關性特徵P95減少為原本的1/2,整體搜尋入口下降40ms,模型部署所需的GPU機器數也減少了一半,降低了資源消耗。
表示模型語義索引儲存規模title減為1/4,content維度從768維壓縮至64維,雖然維度減少了12倍,但增加了倒排索引doc的數量,所以content最終減為1/6,
語義索引召回也有比較大的提升,title減少為1/3,content減少為1/2。精排模組需要線上實時查詢離線計算好的向量,所以查詢服務也有提升。
離線方面:
表示模型語義索引的構建時間減少為1/4,底層知乎自研的TableStore/TIDB儲存減為原來的1/6,LTR訓練資料和訓練時間都有很大的提升,粗排早期用的是BM25等基礎特徵,後來引入了32維的BERT向量,提升了精排精度。