論文名稱(chēng)KBLAM: KNOWLEDGE BASE AUGMENTED LANGUAGE MODEL

ICLR 2025的論文，MS系的這篇論文介紹了一種名為 KBLAM（Knowledge Base augmented Language Model，知識(shí)庫(kù)增強(qiáng)語(yǔ)言模型） 的新方法，用于將外部知識(shí)庫(kù)（KB）高效地集成到預(yù)訓(xùn)練的大型語(yǔ)言模型（LLM）中，以提升其在特定任務(wù)中的表現(xiàn)，同時(shí)避免傳統(tǒng)方法的局限性。

核心問(wèn)題與背景

LLM雖然在知識(shí)和推理方面表現(xiàn)優(yōu)異，但當(dāng)需要處理超出其參數(shù)存儲(chǔ)知識(shí)的外部信息時(shí)，往往需要額外的增強(qiáng)手段。傳統(tǒng)方法包括：

1. 監(jiān)督微調(diào)（Fine-tuning）：通過(guò)更新模型權(quán)重來(lái)融入外部知識(shí)，但效率低，且可能導(dǎo)致“災(zāi)難性遺忘”（對(duì)通用任務(wù)性能下降）。

2. 檢索增強(qiáng)生成（RAG）：利用外部檢索模塊從文檔中提取相關(guān)信息并拼接到輸入中，雖然有效，但依賴(lài)獨(dú)立的檢索系統(tǒng)，增加了復(fù)雜性。

3. 下文學(xué)習(xí)（In-context Learning）：直接將整個(gè)外部語(yǔ)料放入模型上下文，雖然簡(jiǎn)單，但計(jì)算和內(nèi)存開(kāi)銷(xiāo)隨上下文長(zhǎng)度呈二次方增長(zhǎng)，且難以動(dòng)態(tài)更新，其實(shí)和rag一樣他倆都是O2的計(jì)算復(fù)雜度，因?yàn)樽叩氖俏谋緸榻o大模型。

這些方法都又局限性，所以作者就提出一種新的解決方案：KBLAM。

KBLAM 的核心思想

KBLAM 的目標(biāo)是通過(guò)將外部知識(shí)轉(zhuǎn)化為結(jié)構(gòu)化的“知識(shí)令牌”（knowledge tokens），并利用一種高效的“矩形注意力機(jī)制”（rectangular attention）將其融入預(yù)訓(xùn)練的 LLM 中。主要特點(diǎn)包括：

• 知識(shí)表示：將外部語(yǔ)料轉(zhuǎn)化為知識(shí)三元組（<name>, <property>, <value>），如“Apple, founder, Steve Jobs”，并通過(guò)預(yù)訓(xùn)練的句子編碼器（帶線(xiàn)性適配器）將每個(gè)三元組編碼為固定長(zhǎng)度的鍵值向量對(duì)（key-value vector pairs），稱(chēng)為“知識(shí)令牌”。
• 注意力機(jī)制：使用矩形注意力結(jié)構(gòu)，使輸入提示（prompt）的查詢(xún)（query）可以同時(shí)關(guān)注自身令牌和所有知識(shí)令牌，而知識(shí)令牌之間互不關(guān)注。這種設(shè)計(jì)使得計(jì)算復(fù)雜度隨知識(shí)庫(kù)大小呈線(xiàn)性增長(zhǎng)（而非二次方），并支持動(dòng)態(tài)更新。
• 訓(xùn)練方式：通過(guò)指令微調(diào)（instruction tuning）在合成數(shù)據(jù)上學(xué)習(xí)線(xiàn)性適配器（你就當(dāng)整出來(lái)個(gè)lora就可以了），將編碼器空間映射到 LLM 的嵌入空間，而無(wú)需微調(diào)整個(gè)模型。

注意力是這樣的。

與現(xiàn)有方法的對(duì)比

• 對(duì)比 RAG：KBLAM 無(wú)需外部檢索模塊，所有知識(shí)通過(guò)注意力機(jī)制端到端融入模型。這么做肯定是降低延遲提升體驗(yàn)，降低故障點(diǎn)。
• 對(duì)比上下文學(xué)習(xí)：KBLAM 的內(nèi)存和時(shí)間復(fù)雜度更低（線(xiàn)性 vs 二次方），且支持動(dòng)態(tài)更新知識(shí)庫(kù)，而無(wú)需重新計(jì)算整個(gè)上下文。
• 優(yōu)勢(shì)：KBLAM 可以在單張 A100 80GB GPU 上將超過(guò) 10,000 個(gè)三元組的知識(shí)庫(kù)集成到僅 8,000 上下文窗口的 8B 參數(shù) LLM 中，同時(shí)保持可解釋性（通過(guò)注意力分?jǐn)?shù)分析模型如何使用知識(shí)）。

