本文介紹了本小組發表于ICLR 2023的論文UniKGQA，其基于預訓練語言模型設計了一套統一的模型架構，同時適用于多跳KBQA檢索和推理，在多個KBQA數據集上取得顯著提升。

該論文發表于 ICLR-2023 Main Conference：

論文鏈接：https://arxiv.org/pdf/2212.00959.pdf

開源代碼：https://github.com/RUCAIBox/UniKGQA

進NLP群—>加入NLP交流群(備注nips/emnlp/nlpcc進入對應投稿群)

前言

如何結合PLM和KG以完成知識與推理仍然是一大挑戰。我們在NAACL-22關于常識知識圖譜推理的研究 (SAFE) 中發現，PLM 是執行復雜語義理解的核心。因此，我們深入分析了已有的復雜 GNN 建模外部 KG 知識的方法是否存在冗余。最終，基于發現，我們提出使用純 MLP 輕量化建模輔助 PLM 推理的 KG 知識，初步探索了 PLM+KG 的使用方法。

進一步，本文研究了在更依賴 KG 的知識庫問答任務中如何利用 PLM。已有研究通常割裂地建模檢索-推理兩階段，先從大規模知識圖譜上檢索問題相關的小子圖，然后在子圖上推理答案節點，這種方法忽略了兩階段間的聯系。我們重新審視了兩階段的核心能力，并從數據形式，模型架構，訓練策略三個層面進行了統一，提出UniKGQA。同時受 SAFE 啟發，我們認為 KG 僅為執行推理的載體，因此 UniKGQA 架構的設計思考為：核心利用 PLM 匹配問題與關系的語義，搭配極簡 GNN 在 KG 上傳播匹配信息，最終推理答案節點。針對這樣的簡潔架構，我們同時設計了一套高效的訓練方法，使得 UniKGQA 可以將檢索的知識遷移到推理階段，整體性能更高效地收斂到更好的表現。實驗證明，在多個標準數據集上相較于已有 SOTA，取得顯著提升。

一、研究背景與動機

1、多跳知識庫問答

給定一個自然語言問題 和一個知識圖譜 ，知識圖譜問答 (KGQA) 旨在從知識圖譜上尋找答案集合，表示為 。我們在圖1 (a)中展示了一個例子。給定問題：Who is the wife of the nominee for The Jeff Probst Show?，該任務的目標是從主題實體 The Jeff Probst Show 開始，尋找匹配問題語義的推理路徑 nominee --> spouse，最終得到答案實體 Shelley Wright 和 Lisa Ann Russell。已有研究通常假設問題中提到的實體 (例如圖1 (a)中的The Jeff Probst Show) 被標記并鏈接到知識圖譜上，即主題實體，表示為 。

本文關注多跳 KGQA 任務，即答案實體和主題實體在知識圖譜上距離多跳。考慮到效率和精度之間的平衡，我們遵循已有工作，通過檢索-推理兩階段框架解決此任務。具體而言，給定一個問題 和主題實體 ，檢索階段旨在從超大知識圖譜 中檢索出一個小的子圖 ，而推理階段則在檢索子圖 上推理答案實體 。

2、研究動機

雖然兩個階段的目的不同，但是兩個階段都需要評估候選實體與問題的語義關聯性 (用于檢索階段remove或推理階段rerank)。本質上，上述過程可以被視為一個語義匹配問題。由于 KG 中實體與實體間的聯系通過關系表示，為了衡量實體的相關性，在構建語義匹配模型時，基于關系的特征(直接的一跳關系或復合的多跳關系路徑)，都被證明是特別有用的。如圖1 (a)所示，給定問題，關鍵是要在知識圖譜中識別出與問題語義匹配的關系及其組成的關系路徑 (例如nominee --> spouse)，以找到答案實體。

由于兩個階段處理知識圖譜時，面臨的搜索空間尺度不同 (例如，檢索時數百萬個實體結點與推理時數千個實體結點)，已有方法通常割裂地為兩階段考慮對應的解決方案：前者關注如何使用更高效的方法提升召回性能，而后者關注如和利用更細粒度的匹配信號增強推理。這種思路僅將檢索到的三元組從檢索階段傳遞到推理階段，而忽略了整個流程中其他有用的語義匹配信號，整體性能為次優解。由于多跳知識圖譜問答是一項非常具有挑戰性的任務，我們需要充分利用兩個階段習得的各種能力。

