如何改進(jìn)雙塔模型才能更好的提升你的算法效果

來自：對白的算法屋

今天寫點技術(shù)干貨來回饋一下我的粉絲們。本來想繼續(xù)寫對比學(xué)習(xí)（Contrastive Learing）相關(guān)類型的文章，以滿足我出一本AI前沿技術(shù)書籍的夢想，但奈何NIPS2021接收的論文一直未公開，在arxiv上不停地刷，也只翻到了零碎的幾篇。于是，我想到該寫一下雙塔模型了，剛進(jìn)美團(tuán)的第一個月我就用到了Sentence-BERT。

為什么呢？因為雙塔模型在NLP和搜廣推中的應(yīng)用實在太廣泛了。不管是校招社招，面試NLP或推薦算法崗，這已經(jīng)是必問的知識點了。

接下來，我將從模型結(jié)構(gòu)，訓(xùn)練樣本構(gòu)造，模型目標(biāo)函數(shù)三個方面介紹雙塔模型該如何改進(jìn)，才能更好的提升業(yè)務(wù)中的效果。

一、雙塔模型結(jié)構(gòu)改進(jìn)

如圖所示，目前主流的雙塔模型結(jié)構(gòu)主要可以歸為三類。

第一類在離線階段直接將BERT編碼的document映射為固定長度的向量，在線階段將query映射為固定長度的向量，然后通過打分函數(shù)計算最后的得分，例如：Sentence-BERT，DPR。

第二類

模型在離線階段將BERT編碼document得到的多個向量（每個向量對應(yīng)一個token）全部保留，在線階段利用BERT將query編碼成多個向量，和離線階段保留的document的多個向量進(jìn)行交互打分（復(fù)雜度O（mn）），得到最后的得分，代表工作，Col-BERT。

第三類

模型是前兩種的折中，將離線階段BERT編碼得到的document向量進(jìn)行壓縮，保留k個（k《m）個向量，并且使用一個向量來表示query（一般query包含的信息較少），在線階段計算一個query向量和k個document向量的交互打分（復(fù)雜度O（k）），代表工作：Poly-BERT，PQ-BERT。

總結(jié)這類工作的主要思想是增強(qiáng)雙塔模型的向量表示能力，由于document較長，可能對應(yīng)多種語義，而原始雙塔模型對query和document只使用一個向量表示，可能造成語義缺失。那么可以使用多個向量來表示document，在線階段再進(jìn)行一些優(yōu)化來加速模型的推斷。

二、訓(xùn)練樣本構(gòu)造

檢索任務(wù)中，相對于整體document庫，每個query所對應(yīng)的相關(guān)document是很少的一部分。在訓(xùn)練時，模型往往只接收query對應(yīng)的相關(guān)文檔（正樣本）以及少量query的不相關(guān)文檔（負(fù)樣本），目標(biāo)函數(shù)是區(qū)分正樣本和負(fù)樣本。然而在模型推斷時，模型需要對document庫中的所有document進(jìn)行打分。如果模型在訓(xùn)練時讀取的document和document庫中的一些document之間的語義距離相差較大，則可能造成模型在推斷階段表現(xiàn)不佳。因此，如何構(gòu)造訓(xùn)練樣本是一個重要的研究方向。

方法一：（1） 首先介紹一個比較簡單的trick，In-batch negatives。顧名思義，在訓(xùn)練時，假設(shè)一個batch中包含b個query，每個query（q_i）都有一個對應(yīng)的正樣本dp_i和負(fù)樣本dq_i，那么在這個batch中，每個q_i除了自己所對應(yīng)的負(fù)樣本，還可以將batch中其他query所對一個的正樣本和負(fù)樣本都作為當(dāng)前query所對應(yīng)的負(fù)樣本，大大提高了訓(xùn)練數(shù)據(jù)的利用率。實驗表明，該trick在各種檢索任務(wù)上都能提高模型的效果。

方法二：（2） 上述方法的目標(biāo)是在訓(xùn)練過程中利用更多的負(fù)樣本，讓模型的魯棒性更強(qiáng)。然而訓(xùn)練過程能遍歷的負(fù)樣本始終是有限的，那么如何在有限的訓(xùn)練樣本中構(gòu)造更有利于模型訓(xùn)練的負(fù)樣本是一個重要的研究問題。

ANCE提出了一種迭代式生成負(fù)樣本的思路：隨著訓(xùn)練的進(jìn)行，模型對于文本的表示也會變化，之前對于模型較難的負(fù)樣本可能變得不那么難，而之前沒見過的負(fù)樣本對于模型可能會較難區(qū)分。該工作以此為出發(fā)點，同時進(jìn)行train和inference，在訓(xùn)練的同時，利用上一個checkpoint中的模型進(jìn)行inference，對訓(xùn)練數(shù)據(jù)生成新的負(fù)樣本，在inference完成后，使用新的負(fù)樣本進(jìn)行訓(xùn)練。這樣可以漸進(jìn)的訓(xùn)練模型，保持負(fù)樣本的難度，更充分的訓(xùn)練模型。