如圖：

人話(huà)解釋就是：

• SFT效果就很難好，其實(shí)這論文里說(shuō)的只是一方面，最核心的是數(shù)據(jù)分布，加上專(zhuān)有知識(shí)領(lǐng)域就那么點(diǎn)東西，很難對(duì)parameter有什么好作用，要不就不行，要不就再難遺忘，過(guò)擬合
• RAG和prompt enginering+inconext-learning就沒(méi)什么可說(shuō)的了
• 咱們這個(gè)方案理解起來(lái)某種程度上其實(shí)有點(diǎn)復(fù)雜

為什么說(shuō)復(fù)雜呢？

我先問(wèn)你rag送給llm的是啥？

是L2距離或者余弦距離比對(duì)的最短的那塊吧，然后你也不是給向量表達(dá)，而是把對(duì)應(yīng)最短的index的value也就是你對(duì)應(yīng)的那段chunk文本+prompt一起發(fā)給LLM對(duì)不，說(shuō)白了，你做的還時(shí)inconext-learning，區(qū)別只是給多給少罷了。

那咱們這個(gè)是什么呢？

你給的是整個(gè)的知識(shí)庫(kù)轉(zhuǎn)換的知識(shí)令牌（KBtokens）

啥是知識(shí)令牌呢？

以下進(jìn)入抽象環(huán)節(jié)

知識(shí)令牌是外部知識(shí)庫(kù)中每個(gè)知識(shí)三元組（例如，<name>, <property>, <value>，如 <Apple, founder, Steve Jobs>）的連續(xù)向量表示。具體生成過(guò)程如下：

• 鍵向量（key vector）：通過(guò)預(yù)訓(xùn)練的句子編碼器（帶有線(xiàn)性適配器）將三元組中的 <name> 和 <property> 編碼為一個(gè)固定長(zhǎng)度的向量。
• 值向量（value vector）：將 <value> 編碼為另一個(gè)固定長(zhǎng)度的向量。
• 這些鍵值向量對(duì)（key-value vector pairs）被稱(chēng)為知識(shí)令牌，是在預(yù)處理階段生成并存儲(chǔ)的，而不是直接以文本形式存在。

這玩意有啥用呢？

在 KBLAM 中，知識(shí)令牌與提示中的令牌通過(guò) 矩形注意力機(jī)制 進(jìn)行聯(lián)合計(jì)算。這種機(jī)制的核心思想是讓提示中的令牌能夠同時(shí)關(guān)注自身和其他提示令牌，以及所有知識(shí)令牌，而知識(shí)令牌之間不相互關(guān)注。具體過(guò)程如下：

矩形注意力機(jī)制的結(jié)構(gòu)

圖片

還是這個(gè)圖

• 提示令牌的注意力

a.提示中所有先前令牌的鍵向量（key vector）。

b.所有知識(shí)令牌的鍵向量。

• 提示中的每個(gè)令牌生成一個(gè)查詢(xún)向量（query vector），說(shuō)白了就是QKV里的Q（你可以這么理解）他能關(guān)注到：1- 提示中所有先前令牌的鍵向量（key vector）。
  2- 所有知識(shí)令牌的鍵向量。
• 通過(guò)計(jì)算注意力分?jǐn)?shù)，提示令牌（prompt的token）可以從這些知識(shí)令牌的鍵向量對(duì)應(yīng)的值向量（value vector）中獲取信息，這不就相當(dāng)于學(xué)到了嗎？但是知識(shí)令牌之間是不做運(yùn)算的，只有prompt的token和KB token做運(yùn)算，知識(shí)令牌之間不相互關(guān)注，彼此獨(dú)立。這種設(shè)計(jì)降低了計(jì)算復(fù)雜度。

對(duì)于提示中的第 n個(gè)令牌，在模型第L層的注意力輸出計(jì)算如下：

就是一頓softmax。

計(jì)算過(guò)程解釋

• 注意力分?jǐn)?shù)提示令牌的查詢(xún)向量分別與提示中的鍵向量和知識(shí)令牌的鍵向量計(jì)算點(diǎn)積，得到注意力分?jǐn)?shù)。
• 加權(quán)求和根據(jù)注意力分?jǐn)?shù)，對(duì)提示中先前令牌的值向量和所有知識(shí)令牌的值向量進(jìn)行加權(quán)求和，生成最終的注意力輸出。（另一個(gè)維度的QKV）
• 這最后就能實(shí)現(xiàn)往prompt tokens和知識(shí)庫(kù)的KB tokens實(shí)現(xiàn)了信息融合，就相當(dāng)于把外部知識(shí)給加載進(jìn)來(lái)了。

