TokenSkip：当推理链遭遇选择性遗忘——可控思维链压缩的方法论与启示

# TokenSkip：当推理链遭遇选择性遗忘——可控思维链压缩的方法论与启示 > 2025 年 2 月，香港理工大学与中国科学技术大学的联合团队提出了一种名为 TokenSkip 的方法，通过分析 Chain-of-Thought（CoT）序列中 token 级别的语义重要性，实现了可控的推理链压缩。该方法在 Qwen2.5-14B-Instruct 上实现了 40% 的 token 减量，同时保持性能下降低于 0.4%，为推理效率优化提供了一条区别于训练时压缩和暴力截断的第三条路径。 --- ## 1. 问题背景：推理链的冗余性 ### 1.1 CoT 的效率困境 Chain-of-Thought（CoT）提示通过引导模型生成逐步推理过程显著提升了复杂任务的表现。然而，随着 OpenAI o1 和 DeepSeek-R1 等模型将 CoT 长度从数百扩展到数千甚至上万 token，推理延迟和计算成本呈线性增长。 | 模型 | 典型 CoT 长度 | 推理延迟影响 | 关键瓶颈 | |:---:|:---:|:---:|:---| | 传统 CoT | 50-200 tokens | 低 | — | | o1/R1 类模型 | 1,000-10,000+ tokens | **高** | KV Cache 内存、注意力二次成本 | > **核心矛盾**：更长的 CoT 可能带来更好的推理性能，但自回归解码的线性延迟增长使得长 CoT 在实际部署中成本高昂。 ### 1.2 现有压缩策略的局限 | 方法类别 | 代表工作 | 机制 | 核心问题 | |:---|:---|:---|:---| | 提示约束 | Prompt-based reduction | 在提示中要求模型缩短输出 | 压缩比不可控，模型响应不一致 | | 暴力截断 | Truncation | 限制最大生成长度 | 可能截断关键推理步骤 | | 训练时压缩 | DAST, MRT, LIMR | 在训练阶段优化推理长度 | 需要重新训练模型 | | **推理时压缩** | **TokenSkip** | **选择性跳过低重要性 token** | **需额外训练但成本低** | --- ## 2. Token 重要性：CoT 不是均匀的信息流 ### 2.1 语义重要性的差异性 Xia 等人通过 LLMLingua-2（Pan et al., 2024）对 CoT 序列中的 token 进行重要性评估，发现： | Token 类别 | 典型示例 | 重要性分布 | 功能 | |:---|:---|:---:|:---| | 数学表达式 | `26 - 5 = 21` | 高 | 核心推理步骤 | | 关键实体 | `Deanna`, `Marcus` | 高 | 问题要素识别 | | 数值 | `26`, `21`, `42` | 中高 | 计算节点 | | 逻辑连接词 | `so`, `since`, `therefore` | 低 | 语义衔接 | | 填充语 | `Let's break it down` | 极低 | 格式框架 | > **关键发现**：CoT 序列中的信息分布高度不均匀。数学公式和关键实体承载了推理的核心逻辑，而连接词和填充语主要服务于语言流畅性，对最终答案的贡献有限。 ### 2.2 重要性度量方法 TokenSkip 采用双向语言模型（BERT-like）进行重要性评估，而非因果语言模型的自回归困惑度： $$I_2(x_i) = P(x_i \mid \bm{x}_{\leq n}; \bm{\theta}_{\mathcal{M}_B})$$ > **方法选择依据**：因果语言模型的自回归注意力存在位置偏差——句子末尾的 token 由于获得了更多上下文信息，天然具有更高的模型置信度（更低的困惑度）。双向模型消除了这一偏差，能够更公平地评估每个 token 的独立重要性。 --- ## 3. TokenSkip 方法：选择性跳过与可控压缩 ### 3.1 Token 剪枝策略 给定 CoT 序列 $\bm{c} = \{c_i\}_{i=1}^{m}$ 和目标压缩比 $\gamma \in [0, 1]$，TokenSkip 执行以下操作： 1. 计算每个 token 的重要性 $I(c_i)$ 2. 确定重要性阈值：$I_\gamma = \text{percentile}([I(c_1), \dots, I(c_m)], \gamma)$ 3. 保留重要性高于阈值的 token：$\widetilde{\bm{c}} = \{c_i \mid I(c_i) \geq I_\gamma\}$ > **结果示例**：原始 CoT "Let's break it down step by step. Deanna is 26 years old..." 被压缩为 "break down Deanna 26 Marcus five younger 26 - 5 21..."——保留了所有数学关系和实体，去除了连接词和格式框架。 ### 3.2 训练范式 TokenSkip 的训练数据格式为： $$\mathcal{Q} \ [\text{EOS}] \ \gamma \ [\text{EOS}] \ \text{Compressed CoT} \ \mathcal{A}$$ 其中 $\gamma$ 从 $\{0.5, 0.6, 0.7, 0.8, 0.9, 1.0\}$ 中随机采样。训练时混合不同压缩比，使模型学会在全谱系上生成压缩推理链。 关键训练参数： | 配置 | 设置 | 优势 | |:---|:---|:---| | 微调方法 | LoRA (rank=8, α=16) | 仅训练 0.2% 参数 | | 训练数据 | GSM8K 7,473 + MATH 7,500 | 规模小，质量高 | | 训练时间 | 7B: ~2h, 14B: ~2.5h | 2×3090 GPU 即可 | | 保留原始 CoT | γ=1.0 的样本混合 | 防止过度压缩导致能力退化 | ### 3.3 推理时的可控性 推理阶段，用户可通过指定 $\gamma$ 来控制压缩程度： - $\gamma = 1.0$：原始 CoT，无压缩 - $\gamma = 0.7$：中等压缩，平衡效率与性能 - $\gamma = 0.5$：高压缩，最大化效率 > **可控性价值**：不同应用场景对效率和准确率的权衡需求不同。交互式应用可能偏好高压缩（低延迟），而关键决策场景可能选择低压缩（高准确率）。 --- ## 4. 实验评估：与基线的系统性对比 ### 4.1 GSM8K 上的性能 | 方法 | 目标压缩比 | 实际压缩比 | 准确率 | Token 数 | 延迟 | |:---|:---:|:---:|:---:|:---:|:---:| | 原始 CoT | — | — | 86.2% | 213 | 5.96s | | Prompt 0.5 | 0.5 | **0.89** ❌ | 83.7% | 189 | 4.97s | | Truncation 0.5 | 0.5 | 0.49 ✅ | **7.0%** 💀 | 104 | 2.95s | | TokenSkip 0.5 | 0.5 | **0.53** ✅ | **78.2%** ✅ | 113 | 3.40s | | TokenSkip 0.7 | 0.7 | **0.70** ✅ | **82.5%** ✅ | 150 | 4.36s | > **基线分析**：Prompt-based reduction 的实际压缩比严重偏离目标（目标 0.5，实际 0.89），说明模型对压缩指令的理解不一致。Truncation 虽然压缩比准确，但 0.5 压缩导致准确率从 86% 暴跌至 7%——验证了**盲目截断会摧毁推理能力**。TokenSkip 在压缩比准确性和性能保持之间取得了最佳平衡。 ### 4.2 模型规模的扩展性 Qwen2.5-Instruct 系列在 GSM8K 上的表现： | 模型 | 压缩比 | 准确率变化 | 关键观察 | |:---|:---:|:---:|:---| | 7B | 0.6 | -5.1% | 较小模型压缩敏感度较高 | | 14B | 0.6 | **< -0.4%** | **大模型几乎不受压缩影响** | | 14B | 0.5 | -2.0% | 即使高压缩，性能仍稳健 | > **规模效应**： larger models 在压缩后的表现更稳定，表明它们更擅长在稀疏表示中识别关键推理路径。 ### 4.3 MATH-500 上的挑战 | 方法 | 压缩比 | 准确率 | 性能变化 | 加速比 | |:---|:---:|:---:|:---:|:---:| | 原始 | — | 48.6% | — | 1.0x | | TokenSkip | 0.7 | 46.7% | -1.9% | **1.4x** | | TokenSkip | 0.5 | 40.2% | -8.4% | **1.7x** | > 在更难的 MATH-500 上，TokenSkip 0.7 实现了 30% 压缩 + 1.9% 性能下降 + 1.4x 加速，展现了跨难度任务的稳健性。 --- ## 5. 压缩 CoT 的可恢复性：效率与可解释性的兼得 TokenSkip 的一个独特优势是**非破坏性压缩**。研究者验证了模型可以从压缩后的 CoT 恢复完整的推理过程： | 版本 | 内容 | |:---|:---| | 压缩版 | "break down Deanna 26 Marcus five younger 26 - 5 21 Marcus half Leo's age twice Marcus Marcus 21, Leo's age 2 x 21 = 42" | | 恢复版 | "Let's break it down step by step. Deanna is 26 years old. Marcus is five years younger than Deanna: M = D - 5. Marcus's age: M = 26 - 5 = 21..." | > **方法论意义**：这证明了压缩 CoT 保留了足够的语义信息，使得恢复完整推理链成为可能。在实际部署中，可以先用压缩版快速获取答案，在需要审计或解释时恢复完整版本。 --- ## 6. 讨论：TokenSkip 在推理效率谱系中的定位 ### 6.1 三条路径的对比 当前推理效率优化存在三条互补路径： | 路径 | 代表方法 | 干预时机 | 核心思想 | 复杂度 | |:---|:---|:---:|:---|:---:| | **训练时优化** | DAST, MRT, LIMR | 训练阶段 | 改变模型的推理行为模式 | 高 | | **后处理压缩** | TokenSkip | 后训练微调 | 教会模型选择性跳过 token | 中 | | **架构/解码优化** | Speculative Decoding, Medusa | 推理阶段 | 加速解码过程 | 中 | > TokenSkip 的优势在于**低门槛和可控性**。它不需要重新设计训练流程（如 MRT），也不需要修改模型架构（如 Medusa），仅需轻量级的 LoRA 微调即可实现显著效率提升。 ### 6.2 开放问题 1. **动态压缩**：当前 TokenSkip 使用全局压缩比。能否根据推理进程动态调整——在关键步骤保留更多 token，在常规步骤允许更高压缩？ 2. **自评估重要性**：当前依赖外部双向 LM（LLMLingua-2）评估 token 重要性。模型能否学会自评估并动态决定跳过策略？ 3. **跨领域泛化**：数学推理上的成功是否可迁移到代码生成、科学推理、多模态推理等领域？ 4. **与训练时方法的协同**：TokenSkip 与 DAST 的难度自适应预算能否结合，实现"先自适应分配预算，再在预算内选择性跳过"？ --- ## 7. 结论 TokenSkip 通过系统分析 CoT 序列中 token 级别的语义重要性，证明了推理链中存在显著的可压缩冗余。其核心方法论——基于重要性的选择性跳过——在保持推理性能的同时实现了可观的效率提升，且实现成本极低。 在推理成本日益成为 LLM 部署瓶颈的背景下，TokenSkip 代表了一种务实且可扩展的解决方案：不需要重新训练模型，不需要修改架构，只需要教会模型"跳着思考"。 --- ## 论文详情 | 项目 | 内容 | |:---|:---| | **标题** | TokenSkip: Controllable Chain-of-Thought Compression in LLMs | | **作者** | Heming Xia, Yongqi Li, Chak Tou Leong, Wenjie Wang, Wenjie Li | | **机构** | The Hong Kong Polytechnic University, University of Science and Technology of China | | **arXiv ID** | 2502.12067 | | **日期** | 2025-02-17 | | **核心贡献** | Token 重要性分析；可控 CoT 压缩；低训练成本；压缩可恢复性 | | **关键结果** | Qwen2.5-14B GSM8K: -40% token, < 0.4% 性能下降；LLaMA-3.1-8B MATH-500: -30% token, 1.4x 加速 | | **训练配置** | LoRA (rank=8, α=16)，0.2% 参数；2x 3090 GPU；2-2.5 小时 | | **代码** | https://github.com/hemingkx/TokenSkip | #Research #ChainOfThought #Compression #Efficiency #TokenSkip #智柴 🔬