因此，本文探討能否設計一個統一的模型架構來為兩個階段提供更好的性能？如果這樣，我們可以緊密關聯兩階段并增強習得能力的共享，從而提升整體性能。

二、UniKGQA：適用于檢索和推理的統一架構

然而，實現統一的多跳 KGQA 模型架構面臨兩個主要挑戰: (1) 如何應對兩個階段的搜索空間尺度差異很大的問題? (2) 如何在兩個階段之間有效地共享或傳遞習得的能力? 考慮到這些挑戰，我們從數據形式，模型架構，訓練策略三方面進行探索，最終對兩階段的模型架構進行了統一。

1、數據形式

在 KG 中，存在大量的一對多現象，例如，頭實體為中國，關系為城市，那么存在多個尾實體，每個尾實體又會存在各自的一對多現象，使得圖的規模隨跳數成指數級增長。實際上，在檢索階段，我們僅需要通過關系或關系路徑召回一批相關的實體，而不需要細粒度關注實體本身的信息。結合以上思考，我們針對檢索階段提出了Abstract Subgraph(抽象子圖) 的概念，核心是將同一個頭實體和關系派生出的尾實體聚合在一起，得到對應的抽象結點，如圖1 (b)即為圖1 (a)的抽象子圖表示，這樣可以顯著降低原始知識圖譜的規模。因此，檢索階段通過關系或關系路徑判斷抽象節點的相關性，檢索完成后，將含有抽象節點的子圖進行還原，得到包含原始節點的子圖；推理階段通過關系或關系路徑同時考慮具體的節點信息推理最終的答案節點。這樣，我們就可以減緩兩個階段面臨的搜索空間尺度過大的問題。

基于抽象子圖，我們針對兩階段提出一個評估實體相關性的通用形式，即給定問題 和候選實體的子圖 。對于檢索階段， 是抽象子圖，包含抽象節點以合并同一關系派生的尾實體。對于推理階段， 是基于檢索階段的檢索子圖構建的，還原后沒有抽象節點。這種通用的輸入格式為開發統一的模型架構提供了基礎。接下來，我們將以一般方式描述針對這種統一數據形式設計的模型架構，而不考慮特定的檢索或推理階段。

2、模型架構

基于上述統一數據形式，我們開發的模型架構包含兩個基礎模塊：(1) 語義匹配 (SM) 模塊，利用 PLM 執行問題和關系之間的語義匹配；(2) 匹配信息傳播 (MIP) 模塊，在知識圖譜上傳播語義匹配信息。我們在圖2中展示了模型架構的概覽。

語義匹配 (SM)：SM 模塊旨在生成問題 與給定子圖 中的三元組間的語義匹配特征。具體而言，我們首先利用 PLM 對和的文本進行編碼，然后使用 [CLS] 令牌的輸出表示作為它們的表示:

基于 和 ，受 NSM 模型的啟發，我們通過對應的投影層，在第 步獲得問題 和三元組間語義匹配特征的向量間語義匹配特征的向量:

其中，, 是第 步投影層的參數， 和 分別是 PLM 和特征向量的隱藏層維度， 是 sigmoid 激活函數，而 是 hadamard 積。

匹配信息傳播 (MIP)：基于語義匹配特征，MIP 模塊首先將它們聚合起來以更新實體表示，然后利用它來獲取實體匹配得分。為了初始化匹配得分，對于給定問題 和子圖 中的每個實體 ，我們將 和 之間的匹配分數設置為：如果 是主題實體，則 ，否則 。在第 步，我們利用上一步計算出的頭實體的匹配分數作為權重，聚合相鄰三元組的匹配特征，以獲得尾實體的表示：

其中， 是第 步中實體 的表示， 是可學習的矩陣。在第一步中，由于沒有匹配分數，我們按照NSM模型的方法，直接將其一跳關系的表示聚合為實體表示：，其中 是可學習的矩陣。基于所有實體 的表示，我們使用 softmax 函數更新它們的實體匹配分數：

其中， 是一個可學習的向量。

經過 步迭代，我們可以獲得最終的實體匹配得分 ，它是子圖 中所有實體的概率分布。這些匹配分數可以用來衡量實體作為給定問題 答案的可能性，并將在檢索和推理階段中使用。

訓練策略