知識(shí)令牌怎么加載的呢？

存儲(chǔ)形式：知識(shí)令牌是以鍵值向量的形式存儲(chǔ)的，而不是文本令牌。這些向量在預(yù)處理階段生成并保存在內(nèi)存中。

推理過(guò)程：

• 在推理時(shí)，提示中的文本令牌被輸入到 LLM 中，生成對(duì)應(yīng)的查詢(xún)、鍵和值向量。
• 知識(shí)令牌的鍵值向量被直接加載到模型的注意力層，與提示令牌的鍵值向量一起參與計(jì)算。
• 效率優(yōu)勢(shì)：這種方式不僅減少了輸入處理的開(kāi)銷(xiāo)，還支持動(dòng)態(tài)更新——當(dāng)知識(shí)庫(kù)變化時(shí)，只需更新對(duì)應(yīng)的知識(shí)令牌向量，而無(wú)需重新處理整個(gè)上下文

前面這些應(yīng)該講明白了，那你原生LLM怎么能認(rèn)這些KB tokens?。?/p>

就是我前面提的訓(xùn)練。

• 訓(xùn)練目標(biāo)

a.KBLAM 在訓(xùn)練時(shí)會(huì)學(xué)習(xí)如何將知識(shí)庫(kù)中的三元組（比如 <名字>, <屬性>, <值>）通過(guò)句子編碼器和線(xiàn)性適配器映射到 LLM 的嵌入空間。

• 訓(xùn)練的目標(biāo)不是讓模型記住某個(gè)具體的知識(shí)庫(kù)，而是學(xué)會(huì)一種通用的方法，能夠處理各種不同的知識(shí)庫(kù)。
• 知識(shí)令牌的生成

a.每個(gè)知識(shí)三元組會(huì)被獨(dú)立編碼成一個(gè)知識(shí)令牌（knowledge token），這些令牌是模型在推理時(shí)用來(lái)提取信息的。

• 動(dòng)態(tài)更新

a.KBLAM 的設(shè)計(jì)支持知識(shí)庫(kù)的動(dòng)態(tài)更新。如果您想注入一個(gè)新的知識(shí)庫(kù)，只需要將新知識(shí)庫(kù)中的三元組生成對(duì)應(yīng)的知識(shí)令牌，然后加載到模型中即可。

b.如果知識(shí)庫(kù)內(nèi)容有變化（比如添加、刪除或修改三元組），只需更新對(duì)應(yīng)的知識(shí)令牌，不需要重新訓(xùn)練整個(gè)模型。

那訓(xùn)啥呢？

在 KBLAM 的推理過(guò)程中，以下幾個(gè)部分是需要通過(guò)訓(xùn)練來(lái)學(xué)習(xí)的：

1. 線(xiàn)性適配器（Linear Adapters）

a.鍵適配器 WK將知識(shí)三元組的鍵向量（從 <name> 和 <property> 編碼生成）映射到 LLM 的鍵向量空間。

b.值適配器 WV將知識(shí)三元組的值向量（從 <value> 編碼生成）映射到 LLM 的值向量空間。

• 作用線(xiàn)性適配器負(fù)責(zé)將句子編碼器生成的知識(shí)令牌向量映射到 LLM 的嵌入空間。具體來(lái)說(shuō)：
  句子編碼器和 LLM 的向量空間天然是不同的。如果不訓(xùn)練線(xiàn)性適配器，知識(shí)令牌的向量無(wú)法與 LLM 的內(nèi)部表示對(duì)齊，導(dǎo)致注意力機(jī)制無(wú)法正確識(shí)別和利用這些知識(shí)。

2. 查詢(xún)頭（Query Head, WQ）

• 作用KBLAM 在 LLM 的每個(gè)注意力層中引入了一個(gè)額外的查詢(xún)頭，用于生成提示令牌的查詢(xún)向量（query vector）。這些查詢(xún)向量與知識(shí)令牌的鍵向量進(jìn)行注意力計(jì)算，從而決定提示令牌關(guān)注哪些知識(shí)令牌。
• 未經(jīng)訓(xùn)練的查詢(xún)頭無(wú)法生成合適的查詢(xún)向量，提示令牌就無(wú)法準(zhǔn)確地從知識(shí)令牌中提取相關(guān)信息。例如，如果查詢(xún)頭參數(shù)是
  是隨機(jī)的，注意力機(jī)制可能會(huì)關(guān)注錯(cuò)誤的知識(shí)令牌，導(dǎo)致答案不準(zhǔn)確。

