K-Token Merging：在潜在嵌入空间压缩序列，实现大语言模型的高效推理

K-Token Merging是一种创新的提示词压缩方法，通过在潜在嵌入空间合并连续token块，在保持模型性能的同时显著减少输入序列长度，为大语言模型的高效推理开辟新路径。

大型语言模型在处理长文本时面临一个根本性挑战：预填充（prefill）阶段的计算成本与输入序列长度呈线性关系。当用户提交一个包含数万token的长文档时，模型需要消耗大量计算资源来处理这些输入，才能开始生成回复。这种"长上下文惩罚"严重限制了LLM在文档分析、代码理解和多轮对话等场景中的应用效率。

传统的解决方案包括稀疏注意力机制、滑动窗口和分层处理等方法，但这些方法往往需要在模型架构层面进行改动。K-Token Merging另辟蹊径，提出了一种在潜在嵌入空间进行提示词压缩的全新思路。

K-Token Merging的核心洞察是：自然语言中存在大量冗余信息，相邻的token往往携带相似的语义信息。如果在嵌入层面将这些冗余token合并，可以在不显著损失信息的前提下大幅缩短序列长度。

具体来说，该方法将每K个连续的输入token视为一个块，通过一个轻量级编码器将这K个token的嵌入合并成单个潜在嵌入。这个压缩后的前缀随后被送入大语言模型进行预填充，而生成阶段仍在原始token空间进行。

K-Token Merging的工作流程分为两个明确的阶段：

在预填充阶段，编码器f接收每K个连续的输入token，产生单个压缩后的token嵌入。具体流程如下：

1. 对输入提示词进行tokenize
2. 从缓存的嵌入表中查找基础模型的token嵌入
3. 将提示词嵌入分割成大小为K的连续块
4. 使用轻量级编码器合并每个块（编码器初始化为均值池化行为，随后与LoRA适配器联合训练）
5. 将压缩后的前缀送入基础LLM

在生成阶段，LLM输出原始的未压缩token。每个新生成的token被追加到混合的压缩/未压缩前缀之后，随后继续进行标准的自回归生成。这种设计确保了生成质量不受影响，同时享受了预填充加速的好处。

K-Token Merging采用了一种巧妙的编码器初始化策略。编码器被初始化为表现得像均值池化，然后与LoRA适配器一起进行端到端训练。这种设计带来了几个优势：

• 稳定性：均值初始化提供了一个合理的起点，避免了训练初期的梯度不稳定
• 灵活性：LoRA适配器允许在保持基础模型冻结的情况下学习压缩策略
• 效率：轻量级编码器的计算开销很小，不会抵消压缩带来的收益

研究团队在三个基准测试上验证了K-Token Merging的有效性：Textualized Tree、Amazon Reviews和CommitPackFT。

以Textualized Tree基准为例，当合并因子K=4时，该方法实现了：

• 75%的输入长度减少：原始序列被压缩为原来的1/4
• 仅1.59%的准确率下降：在显著缩短序列的同时保持了极高的性能

这一结果证明，K-Token Merging成功地利用了潜在嵌入空间中的冗余，同时保留了模型的大部分推理能力。