我們在多跳知識圖譜問答的推理和檢索階段都使用了前述的模型架構，分別為推理模型和檢索模型。由于這兩個模型采用相同的架構，我們引入 和 來分別表示用于檢索和推理的模型參數。如前所述，我們的架構包含兩組參數，即基礎 PLM 以及用于匹配和傳播的其他參數。因此， 和 可以分解為 和 ，其中下標 和 分別表示我們架構中的 PLM 參數和其他參數。為了學習這些參數，我們基于統一架構設計了預訓練 (即問題-關系匹配)和微調 (即面向檢索和推理的學習)策略。下面，我們描述模型訓練方法。

問題-關系匹配的預訓練 (QRM): 對于預訓練，我們主要關注學習基礎 PLMs (即 和 ) 的參數。在實現中，我們讓兩個模型共享相同的 PLM 參數，即 。語義匹配模塊的基本功能是對一個問題和一個單獨的關系進行相關性建模 (式2)。因此，我們設計了一個基于問題-關系匹配的對比預訓練任務。具體來說，我們采用對比學習目標來對齊相關問題-關系對的表示，同時將其他不相關的對分開。為了收集相關問題-關系對，對于一個由問題 、主題實體 和答案實體 組成的例子，我們從整個知識圖譜中提取 和 之間的所有最短路徑，并將這些路徑中的所有關系視為與 相關的關系，表示為 。這樣，我們就可以獲得許多弱監督樣例。在預訓練期間，對于每個問題 ，我們隨機采樣一個相關的關系 ，并利用對比學習損失進行預訓練:

其中，是一個溫度超參數，是一個隨機采樣的負關系，是余弦相似度，、是由SM模塊（式1）中的 PLM 編碼的問題和關系。這樣，通過預訓練 PLM 參數，問題-關系匹配能力將得到增強。請注意，在預訓練之后，PLM 參數將被固定。

在抽象子圖上微調檢索 (RAS)：在預訓練之后，我們在檢索任務上學習參數 。回憶一下，我們將子圖轉化為一種抽象子圖的形式，其中包含抽象節點，用于合并來自同一關系派生的尾實體。由于我們的 MIP 模塊可以生成子圖中節點的匹配分數 (式4)，其中下標 表示節點來自抽象子圖。此外，我們利用標注的答案來獲取標簽向量，表示為 。如果抽象節點中包含答案實體，則在 中將抽象節點設置為1。接下來，我們最小化學習匹配得分向量和標簽向量之間的KL散度，如下式所示：

通過RAS損失微調后，可以有效地學習檢索模型。我們通過它們的匹配得分選擇排名前 個節點，利用它們來檢索給定問題 的子圖。請注意，僅選擇與主題實體距離合理的節點進入子圖，這可以確保推理階段使用的子圖 相對較小但與問題相關。

在檢索子圖上微調推理 (RRS)：在微調檢索模型后，我們繼續微調推理模型，學習參數 。通過微調后的檢索模型，我們可以獲得每個問題 的較小子圖 。在推理階段，我們專注于執行準確的推理，以找到答案實體。因此，我們還原抽象節點中的原始節點及其原始關系。由于檢索和推理階段高度依賴，我們首先使用檢索模型的參數來初始化推理模型的參數： 。然后，根據式4，我們采用類似的方法使用KL損失函數來使學習到的匹配得分 (表示為 ) 擬合標簽向量 (表示為 )：

其中，下標 表示節點來自檢索子圖。通過RRS損失的微調后，我們可以利用學習的推理模型選擇排名前個實體。

如圖1 (c)所示，整體的訓練過程由以下三個步驟組成：(1) 與 共享參數，(2) 使用問題-關系匹配預訓練 ，(2) 使用抽象子圖微調 以進行檢索，(3) 使用子圖微調 以進行推理，其中 使用 進行初始化。

討論

據我們所知，我們是KGQA領域首次提出使用統一模型在檢索和推理階段共享推理能力。在表格1中，我們總結了我們的方法和幾種流行的多跳知識庫問答方法（包括 GraphfNet、PullNet、NSM 和 SR+NSM 之間的區別。我們可以看到，現有方法通常針對檢索和推理階段采用不同的模型，而我們的方法更為統一。統一帶來的一個主要優點是，兩個階段之間的信息可以有效地共享和復用，即，我們使用學習的檢索模型來初始化推理模型。

