RAG：KBLaM 高效整合 LLM 外部知识传统的微调与检索增强生成（RAG）方法各有弊端，微调成本高昂，RAG则因引

【参考原文】KBLAM: KNOWLEDGE BASE AUGMENTED LANGUAGE MODEL

对比项目	RAG	KBLaM (KB)
知识位置	外部知识库	内化于 LLM 内部
检索机制	显式的检索模块	隐式地通过注意力机制进行“检索”
知识格式	通常为非结构化的文本	结构化的知识三元组编码为向量
训练方式	通常检索和生成组件分开训练或有限的联合优化	单一模型进行端到端训练
计算成本	检索成本 + LLM 推理成本	LLM 推理成本，知识库大小呈线性扩展
知识更新机制	重新索引外部知识库	动态更新，无需重新训练整个模型

1， fine-tune，RAG，in-context，KB

传统的微调与检索增强生成（RAG）方法各有弊端，微调成本高昂，RAG则因引入独立检索模块增加了复杂度。而随着知识库规模扩大，上下文学习的效率也愈发低下。KBLaM可以随时更新，不用重新训练。它把外部知识转换成连续的键值向量对，直接嵌入到模型的注意力层中，从而实现隐式检索，当知识库更新时，只需要更新对应的键值对即可，无需重新训练整个模型，它的计算开销随着知识库大小呈线性增长，而不是二次增长。

这种方法比RAG更直接，不需要外部检索；
比上下文学习更高效，可以处理更多知识；
比微调更灵活，可以随时更新知识；微调的目标是将知识“记忆”进模型的参数中；而 KBLAM 的学习目标是在预训练句子编码器的向量空间与 LLM 的嵌入空间之间建立投影关系。

监督微调： 通过修改模型参数来引入新的知识。但这种方法效率较低，因为每次更新知识时都需要重新训练模型参数。而且，微调还可能导致“灾难性遗忘”，即模型在通用任务上的表现下降。
RAG： 假设外部知识以非结构化文档的形式存在。当接收到一个输入提示时，RAG 首先使用一个检索模块从文档中提取相关的语料片段，然后将这些提取结果与输入问题拼接成一个完整的提示，最后输入到 LLM 中进行处理。
长上下文语言模型： 可以直接将整个外部语料库放入上下文中，从而无需检索模块进行文档选择，简化了整个处理流程。然而，in-context learning 的时间和内存复杂度会随着上下文长度呈二次增长。此外，上下文中各个 token 之间的依赖关系，也使得注意力矩阵的解释和动态更新 KV 缓存（KV cache）变得更加困难。

2，Self-Attention & KV Cache

【self-attention】 KBLAM 也与一些使用结构化注意力掩码（而非标准的下三角因果注意力掩码）的方法类似。这些工作基于上下文中的独立性假设，例如不同随机选取的 few-shot 示例之间的独立性，或者来自不同来源的文档之间的独立性，从而让彼此独立的上下文不相互注意，以降低注意力运算的开销。在 KBLAM 中，我们也采用了类似的输入之间的独立性假设：我们认为不同的 KB 三元组是相互独立的，这一假设为 KBLAM 带来了更好的可扩展性和可解释性。

标准的自注意力实现的时间复杂度为 $O(N^2D)$ ，存储复杂度为 $O(N^2)$ ，用于计算和存储所有 $w_{n,i}$ 。此外，FFN 的计算也会引入显著的计算开销，规模为 $O(ND^2)$ 。由于这些因素，随着序列长度的增加，自注意力会面临较高的内存消耗和较慢的计算速度。

【KV Cache】 在给定上下文字符串的情况下，KV 缓存机制会在每一层缓存 key 和 value 嵌入，这一阶段也被称为“预填充（prefill）”阶段。使用 KV 缓存后，生成新 token 时的复杂度从上下文长度的平方降低为线性。KBLAM 编码后的知识token的工作方式与 KV 缓存类似，但其 key 和 value 不是通过自注意力机制计算得到的，而是由外部编码器生成的。此外，在标准 KV 缓存机制中，如果上下文被修改，则需要重新计算整个 KV 缓存，这是由于因果注意力掩码的存在；而在KBLAM中，由于知识token之间不相互注意，仅需更新对应的知识token即可。

