基于訓(xùn)練階段使用知識庫+KNN檢索相關(guān)信息輔助學(xué)習(xí)方法

在prompt learning中一個(gè)核心問題是模型存在死記硬背現(xiàn)象。Prompt learnin主要應(yīng)用在few-shot learning場景，先將訓(xùn)練數(shù)據(jù)轉(zhuǎn)換成prompt的形式，在訓(xùn)練過程模型側(cè)重于記憶訓(xùn)練數(shù)據(jù)，然后使用記憶的信息做預(yù)測。這個(gè)過程會(huì)導(dǎo)致模型缺乏泛化能力，一些長尾的case預(yù)測效果不好。

NIPS 2022中浙大和阿里提出使用檢索方法增強(qiáng)prompt learning，利用訓(xùn)練數(shù)據(jù)構(gòu)造知識庫，在訓(xùn)練階段使用知識庫+KNN檢索相關(guān)信息輔助學(xué)習(xí)，通過這種方式將需要記憶的信息從模型中拆分出來，直接輸入到模型中。通過這種方式，可以讓模型參數(shù)更側(cè)重泛化信息的學(xué)習(xí)，而不是過擬合訓(xùn)練數(shù)據(jù)。下面為大家詳細(xì)介紹一下這篇工作。

NLP Prompt系列——Prompt Engineering方法詳細(xì)梳理

1 Prompt Learning回顧

Prompt learning主要面向的是訓(xùn)練數(shù)據(jù)較少的場景。首先需要一個(gè)預(yù)訓(xùn)練模型，然后將下游任務(wù)轉(zhuǎn)換成完形填空的形式。對于分類問題，判斷某個(gè)text屬于哪個(gè)label，轉(zhuǎn)換成如下的文本輸入到預(yù)訓(xùn)練語言模型中：

利用預(yù)訓(xùn)練語言模型，預(yù)測出[MASK]對應(yīng)的文本，后面接一個(gè)文本到label的映射函數(shù)，即可實(shí)現(xiàn)文本分類任務(wù)。Prompt learning的好處是可以充分利用預(yù)訓(xùn)練語言模型的知識，讓下游任務(wù)和預(yù)訓(xùn)練任務(wù)更加適配，以提升樣本量不足情況下的效果。我在之前的文章NLP Prompt系列——Prompt Engineering方法詳細(xì)梳理詳細(xì)介紹過prompt相關(guān)工作，感興趣的同學(xué)可以進(jìn)一步閱讀。

雖然這種方法充分運(yùn)用了預(yù)訓(xùn)練語言模型的知識，但畢竟finetune的數(shù)據(jù)少，模型更像在死記硬背訓(xùn)練prompt數(shù)據(jù)中的信息。這對于長尾樣本或非典型的句子的預(yù)測效果不友好。為了解決上述問題，本文的核心思路是，如果我們把這些需要記憶的信息單獨(dú)拿出來存儲(chǔ)到一個(gè)知識庫中，在需要的時(shí)候檢索它們并作為模型額外輸入，就能讓模型參數(shù)沒必要再死記硬背這些信息了，從而實(shí)現(xiàn)記憶和泛化更好的平衡，有點(diǎn)【好記性不如爛筆頭】的感覺。下圖是本文提出的基本框架示意圖。

2 從知識庫中檢索信息

實(shí)現(xiàn)上面的框架核心是從知識庫中檢索信息，這也就涉及到兩個(gè)問題，一個(gè)是如何構(gòu)造知識庫，另一個(gè)是如何進(jìn)行信息檢索和利用。

在知識庫的構(gòu)造上，文中構(gòu)造的是一個(gè){K, V}格式的數(shù)據(jù)，訓(xùn)練集中的每條樣本對應(yīng)一個(gè){K, V}。K代表這個(gè)樣本的prompt輸入模型后[MASK]位置的隱向量，V代表這個(gè)樣本的label對應(yīng)的單詞。由于K是模型輸出的向量，因此每訓(xùn)練幾輪，就會(huì)動(dòng)態(tài)更新知識庫中的Key，避免Key和模型最新參數(shù)隔代太多不匹配。

在信息檢索和利用上，對于當(dāng)前樣本模型先得到其[MASK]位置的向量，然后用這個(gè)向量在知識庫中進(jìn)行KNN檢索，每個(gè)類別的樣本都取出topK個(gè)，檢索的距離度量是向量內(nèi)積。對于每個(gè)類別檢索出的向量，使用內(nèi)積做softmax后進(jìn)行加權(quán)融合，得到這個(gè)類別最終向量，拼接到當(dāng)前樣本embedding后面輸入到模型中：

這部分檢索出來的樣本就是將需要記憶的知識直接引入到了當(dāng)前樣本中，不再需要模型參數(shù)去記憶了。此外，這種將向量引入而不是引入對應(yīng)的token，可以讓信息的擴(kuò)展更方便，直接引入樣本的token會(huì)拉長輸入樣本長度，導(dǎo)致模型性能下降，且長度也有上限。

3 使用KNN指導(dǎo)模型訓(xùn)練和預(yù)測

上面收的引入知識庫+KNN的方法，緩解了模型參數(shù)需要強(qiáng)記憶訓(xùn)練樣本的問題。此外，文中還通過KNN檢索結(jié)果來指導(dǎo)模型的學(xué)習(xí)過程。KNN檢索的好處是不需要模型訓(xùn)練，直接根據(jù)預(yù)訓(xùn)練的表示計(jì)算距離，利用鄰居樣本的label，就能預(yù)測當(dāng)前樣本的label。這對于模型來說是另一個(gè)維度的信息補(bǔ)充，文中通過區(qū)分難樣本指導(dǎo)訓(xùn)練和在inference階段指導(dǎo)預(yù)測兩個(gè)方面進(jìn)一步指導(dǎo)模型的訓(xùn)練和預(yù)測。

KNN的檢索結(jié)果可以用來區(qū)分難樣本和簡單樣本。通過KNN檢索以及檢索鄰居的label，可以得到當(dāng)前樣本各個(gè)類別的預(yù)測概率。這個(gè)KNN的預(yù)測結(jié)果可以作為是否是難樣本的參考，如果模型預(yù)測預(yù)測結(jié)果和KNN結(jié)果不一致，就是難樣本。對于難樣本，加大其學(xué)習(xí)權(quán)重，通過將KNN預(yù)測概率引入到交叉熵?fù)p失中實(shí)現(xiàn)：

另一方面，在預(yù)測階段，也直接將KNN的預(yù)測結(jié)果拿出來和模型對于[MASK]的預(yù)測結(jié)果做插值，得到最終的預(yù)測結(jié)果：

4 實(shí)驗(yàn)結(jié)果

文中在9個(gè)NLU數(shù)據(jù)集的few-shot和zero-shot learning上對比了效果，可以看到本文提出的方法對于效果的提升還是非常明顯的。

此外，文中也對比了跨領(lǐng)域的效果，在source domain進(jìn)行prompt learning，對比在target domain上的效果：

5 總結(jié)

檢索在NLP各類任務(wù)中的應(yīng)用越來越多，本文也將檢索用于分離可記憶的信息來提升模型的泛化能力，并取得了顯著效果。檢索通過信息記憶+查詢的方式，引入了豐富的外部信息，能夠讓模型更多的容量服務(wù)于學(xué)習(xí)泛化性，而非簡單的記住訓(xùn)練數(shù)據(jù)。