論文名稱KBLAM: KNOWLEDGE BASE AUGMENTED LANGUAGE MODEL

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

核心問題與背景

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

1. 監(jiān)督微調(diào)（Fine-tuning）：通過更新模型權重來融入外部知識，但效率低，且可能導致“災難性遺忘”（對通用任務性能下降）。

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

3. 下文學習（In-context Learning）：直接將整個外部語料放入模型上下文，雖然簡單，但計算和內(nèi)存開銷隨上下文長度呈二次方增長，且難以動態(tài)更新，其實和rag一樣他倆都是O2的計算復雜度，因為走的是文本為給大模型。

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

KBLAM 的核心思想

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

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

注意力是這樣的。

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

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

如圖：

人話解釋就是：

• SFT效果就很難好，其實這論文里說的只是一方面，最核心的是數(shù)據(jù)分布，加上專有知識領域就那么點東西，很難對parameter有什么好作用，要不就不行，要不就再難遺忘，過擬合
• RAG和prompt enginering+inconext-learning就沒什么可說的了
• 咱們這個方案理解起來某種程度上其實有點復雜

為什么說復雜呢？

我先問你rag送給llm的是啥？

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

那咱們這個是什么呢？

你給的是整個的知識庫轉(zhuǎn)換的知識令牌（KBtokens）

啥是知識令牌呢？

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

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

• 鍵向量（key vector）：通過預訓練的句子編碼器（帶有線性適配器）將三元組中的 <name> 和 <property> 編碼為一個固定長度的向量。
• 值向量（value vector）：將 <value> 編碼為另一個固定長度的向量。
• 這些鍵值向量對（key-value vector pairs）被稱為知識令牌，是在預處理階段生成并存儲的，而不是直接以文本形式存在。

這玩意有啥用呢？

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

矩形注意力機制的結(jié)構

圖片

還是這個圖

• 提示令牌的注意力

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

b.所有知識令牌的鍵向量。

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

對于提示中的第 n個令牌，在模型第L層的注意力輸出計算如下：

就是一頓softmax。

計算過程解釋

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

知識令牌怎么加載的呢？

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

推理過程：

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

前面這些應該講明白了，那你原生LLM怎么能認這些KB tokens啊？

就是我前面提的訓練。

• 訓練目標

a.KBLAM 在訓練時會學習如何將知識庫中的三元組（比如 <名字>, <屬性>, <值>）通過句子編碼器和線性適配器映射到 LLM 的嵌入空間。

• 訓練的目標不是讓模型記住某個具體的知識庫，而是學會一種通用的方法，能夠處理各種不同的知識庫。
• 知識令牌的生成

a.每個知識三元組會被獨立編碼成一個知識令牌（knowledge token），這些令牌是模型在推理時用來提取信息的。

• 動態(tài)更新

a.KBLAM 的設計支持知識庫的動態(tài)更新。如果您想注入一個新的知識庫，只需要將新知識庫中的三元組生成對應的知識令牌，然后加載到模型中即可。

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

那訓啥呢？

在 KBLAM 的推理過程中，以下幾個部分是需要通過訓練來學習的：

1. 線性適配器（Linear Adapters）

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

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

• 作用線性適配器負責將句子編碼器生成的知識令牌向量映射到 LLM 的嵌入空間。具體來說：
  句子編碼器和 LLM 的向量空間天然是不同的。如果不訓練線性適配器，知識令牌的向量無法與 LLM 的內(nèi)部表示對齊，導致注意力機制無法正確識別和利用這些知識。

2. 查詢頭（Query Head, WQ）

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

這些所有要被訓練的東西，我們叫

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

然后你這些整個KB tokens是不是要加載啊，比如你放在外面，你和rag有什么區(qū)別？這個時候的主要區(qū)別就是rag是text，需要外置查詢，我這個放外面也是一堆tokens，而且rag因為你要從走一遍transformer的邏輯，所以還是O2的復雜度，而我這個是O

第二點，我這東西比起整個rag還是小的，舉個例子：

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

實驗與結(jié)果

• 任務：包括問答（Q&A）和開放式推理。
• 表現(xiàn)：在合成數(shù)據(jù)集上，KBLAM 與上下文學習性能相當，但在內(nèi)存占用更低的情況下能擴展到更大規(guī)模的知識庫；在 Enron（真實數(shù)據(jù)集）上性能略有下降，但仍優(yōu)于零樣本基線。
• 可解釋：注意力矩陣顯示模型能準確檢索相關知識三元組。
• 抗幻覺：KBLAM 能在知識庫中無答案時拒絕回答，減少幻覺現(xiàn)象。

局限性與未來工作

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

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

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