3，知识表示

KBLaM 将知识三元组（<name> - <property> - <value>）通过预训练的句子编码器和轻量级线性适配器，转化为连续的键值向量对，将非结构化的外部语料通过现有工具转化为结构化的知识库（KB）。

公式（1）： $\{(<name>_m;<property>_m;<value>_m)\}_{m=1}^M$

【首先】KBLAM 将每个知识三元组映射为一个固定长度的键值向量对，称为知识 token，其大小与 LLM 中单个 token 的 KV 缓存（Key-Value Cache）一致，因此可以无缝地集成到 LLM 的每一层注意力机制中。具体来说，KBLAM 将 <name> 和 <property> 编码为一个key向量，作为标识符，模仿 token 的 key 嵌入；将 <value> 编码为一个value向量，提供实际内容，类似于 token 的 value 嵌入。

【接着】KBLAM 利用三元组之间的结构关系，通过一种简单的矩形注意力结构，以可扩展和动态的方式将知识 token 注入 LLM 的注意力机制中：具有不同 <name> 和 <property> 的三元组可以被视为相互独立的信息，因此每个三元组生成的知识 token 都可以被独立编码并注入到预训练的 LLM 中。这种设计使得 KBLAM 的计算复杂度（无论是内存还是时间）随三元组数量线性增长，相比之下，上下文内学习的开销是与上下文长度成二次关系的，因此 KBLAM 具有更好的可扩展性。

【此外】这种独立性设计也使得我们可以通过仅修改对应的知识 token 来更新、删除或添加某个三元组，无需进行其他改动，这是标准KV缓存机制无法实现的。更重要的是，这种注意力结构使得模型对知识token的使用高度可解释。模型对知识token的使用情况可以通过注意力得分直接观察到。

4，知识 token

给定一个符合公式 (1) 格式的知识库（KB），对于每个三元组，首先采用一个预训练的句子编码器模型，记为 $f(\cdot)$ ，该模型将字符串转换为 P 维的连续嵌入向量。通过编码器，将每个三元组转换为基础的键嵌入和基础的值嵌入。特别地，对于第 m 个三元组：

$k_m=f(The\, <property>_m\, of \, <name>_m\, )\in R^P$

$v_m=f(<value>_m)\in R^P$

线性KV适配器： 其中， $L$ 表示模型中的注意力层数。借助适配器，在句子编码器空间中的 $k_m$ ， $v_m$ 映射到大型语言模型（LLM）每一层注意力机制中的键和值的嵌入空间。

$\tilde{W}_K\in R^{L\times D\times P}$

$\tilde{W}_V\in R^{L\times D\times P}$

具体来说，对于第 $m$ 个知识三元组，我们将其基础键和值嵌入转换为：

$\tilde{k}_m=[\tilde{k}_m^1,...,\tilde{k}_m^l,...,\tilde{k}_m^L]^T=\tilde{W}_Kk_m\in R^{L\times D}$

$\tilde{v}_m=[\tilde{v}_m^1,...,\tilde{v}_m^l,...,\tilde{v}_m^L]^T=\tilde{W}_Vv_m\in R^{L\times D}$

由于每个 $\tilde{k}^l_m$ 和 $\tilde{v}^l_m$ 的维度与大语言模型（LLM）中的 key 和 value 向量中的 $k^l_n$ 和 $v^l_n$ ）相同，因此它们可以直接融入注意力计算中。因此，这一编码过程等价于将一个知识三元组从由多个 token 构成的字符串转换为一个特殊的 token，其在每一层的 key 和 value 向量不是通过自注意力机制生成的，而是通过一个编码器生成的。因此，我们将 $(\tilde{k}_m, \tilde{v}_m)$ 这一对称为知识 token（knowledge token）。

这个编码过程会应用于知识库（KB）中的所有三元组，从而将知识库中的信息转化为一组 知识 token（knowledge tokens） 的集合。

