如何讓大模型感知知識圖譜知識？

螞蟻聯(lián)合實驗室：利用多詞元并行預(yù)測給它“上課”。

大語言模型的飛速發(fā)展打破了許多自然語言處理任務(wù)間的壁壘。通常情況下，大語言模型以預(yù)測下一個詞元（Token）為訓(xùn)練目標(biāo)，這與許多自然語言處理任務(wù)十分契合。

但對于知識圖譜而言，實體作為最基本的數(shù)據(jù)單元，往往需要多個自然語言詞元才能準(zhǔn)確描述，這導(dǎo)致知識圖譜與自然語言之間存在明顯的粒度不匹配。

為了解決這一問題，螞蟻團隊提出了一種基于大語言模型的多詞元并行預(yù)測方法K-ON，其利用多詞元并行預(yù)測機制能夠一次生成對所有實體的評估結(jié)果，進而實現(xiàn)語言模型實體層級的對比學(xué)習(xí)。

其結(jié)果收錄于AAAI 2025 Oral。論文一作目前在浙江大學(xué)攻讀博士。

實驗結(jié)果表明，本文方法在多個數(shù)據(jù)集上的知識圖譜補全任務(wù)中均優(yōu)于現(xiàn)有方法。

基于多詞元并行預(yù)測的實體對比學(xué)習(xí)

詞元是語言模型所能處理的最基本元素，通常需要數(shù)個詞元組成的文本標(biāo)簽才能準(zhǔn)確描述和鑒別知識圖譜中的實體。雖然為每個實體創(chuàng)建一個新的詞元并在微調(diào)過程中學(xué)習(xí)這些詞元的表示不失為一種替代方案，但這種方式訓(xùn)練調(diào)優(yōu)成本較高，且可能會對大模型的性能產(chǎn)生負(fù)面影響，通用性也受到限制。

本文探討了如何高效利用多個詞元描述知識圖譜中的實體以解決知識圖譜相關(guān)問題的方法。首先，直接優(yōu)化經(jīng)典的序列預(yù)測損失可能會導(dǎo)致大模型缺乏對知識圖譜實體的認(rèn)識，從而出現(xiàn)生成知識圖譜中不存在的實體的問題；且考慮到知識圖譜中實體的數(shù)量，將所有實體以文本上下文的方式輸入給大模型顯然也是不現(xiàn)實的。

以上圖為例，假設(shè)任務(wù)是給定不完整三元組以預(yù)測目標(biāo)實體Matt Damon。左圖中使用常規(guī)連續(xù)單詞元預(yù)測方式生成結(jié)果需要多個子步驟，且無法直接處理多個實體。因此，現(xiàn)有大多數(shù)知識圖譜相關(guān)方法僅將大模型應(yīng)用于簡單任務(wù)上，如驗證三元組的正確性或從有限數(shù)量的候選實體中選擇正確答案。

相比之下，本文提出的K-ON方法使用K個輸出層并行預(yù)測多個實體不同位置詞元的概率，這與目前DeepSeek等大模型中使用的多詞元預(yù)測技術(shù)有著一定的相似性，且本文方法進一步借助了實體層級的對比學(xué)習(xí)在模型輸出層上累加知識圖譜知識。

K-ON完成知識圖譜補全任務(wù)的方法論

如下圖所示，K-ON并行評估知識圖譜候選實體分?jǐn)?shù)的過程可分為五步：

(1)與現(xiàn)有微調(diào)大模型的方法相似，K-ON將知識圖譜補全問題以文本指令的方式輸入大模型；

(2) 經(jīng)大模型Transformer模塊處理后的輸出狀態(tài)被輸入至K-ON模塊中，該模塊由多個原大模型輸出層MLP構(gòu)成，對應(yīng)為要預(yù)測實體的不同位置的詞元；

(3) 接著，K-ON使用Conditional Transformer混合不同位置的信息，并考慮到詞元前后的順序依賴性；

(4) 然后，使用低秩適應(yīng)技術(shù)（LoRA）將原大模型評分層構(gòu)造為K個新的評分層，從而把上一步的輸出結(jié)果轉(zhuǎn)換為對實體K個連續(xù)詞元的概率預(yù)測分布；

