通过解码器-only大语言模型内部键值重路由实现免训练的文本嵌入 KV-Embedding: Training-free Text Embedding via Internal KV Re-routing in Decoder-only LLMs
利用LLM内部KV重路由在无需训练情况下提取高质量文本嵌入
前置知识
Decoder-only LLM架构
仅包含解码器部分的自回归模型，如GPT、Llama、Qwen等。采用因果注意力机制，每个位置只能关注前面的token，通过下一个token预测任务进行预训练。这种架构高效且可扩展，成为当前大语言模型的主流设计。
本文的核心目标就是从这类冻结的模型中提取嵌入，需要理解其因果注意力的结构局限性
键值（KV）对与注意力机制
在Transformer注意力机制中，输入通过三个线性投影变换为Query（Q）、Key（K）、Value（V）。注意力权重通过Q与K的点积计算，然后用于加权聚合V。KV对本质上构成了联想记忆：K编码查询模式，V存储语义内容。最终位置（最后一个token）的KV状态包含了对整个序列的压缩视图。
本文的创新点正是利用最终token的KV对作为全局上下文摘要，重新路由到其他位置
内在维度性
描述数据流形固有维度的数学概念，不同于空间的维度。对于神经网络隐藏状态，ID反映了表示流形的复杂度和压缩程度。TwoNN估计器通过分析最近邻距离比率来估计ID。研究观察表明，ID的最小值对应最大语义抽象层。
本文使用ID来自动选择最优的KV重路由层，避免手工调参，确保模型无关的适用性
各向异性问题
在嵌入空间中，向量倾向于聚集在狭窄锥体内而非均匀分布在超球面上，导致相似度计算区分度低。这通常由解码器模型的训练目标（next-token prediction）和词汇预测偏差导致。各向异性使得检索和语义匹配任务性能下降。
本文通过KV重路由产生更均匀的嵌入空间，改善了alignment和uniformity指标
因果注意力与信息不对称
因果注意力通过上三角掩码限制每个token只关注前面的位置，保证自回归性。但这造成信息不对称：早期token无法获取后续上下文。例如在'the bank of the river'中，'bank'在看到'river'之前无法消歧。这种不对称导致早期token的表示质量下降。
这是KV-Embedding要解决的核心问题之一，通过重路由最终KV让所有位置获得全局上下文
研究动机
现有的解码器-only大语言模型作为嵌入backbone时，在免训练设置下面临两个结构性挑战：首先，因果注意力限制了早期token访问后续上下文，造成信息不对称。例如在处理'the bank of the river'这类歧义短语时，'bank'这个词在没有看到'river'之前无法判断是'银行'还是'河岸'，导致其表示模糊。其次，next-token预测的训练目标使最终token的表示偏向预测未来内容而非压缩语义本质。现有训练-free方法通过PromptEOL设计提示词、Echo重复输入序列、Token Prepending插入特殊token来缓解这些问题，但都有明显缺陷：Echo将序列长度加倍导致注意力复杂度平方增长（从O(n²)到O(4n²)），并可能引发'lost-in-the-middle'效应；Token Prepending依赖词汇表外的特殊token，产生不可预测的表示；这些方法都依赖输入级修改而非内部状态操作。
本文的目标是本文的目标是开发一个训练-free框架，能够从冻结的解码器-only大语言模型中提取高质量文本嵌入，同时避免现有方法的局限性。具体而言，需要在不增加序列长度、不引入out-of-vocabulary token、不修改模型参数的前提下，让所有位置都能获得序列级别的全局上下文信息，并缓解最终token的预测偏差。此外，框架应该是模型无关的，能够自适应不同架构而不需要手工调参。
与已有工作不同的是，本文的独特切入角度在于：与其通过输入级修改（如重复输入、插入特殊token）来模拟双向上下文，不如直接利用模型内部已经计算好的全局信息。作者观察到，由于因果注意力结构，每一层最终位置的KV对编码了该层对整个序列的压缩视图——这正是我们需要的全局上下文摘要。通过将这些内部状态重新路由为前缀插入注意力机制，可以让所有位置在一次前向传播中就访问到序列级信息。这种方法既保持了因果约束（不破坏预训练模式），又实现了全局上下文访问。此外，作者还创新性地引入内在维度性来自动选择重路由层，避免了人工经验调参，实现了真正的模型无关性。
核心方法
KV-Embedding的整体思路是：解码器-only LLM在处理输入时，每一层的最终token的KV状态实际上已经包含了该层对整个序列的压缩摘要——这是因果注意力带来的自然属性。如果能将这个摘要传递给序列中的其他位置，就能在不打破因果约束的情况下实现全局上下文访问。方法由三个组件构成：（1）压缩导向的提示词，引导模型将最终表示转向语义压缩而非预测；（2）内部KV重路由机制，在选定层将最终位置的KV对插入为全局前缀；（3）基于内在维度性的自动层选择策略，找到语义压缩程度最高的层进行重路由。整个过程只需一次前向传播，无需修改输入序列或模型参数。
核心创新点是将最终token的KV状态作为内部全局摘要，在注意力计算前将其前缀插入到KV矩阵中。这与现有方法有本质区别：PromptEOL只改变提示词但不解决信息不对称；Echo和Token Prepending通过修改输入来提供全局上下文，但会增加序列长度或依赖特殊token；而KV-Embedding直接操作模型的内部状态，让摘要信息在正确的层注入到注意力机制中。这就像给每个位置都发了一份'序列笔记'，而不是要求每个人都从头读一遍整个文档。另一个关键创新是使用内在维度性来选择重路由层——这是首次将表示几何学用于训练-free嵌入提取，实现了跨架构的自适应。
方法步骤详情
完整的方法步骤如下：输入文本x后，首先应用压缩导向的提示词模板（如'Context: {text} Compress the context in one word:'）来引导模型。然后进行L层的前向传播，对每一层l：如果l在选定层集合L中，则先计算所有位置的完整KV矩阵K^(l)和V^(l)，提取最终位置的KV对(k^(l)_n, v^(l)_n)，然后将其前缀拼接得到增强矩阵：˜K^(l) = [k^(l)_n ∥ K^(l)]，˜V^(l) = [v^(l)_n ∥ V^(l)]，再使用这些增强矩阵计算注意力；如果l不在L中，则使用标准注意力。层选择通过TwoNN估计器在F2LLM数据集上计算各层的内在维度性，找到ID最小的层l*，然后定义L = {l | l* ≤ l ≤ min(L, l* + ⌊0.1L⌋)}。最后，对最后一层的隐藏状态应用混合池化（平均最后token和mean pooling），然后进行ℓ2归一化得到最终嵌入e。
技术新颖性
技术新颖性体现在多个方面：（1）首次提出在注意力机制内部重新路由KV状态而非修改输入，这保持了模型预训练的因果先验，避免了naive移除causal mask导致的性能崩溃（Table 4显示双向注意力得分仅0.22-0.26）。（2）将表示几何学（内在维度性）应用于训练-free设置，自动识别语义压缩峰值的层，避免了启发式规则。实验表明ID选择的层数（7-10层）比均匀策略（10-11层）更高效，性能更高（0.534 vs 0.529）。（3）结合压缩导向提示词和KV重路由，同时解决两个结构性问题：前者缓解预测偏差，后者解决信息不对称。消融研究显示两者缺一不可——去掉KV重路由性能从0.534降至0.450。（4）使用混合池化策略，平衡了全局摘要和分布证据，优于单一last-token或mean pooling（0.534 vs 0.504/0.444）。
实验结果
KV-Embedding在MTEB和LoCoV1基准上取得了显著性能提升。在MTEB上，相比现有训练-free方法（PromptEOL、Echo、Token Prepending），本文方法在三个backbone上平均提升高达10%。具体而言，在Mistral-7B上达到0.534，超过最佳基线Echo（0.501）6.6%；在Qwen3-4B上达到0.494，超过PromptEOL（0.448）10.3%；在Llama-3.1-8B上达到0.527，超过Echo（0.503）4.8%。提升最显著的类别是Retrieval任务，这表明文档级匹配特别受益于全局上下文访问。在LoCoV1长文本检索上，KV-Embedding在1024、2048、4096 token长度下均大幅领先：Mistral-7B在4096长度下达到0.207，而基线方法均低于0.10；Qwen3-4B在4096长度下达到0.182，超过PromptEOL（0.129）41.1%。嵌入空间质量分析显示，KV-Embedding产生了更均匀的分布：Alignment指标从0.7378（PromptEOL）降至0.6082，Uniformity从-2.2715提升至-2.3899，这缓解了解码器模型的各向异性问题。消融实验证实了每个组件的必要性：去掉KV重路由性能下降15.8%；去掉压缩提示词性能下降约10%；使用早期层重路由性能仅0.409，而ID选择达到0.534。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MTEB平均分数 | Average Score | 0.534 | 0.501 | 6.6% |
| MTEB检索任务（NDCG@10） | Retrieval | 0.301 | 0.241 | 24.9% |
| LoCoV1长文本检索（4096 token） | NDCG@10 | 0.207 | 0.060 | 245% |
| STS语义相似度 | Spearman correlation | 0.772 | 0.733 | 5.3% |
| 嵌入空间均匀性 | Uniformity | -2.3899 | -2.2715 | 5.2% |
局限与改进
作者明确指出两个主要局限性：首先，KV重路由机制增加了推理延迟，因为需要在选定层进行额外的KV拼接操作。其次，作为训练-free方法，KV-Embedding的性能可能无法匹配在大型数据集上用对比目标精细调整的监督模型——这在长文本检索上仍然明显，4096长度下最好的结果也只有0.2左右。此外，从论文可以观察到一些额外限制：方法假设最终token的KV状态包含丰富的语义信息，但这对极短文本可能不成立；ID估计需要在每个新模型上运行（虽然只需一次）；混合池化可能不适合所有任务类型（如表4显示Summarization任务各方法差异很小）；当前方法只考虑单轮重路由，对于超长文本可能需要多轮摘要。
独立分析的弱点
KV-Embedding的几个独立分析弱点：（1）计算开销增加：重路由操作需要在选定层处理额外的KV对，虽然保持了O(n²)复杂度，但常数因子增大。改进方向：限制重路由层数（论文已经只选7-10层）、或使用流式处理渐进注入摘要、或优化KV拼接的CUDA实现。（2）对超长文本（>8k token）的有效性未验证：实验只到4096 token，更长的序列可能需要分层摘要或多轮压缩。改进方向：开发自适应摘要长度策略、或引入压缩梯度来平衡信息量和噪声。（3）单token摘要的局限性：最终token可能无法捕获所有重要信息，特别是当文本有多个主题时。改进方向：使用多个anchor token（如句子边界token）的KV对进行加权组合。（4）对不同架构的ID轨迹差异：Qwen3-4B是U型轨迹而Llama有多个局部最小值，当前策略需要手工处理这种情况。改进方向：开发更鲁棒的ID分析算法，或基于任务的层选择策略。
未来方向
作者和基于成果可延伸的未来研究方向：（1）探索多anchor KV重路由：不只是最终token，还可以考虑句末、段落末等位置的KV对，形成多尺度摘要。（2）动态层选择：根据输入长度或任务类型自适应调整重路由层，短文本用浅层摘要、长文本用深层摘要。（3）与训练方法结合：将KV-Embedding作为初始化或正则化项，指导对比学习的表示空间几何。（4）扩展到其他生成任务：不只是文本嵌入，类似的思路可用于改善长文本生成、推理、QA等需要全局理解的任务。（5）分析其他内部状态：除了KV，还可以探索FFN层的中间激活、attention pattern本身等是否包含有用的摘要信息。（6）理论分析：建立ID与语义压缩之间的理论联系，解释为什么ID最小对应最佳重路由点。（7）跨语言和跨模态：将方法推广到多语言和多模态模型，验证是否同样有效。
复现评估
论文在复现性方面表现良好：提供了详细的算法伪代码（Algorithm 1）和完整的超参数设置（attention bias b=1.0、混合池化、层选择策略）；实验使用公开可用的MTEB和LoCoV1基准，以及三个开源模型；所有实验在4张NVIDIA H800 GPU上运行，batch size 64，随机种子42，这提供了明确的硬件配置；附录提供了丰富的实验细节，包括42个数据集的完整结果表（Tables 8-17）、ID轨迹可视化（Figure 3）、消融实验（Tables 18-21）以及 probing实验（Tables 6-7）。论文还提到会开源代码，虽然当前（2026年1月）arXiv版本可能还没有代码链接，但实验设置的详细程度使得重现成为可能。主要难点可能在于：TwoNN估计器的实现细节（虽然引用了Facco et al. 2017）、ID轨迹分析的主观性（如何定义'清晰U型'和'多个局部最小值'）、以及不同模型的层映射（32层Mistral vs 36层Qwen）。对于有LLM推理经验的团队，复现难度中等；对于初学者，可能需要查阅相关实现细节。算力需求相对合理：无需训练，只需推理，单张A100或H800即可运行主要实验。
论文图表