這些所有要被訓(xùn)練的東西，我們叫

其他的什么原始模型的QKV啊，F(xiàn)FN啊，你都不用訓(xùn)，都是凍結(jié)的。

然后你這些整個(gè)KB tokens是不是要加載啊，比如你放在外面，你和rag有什么區(qū)別？這個(gè)時(shí)候的主要區(qū)別就是rag是text，需要外置查詢(xún)，我這個(gè)放外面也是一堆tokens，而且rag因?yàn)槟阋獜淖咭槐閠ransformer的邏輯，所以還是O2的復(fù)雜度，而我這個(gè)是O

第二點(diǎn)，我這東西比起整個(gè)rag還是小的，舉個(gè)例子：

• 每個(gè)知識(shí)令牌的內(nèi)存需求
  假設(shè)使用的是像 Llama 3 8B 這樣的模型，嵌入維度 D=4096。每個(gè)知識(shí)令牌包含一個(gè)鍵向量和一個(gè)值向量，每個(gè)向量占用 4096×4 字節(jié)（float32 格式），即 16 KB。因此，一個(gè)知識(shí)令牌的總占用約為2×16 KB=32 KB。
• 總內(nèi)存占用估算
  如果知識(shí)庫(kù)包含 M=??10,?000 個(gè)三元組，那么加載所有知識(shí)令牌所需的顯存為：10,000×32 KB=320 MB對(duì)于現(xiàn)代 GPU（例如 NVIDIA A100 80GB），320 MB 僅占顯存的極小部分。即使知識(shí)庫(kù)規(guī)模擴(kuò)大到 100,000 個(gè)三元組，也不過(guò)需要約 3.2 GB 顯存，仍然遠(yuǎn)低于 A100 的容量。
• 實(shí)際驗(yàn)證
  KBLAM 的論文中提到，它可以在單張 A100 80GB GPU 上處理超過(guò) 10,000 個(gè)三元組的知識(shí)庫(kù)，并將其集成到僅 8,000 個(gè)上下文窗口的 8B 參數(shù)語(yǔ)言模型中。這表明，對(duì)于中等規(guī)模的知識(shí)庫(kù)，將所有知識(shí)令牌加載到顯存中是完全可行的。

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

• 任務(wù)：包括問(wèn)答（Q&A）和開(kāi)放式推理。
• 表現(xiàn)：在合成數(shù)據(jù)集上，KBLAM 與上下文學(xué)習(xí)性能相當(dāng)，但在內(nèi)存占用更低的情況下能擴(kuò)展到更大規(guī)模的知識(shí)庫(kù)；在 Enron（真實(shí)數(shù)據(jù)集）上性能略有下降，但仍?xún)?yōu)于零樣本基線(xiàn)。
• 可解釋?zhuān)鹤⒁饬仃囷@示模型能準(zhǔn)確檢索相關(guān)知識(shí)三元組。
• 抗幻覺(jué)：KBLAM 能在知識(shí)庫(kù)中無(wú)答案時(shí)拒絕回答，減少幻覺(jué)現(xiàn)象。

局限性與未來(lái)工作

• 合成數(shù)據(jù)質(zhì)量：當(dāng)前使用合成數(shù)據(jù)訓(xùn)練，面對(duì)分布外數(shù)據(jù)（如 Enron）時(shí)性能下降，未來(lái)可改進(jìn)合成數(shù)據(jù)的多樣性。
• 信息壓縮：知識(shí)令牌固定長(zhǎng)度可能丟失細(xì)節(jié)（如精確數(shù)字），可考慮可調(diào)壓縮率。
• 復(fù)雜推理：未來(lái)可探索支持多跳推理或思維鏈（chain-of-thought）等任務(wù)。

KBLAM 提供了一種高效、可擴(kuò)展且可解釋的方法，將結(jié)構(gòu)化知識(shí)庫(kù)融入 LLM，避免了傳統(tǒng)方法的計(jì)算瓶頸和動(dòng)態(tài)性不足。論文還開(kāi)放了代碼和數(shù)據(jù)集（https://github.com/microsoft/KBLaM），為后續(xù)研究奠定了基礎(chǔ)。

本文轉(zhuǎn)載自??熵減AI??，作者：周博洋