(5) 最后便可以從不同位置的概率預(yù)測分布中抽取各實體詞元對應(yīng)的概率值，進而一次評估所有候選實體的分?jǐn)?shù)。

在獲取候選實體分?jǐn)?shù)后，便可使用知識圖譜表示學(xué)習(xí)領(lǐng)域中最為常用的對比學(xué)習(xí)損失使大模型掌握知識圖譜中實體的分布：

此處 p^e、 p^ej分別代表正樣本和負(fù)樣本的分?jǐn)?shù)，均由K-ON模塊并行生成。除了實體層級的對比學(xué)習(xí)外，本文還進一步考慮使用詞元序列對齊使多詞元并行預(yù)測的結(jié)果與原本大模型單步連續(xù)預(yù)測的結(jié)果相接近。為實現(xiàn)這一目標(biāo)，本文首先引入常用的單步預(yù)測損失以在訓(xùn)練語料上微調(diào)原輸出層參數(shù)：

這里的下標(biāo)k指代組成實體的詞元的序號。

接著便可令K-ON模塊中一次評估的K個詞元的概率分布與常規(guī)連續(xù)單詞元預(yù)測得到的K個概率分布對齊：

此處上標(biāo) k-on、llm 分別指代K-ON和常規(guī)連續(xù)預(yù)測所得到的分?jǐn)?shù)。

最后，訓(xùn)練K-ON完成知識圖譜補全任務(wù)的基本流程可總結(jié)如下：

實驗結(jié)果：效率更高、成本更低、效果更好

下表展示K-ON在知識圖譜補全任務(wù)上的實驗結(jié)果，除傳統(tǒng)方法外，本文還與同樣基于大模型的方法以及多模態(tài)方法進行了比較。不難看出，K-ON在所有數(shù)據(jù)集及指標(biāo)上均取得了優(yōu)于現(xiàn)有方法的結(jié)果，且與一些使用額外圖像數(shù)據(jù)的方法相比，仍具有一定優(yōu)勢。

下圖中進一步分析了K-ON中隨著K值的增加，模型各方面性能的變化，這里K的取值直接決定了單個實體最多可以由多少詞元表達。如圖所示，當(dāng)K取值過小時，由于表達能力不足，K-ON取得的實驗結(jié)果很差，但增加至8以后帶來的性能提升已不明顯，而模型可訓(xùn)練參數(shù)量卻穩(wěn)步上升。

值得注意的是，推理所用單步時間及總訓(xùn)練時間受K值影響不大，這說明了K-ON多詞元并行預(yù)測的高效性。

不僅如此，本文還對K-ON所實現(xiàn)的實體層級的對比學(xué)習(xí)進行了分析，如下圖所示?？梢钥闯?，在幾乎不對訓(xùn)練效率造成影響的前提下，K-ON可輕易實現(xiàn)涉及上千個負(fù)樣本實體的對比學(xué)習(xí)，但負(fù)樣本數(shù)量并不是越多越好，將其設(shè)為128個左右便可取得最優(yōu)結(jié)果。

本文提出了一種多詞元并行預(yù)測方法，通過實體層級的對比學(xué)習(xí)使大模型能夠感知知識圖譜知識。充分的實驗結(jié)果表明，本文方法在知識圖譜相關(guān)任務(wù)上具有顯著性能優(yōu)勢，并且較常規(guī)大模型方案具有更高的訓(xùn)練與推理效率。

論文地址：
https://arxiv.org/pdf/2502.06257

螞蟻有18篇技術(shù)論文被收錄

當(dāng)?shù)貢r間2月25日，AAAI 2025將在美國賓夕法尼亞州費城舉辦，會議為期11天，于3月4日結(jié)束。AAAI 由國際人工智能促進協(xié)會主辦，為人工智能領(lǐng)域的頂級國際學(xué)術(shù)會議之一，每年舉辦一屆。AAAI 2025 共有12957篇有效投稿，錄用3032篇，錄取率為 23.4%。

螞蟻有18篇技術(shù)Paper收錄，其中3篇Oral，15篇Poster，研究領(lǐng)域涉及增強大模型隱私保護、提高推理速度與推理能力、提升大模型訓(xùn)練效率、降低模型幻覺等。