$\{(<name>_m,<property>_m,<value>_m)\}_{m=1}^M \overset{Encode}{\longrightarrow } \{(\tilde{k}_m,\tilde{v}_m)\}_{m=1}^M$

对于一个包含 M 个三元组的知识库，首先为每个三元组构造一个 key 字符串和一个 value 字符串。接着，这两个字符串分别通过一个预训练的句子编码器处理，然后再通过一个可学习的线性适配器。最终获得的知识 token将被存储在磁盘上以供后续使用。

5，矩阵注意力

在将知识库（KB）转换为一组知识 token 之后，通过一种修改过的注意力结构，将知识库的信息注入到预训练大语言模型（LLM）的每一层注意力中。具体而言，假设输入提示（prompt）包含 $N$ 个 token，第 $l$ 层的注意力输入为： $x^l=[x_1^l,...,x_n^l,...,x_N^l]^T$ ，以及第 $l$ 层中来自知识 token 的键值对集合： $\{(\tilde{k}_m^l,\tilde{v}_m^l\}_{m=1}^M$ ，那么该层的注意力输出为： $\tilde{y}^l=[\tilde{y}^l_1,...,\tilde{y}^l_n,...,\tilde{y}^l_N]^T\in R^{N\times D}$

公式（9）： ${\Large\tilde{y}_n^l=\frac{\sum_{m=1}^M exp(\tilde{w}_{n,m}^l)\tilde{v}_m^l+\sum_{i=1}^n exp(w_{n,i}^l)v_i^l}{ \sum_{m=1}^M exp(\tilde{w}_{n,m}^l)+\sum_{i=1}^n exp(\tilde{w}_{n,i}^l) }}$

其中： $\tilde{w}_{n,m}^l=<\tilde{q}_n^l,\tilde{k}_m^l>/\sqrt{D}$ ， $\tilde{q}_n^l=\tilde{W}_Q^l x_n$ ， $w_{n,i}^l=<q_n^l,k_m^l>/\sqrt{D}$ ， $\tilde{W}_Q^l\in R^{D\times D}$ 是一个可学习的查询头（query head） ，它从预训练模型中第 $l$ 层的权重 $W_Q^l$ 初始化而来。

【矩阵注意力】 提示（prompt）中的 token 除了可以关注提示中之前的所有 token，还可以关注所有的知识 token。而另一方面，知识 token 之间是无法相互关注的，也就是说，它们之间没有自注意力，也没有查询（query）向量。

这种设计使得注意力矩阵的尺寸为 $(M+N)\times N$ ，其中 $M$ 是知识库中的三元组数量（即知识 token 的数量）， $N$ 是提示 token 的数量。另外，如果没有引入知识 token（即 $M=0$ ），那么公式 (9) 会退化为 Transform 经典公式，得到原始的预训练 LLM 的注意力结构。

矩形注意力的内存和时间复杂度分别为 $O((M+N)N)$ 和 $O((M+N)ND)$ 。在实际应用中，通常存在 $M\gg N$ 的情况，也就是说外部知识项的数量远大于提示或问题的长度。在这种情况下，开销只会随 $M$ 线性增长（而不是二次增长），这使得 KBLAM 可以扩展到非常大的 $M$ ，从而支持在上下文中引入大量的知识信息。此外，需要注意的是，矩形注意力的输出序列长度不随知识 token 的数量 $M$ 变化，因此中间的前馈网络（FFN）模块的计算开销不会因 $M$ 而改变，保持不变。

知识 token 的顺序变化不会影响输出，因此 KBLAM 不会受到上下文学习中常见的位置偏差（positional bias）问题的影响。 此外，由于所有知识 token 是独立编码的，当知识库发生更新（例如新增或修改某个知识三元组）时，只需使用编码器生成新的知识 token，或更新对应的知识 token 值即可。这一点凸显了 KBLAM 与标准 LLM 中 KV 缓存机制（KV cache） 的不同：后者在缓存内容有任何修改时，都需要重新计算。