推理链中的 Token 冗余与剪枝：消除无意义语气词对注意力权重的稀释

随着大模型深度推理（Reasoning Models）能力的演进，生成超长思维链（Long-CoT）已经成为解决复杂数学、符号规划与长代码生成的标准范式。

然而，在审视模型输出的长达数千 Token 的推导轨迹时，我们经常会看到大量高度冗余的过渡性表达，例如：“Let me pause and think...”、“Wait, is this correct? Let's check again...”、“Hmm, let me consider another perspective...”。

这些口语化的语气词虽然在某种程度上模拟了人类的思考节奏，但从信息论与 Transformer 注意力计算的物理机理来看，无意义的冗余 Token 会在自注意力矩阵中形成严重的注意力扩散，显著稀释关键逻辑实体的表征权重，并造成 30%~50% 的推理 FLOPs 浪费。

深入研究思维链中的 Token 冗余机理并进行动态剪枝，是实现极速低成本推理的关键。

一、语气词对注意力机制的物理侵蚀

在标准的因果多头注意力中，Softmax 归一化具有全局竞争性：

$$\alpha_{t, i} = \frac{\exp(Q_t K_i^T / \sqrt{d})}{\sum_{j=1}^t \exp(Q_t K_j^T / \sqrt{d})}$$

1[冗余 Token 导致的注意力权重被动稀释]
2关键题干约束: [变量 x > 0] (位置 5)
3真实代数推导: [方程 2x + 5 = 15] (位置 12)
4大量冗余语气: ["Let's see", "Wait", "Hmm", "Actually", "Let me check..."] (占据位置 13~150!)
5
6当生成第 151 步时:
7- 分母项累积了 130 多个低信息量语气词的 exp 点积得分;
8- 导致位置 5 的关键约束 [x > 0] 所分配到的有效注意力权重 alpha 从 0.35 骤降至 0.02!
9* 结果: 模型在冗余的自言自语中彻底遗忘了最初的几何边界约束。
10

二、Token 重要性评分的数学度量

为了准确识别并剪除思维链中的冗余 Token，可以从注意力流入（Attention Inflow） 与 梯度显著性（Gradient Saliency） 两个维度定义 Token 的重要性得分 $I(z_i)$：

$$I_{\text{attn}}(z_i) = \frac{1}{T - i} \sum_{t=i+1}^T \sum_{h=1}^H \alpha_{t, i}^{(h)}$$

$$I_{\text{grad}}(z_i) = \left| \frac{\partial \mathcal{L}}{\partial \mathbf{e}_{z_i}} \right|_2$$

高价值 Token：数学常数、变量符号、定理名称、方程操作符。后续所有时间步对其具有极高的持续注意力吸纳量；
低价值冗余 Token：纯过渡性语气短语与格式占位符。其注意力流入量迅速归零，且梯度敏感度极低。

三、动态推理剪枝架构（Inference Pruning）

1[在线思维链压缩与剪枝流]
2自回归生成推理步骤
3       │
4       ▼
5【重要性评估过滤器 (Token Saliency Scorer)】
6  ├── 实时测量历史 Token 的注意力聚集度与熵贡献
7  └── 识别连续低熵冗余块 (如连续 20 个占位语气词)
8       │
9       ▼
10【动态 KV Cache 淘汰与上下文收缩】
11  ├── 从 KV Cache 物理块中剔除冗余 Token 对应的 Key/Value 向量
12  └── 释放显存槽位，重置因果位置索引
13

通过这一机制，模型在保留严密逻辑骨架的同时，推导序列长度被大幅压缩 40%，且首字与尾字之间的因果注意力通路更加纯净。

四、PyTorch 代码实战：思维链 Token 敏感度与注意力分析

以下代码构建了一个轻量级分析器，能够提取序列各 Token 的历史注意力吸纳强度并自动化筛选冗余索引。

1import torch
2import torch.nn.functional as F
3import numpy as np
4from typing import List, Tuple
5
6def analyze_token_importance(
7    tokens: List[str],
8    attn_matrix: torch.Tensor, # [NumHeads, SeqLen, SeqLen]
9    threshold_ratio: float = 0.3
10) -> Tuple[List[int], List[float]]:
11    """
12    计算各 Token 的全局重要性得分并标记可剪枝的冗余位置
13    """
14    H, L, _ = attn_matrix.shape
15    # 对多头取平均: [SeqLen, SeqLen]
16    avg_attn = attn_matrix.mean(dim=0)
17    
18    importance_scores = []
19    for i in range(L):
20        # 统计从第 i+1 步到最后一步对位置 i 的平均注意力流入量
21        if i < L - 1:
22            inflow = avg_attn[i+1:, i].mean().item()
23        else:
24            inflow = avg_attn[i, i].item()
25        importance_scores.append(inflow)
26        
27    mean_imp = np.mean(importance_scores)
28    prune_indices = [idx for idx, score in enumerate(importance_scores) if score < mean_imp * threshold_ratio]
29    
30    return prune_indices, importance_scores
31
32if __name__ == "__main__":
33    # 构造模拟序列
34    simulated_tokens = [
35        "已知", "x", "=", "5", "。", 
36        "让我", "仔细", "想一想", "哈", "。", # 冗余语气块 (索引 5~9)
37        "计算", "x", "^", "2", "得到", "25", "。"
38    ]
39    L = len(simulated_tokens)
40    
41    # 模拟注意力矩阵: 因果下三角
42    torch.manual_seed(42)
43    mock_attn = torch.tril(torch.rand(4, L, L))
44    # 强化关键变量 x (索引 1) 的注意力流入
45    mock_attn[:, :, 1] += 3.0
46    mock_attn = mock_attn / mock_attn.sum(dim=-1, keepdim=True)
47    
48    prune_idx, scores = analyze_token_importance(simulated_tokens, mock_attn, threshold_ratio=0.5)
49    
50    print("================ 思维链 Token 敏感度分析 ================")
51    for idx, (tok, score) in enumerate(zip(simulated_tokens, scores)):
52        status = "✂️ 建议剪枝" if idx in prune_idx else "💎 核心保留"
53        print(f"Token [{idx:02d}]: {tok:8s} | 累积重要性得分: {score:.4f} | {status}")
54    print("=======================================================")
55

五、工程实践与对齐建议

SFT 阶段的“去废话”蒸馏：
- 在构建 Thinking 模型的微调数据时，应引入轻量级规则对标注数据中的无意义语气词进行预清洗，强制模型从一开始就习惯于输出高信息密度的紧凑逻辑链条；
推理引擎中的软剪枝（Soft-Prompt Masking）：
- 在使用 vLLM 部署长推理服务时，可以通过修改 FlashAttention 的注意力掩码（Mask），动态阻断注意力流向已标记为冗余的历史 Block，既保留了 KV Cache 的连续性，又消除了权重稀释。

《推理链中的 Token 冗余与剪枝：消除无意义语气词对注意力权重的稀释》是转载文章，点击查看原文。