三、實驗結果

1、主實驗

我們在3個公開的文檔檢索數據集上進行實驗，分別是 WebQuestionsSP (WebQSP)、Complex WebQuestions 1.1 (CWQ)、和 MetaQA 數據集。實驗結果如下表所示，通過對比可以清晰地看出我們的方法的優勢。例如，在難度較大的數據集 WebQSP 和 CWQ 上，我們的方法遠遠優于現有的最先進基線（例如，WebQSP 的 Hits@1 提高了8.1％，CWQ 的 Hits@1 提高了2.0％）。

在我們的方法中，為了提高效率，我們固定了基于 PLM 的編碼器的參數。實際上，更新其參數可以進一步提高模型性能。這樣的方法使研究人員在實際應用中可以權衡效率和精度。因此，我們提出了兩種 UniKGQA 的變體來研究它：（1） 僅在編碼問題時更新 PLM 編碼器的參數，（2） 同時在編碼問題和關系時更新 PLM 編碼器的參數。事實上，這兩種變體都可以提高我們的 UniKGQA 的性能。只在編碼問題時更新 PLM 編碼器可以獲得與同時更新兩者相當甚至更好的性能。可能的原因是在編碼問題和關系時更新 PLM 編碼器可能會導致過度擬合下游任務。因此，僅在編碼問題時更新PLM 編碼器是更有價值的，因為它可以在相對較少的額外計算成本下實現更好的性能。

2、深入分析

檢索性能：我們從三個方面評估檢索性能：子圖大小、答案覆蓋率和最終 QA 性能。可以看到，在檢索出相同大小的子圖的情況下，UniKGQA 和 SR 的答案覆蓋率顯著高于 PPR 的。這證明了訓練可學習的檢索模型的有效性和必要性。此外，盡管 UniKGQA 和 SR 的曲線非常相似，但我們的 UniKGQA 比 SR+NSM 可以實現更好的最終 QA 性能。原因是 UniKGQA 可以基于統一體系結構將相關信息從檢索階段傳遞到推理階段，學習更有效的推理模型。這一發現可以通過將我們的 UniKGQA 與 UniKGQA+NSM 進行比較來進一步驗證。

微調效率：我們比較了 UniKGQA 和較強基線模型 NSM 在相同檢索的子圖上進行微調時，性能隨迭代輪數的變化。如圖3右側展示。首先，我們可以看到，在微調之前（即迭代輪數為零時），我們的 UniKGQA 已經達到了與 NSM 最佳結果相當的性能。這表明推理模型已經成功利用了檢索模型習得的知識，可以進行一定的推理。迭代兩輪之后，我們的 UniKGQA 已經達到接近收斂的性能。表明我們的模型可以實現高效的微調。

消融實驗：我們提出兩個重要的訓練策略以提高性能：(1) 使用問題-關系匹配進行預訓練，(2) 使用檢索模型的參數初始化推理模型。我們通過消融實驗驗證它們的有效性。我們提出了三種變體：(1) 去除預訓練過程, (2) 去除使用檢索模型參數初始化，(3) 同時去除預訓練和初始化過程。我們在表格4中展示了消融研究的結果。可以看到，所有這些變體的性能都低于完整的 UniKGQA，這表明這兩個訓練策略對最終性能都很重要。此外，這種觀察還驗證了我們的 UniKGQA 確實能夠轉移和重用習得的知識以提高最終性能。

四、總結

在這項工作中，我們提出了一種多跳知識圖譜問答任務新的模型架構。作為主要技術貢獻，UniKGQA 引入了基于 PLMs 的統一模型架構，可同時適用于檢索階段與推理階段。為了應對兩個階段的不同搜索空間規模，我們提出了檢索階段專用的抽象子圖的概念，它可以顯著減少需要搜索的節點數量。此外，我們針對統一模型架構，設計了一套高效的訓練策略，包含預訓練（即問題-關系匹配）和微調（即面向檢索和推理的學習）。得益于統一的模型架構，UniKGQA 可以有效增強兩個階段之間習得能力的共享和轉移。我們在三個基準數據集上進行了廣泛的實驗，實驗結果表明，我們提出的統一模型優于競爭方法，尤其是在更具挑戰性的數據集（WebQSP 和 CWQ）上表現更好。

審核編輯 ：李倩