方法三：（3） 除了利用模型本身來生成負(fù)樣本，還可以利用比雙塔模型復(fù)雜的交互模型來生成訓(xùn)練數(shù)據(jù)。RocketQA提出了基于交互模型來增強(qiáng)數(shù)據(jù)的方法。由于交互模型的表現(xiàn)更強(qiáng)，作者使用交互模型來標(biāo)注可能成為正樣本的文檔（這些文檔未經(jīng)過標(biāo)注），以及篩選更難的訓(xùn)練雙塔模型的樣本。具體的訓(xùn)練過程如下圖所示：

總結(jié)：這類工作從訓(xùn)練數(shù)據(jù)著手，彌補(bǔ)原有的訓(xùn)練模式對于缺少負(fù)樣本優(yōu)化的不足。個人角度認(rèn)為這類工作提升可能更為顯著。

三、訓(xùn)練目標(biāo)改進(jìn)

訓(xùn)練目標(biāo)上的改進(jìn)比較靈活，有多種不同的改進(jìn)方式，首先介紹利用交互模型改進(jìn)雙塔模型的工作。

相對于雙塔模型，交互模型的表現(xiàn)更好，但復(fù)雜度更高，因此很多工作的idea是通過模型蒸餾將交互模型的文本表示能力遷移到雙塔模型中，這方面有很多類似的工作。這里選取一個SIGIR2021的最新文章作為代表。

如上圖所示，該模型不僅蒸餾了一個交互模型（monoBERT），同時還蒸餾了一個基于雙塔的改進(jìn)模型ColBERT。該模型使用monoBERT作為teacher，對模型的CLS位置向量進(jìn)行蒸餾，使用ColBERT作為teacher，對模型的除了［CLS］位置的向量進(jìn)行蒸餾，目標(biāo)函數(shù)為以下三部分的加和：

最后的打分函數(shù)是monoBERT和ColBERT的組合，即，首先使用document和query的CLS位置向量輸入MLP，輸出一個分?jǐn)?shù)，同時使用document和query的其他位置表示向量輸入到和ColBERT相同的打分函數(shù)中，最后使用兩個分?jǐn)?shù)的和作為最后打分。

前文所述的工作都是將query和document的文本映射到稠密向量空間中，然后進(jìn)行匹配。另外還有的工作是直接利用文字進(jìn)行匹配。Doc2query使用一個基于seq2seq的預(yù)訓(xùn)練語言模型（比如T5），利用標(biāo)注的document，query對進(jìn)行finetune，目標(biāo)是輸入document輸出對應(yīng)的query，然后將輸出的query和document本身進(jìn)行拼接，擴(kuò)展document。然后利用傳統(tǒng)的檢索方法，比如BM25，對擴(kuò)展過的document建立索引并查找。過程示意如下圖所示。在MSMARCO上的一些實驗表明，這個方法可以和基于向量的搜索一起使用，提高模型的表現(xiàn)。

四、雙塔模型預(yù)訓(xùn)練

一般的預(yù)訓(xùn)練模型使用的目標(biāo)函數(shù)主要是MLM或者seq2seq，這種預(yù)訓(xùn)練目標(biāo)和雙塔的匹配任務(wù)還是有一些不一致。并且已有的預(yù)訓(xùn)練模型即使有匹配任務(wù)（比如NSP），也是使用交互而非雙塔的工作方式。為了使用預(yù)訓(xùn)練提高雙塔模型的效果，SimCSE通過對比學(xué)習(xí)，提升模型對句子的表示能力。

該方法的實現(xiàn)很簡單，假設(shè)提取一個batch的句子，通過模型自帶的dropout，將每個句子輸入到預(yù)訓(xùn)練模型中，dropout兩次，將同一個句子dropout后的結(jié)果作為正樣本，不同句子的dropout結(jié)果作為負(fù)樣本，拉近正樣本的距離，拉遠(yuǎn)負(fù)樣本的距離，每個句子的向量由BERT的CLS位置向量表示。如下圖所示：

模型雖然很簡單，但是在句子匹配任務(wù)上取得了很好的效果。該模型在檢索任務(wù)上的效果還需實驗。

還有的工作是針對檢索任務(wù)的預(yù)訓(xùn)練。ICLR2020一篇論文Pre-training Tasks for Embedding-based Large-scale Retrieval提出了一些預(yù)訓(xùn)練任務(wù)，這些任務(wù)主要是針對Wikipedia的，不一定具有普適性。如下圖所示，紫色d框出來的代表document，q1，q2，q3代表不同任務(wù)構(gòu)造的的query，q1是ICT，即利用document所在的一句話作為query，q2是BFS，即利用document所在網(wǎng)頁的第一段中的一句話作為query，q3是WLP，使用document中的某個超鏈接頁面的第一句話作為query。任務(wù)目標(biāo)是匹配q1，q2，q3和d。

Condenser

傳統(tǒng)的MLM預(yù)訓(xùn)練任務(wù)如下圖所示，該任務(wù)沒有特別強(qiáng)制訓(xùn)練CLS位置的向量表示能力。為了將整個序列的信息壓縮到CLS位置上，Condenser將模型分成兩部分，第一部分和普通的Transformer一樣，第二部分使用經(jīng)過更多交互后的［CLS］位置向量（黃色部分）來預(yù)測［MASK］的token，強(qiáng)制模型的［CLS］編碼可以具有還原其他token的能力。

編輯：jq