LightThinker：大模型思考压缩

现在的大模型为了解开复杂的数学或逻辑题，会生成非常长的“Chain of Thinking”。注意力机制的计算复杂度随上下文长度呈二次方增长，而 KV 缓存的存储开销则随上下文长度线性增加。因此，生成 token 数量的增加导致内存开销和计算成本急剧上升。

作者观察到：LLM 生成的 token 具有双重目的：既保证语言流畅性，又促进实际推理。当人类解决问题时，通常只在草稿纸上记下关键步骤，而将剩余的思考过程保留在头脑中。

基于以上见解，为了解决这个问题，研究者引入了 LightThinker，一种在生成过程中动态压缩中间思路的方法。在生成一个冗长的思考步骤后，它会被压缩成一个紧凑的表示（如 C Ti），而原始的思考链被丢弃，后续的推理仅基于压缩后的内容继续进行。

数据集重构

为了使 LLM 能够在生成过程中动态压缩，研究者重构了数据集 $D$。对于给定的输入输出 $(X,Y)$ ，作者将 $Y$ 分为 $k$ 个子序列 $S$，即 $Y = {S_i}$。该函数可以基于 token 级别或思路级别。

然后，作者插入特殊 token。

• <w> 是压缩触发器，表示需要压缩前面的思路。
• C = {[ci]} 由 |C| 个特殊 token 组成，作为存储压缩内容的 token。这里我们将 C 称为 cache tokens。 令牌 [o] 是一个强制性的输出令牌，它使得模型能够基于压缩内容继续生成。
• [o] 是一个强制性的输出令牌，它使得模型能够基于压缩内容继续生成。

最终构建的新数据集形如 Y' = {S_1, <w>, C, [o], S_2, <w>, C, [o],...,S_k} 。

基于思路的注意力掩码构建

为了让 LLM 学习如何压缩以及如何基于压缩内容生成（也就是如何理解压缩内容），

• 在压缩期间，$C_i$ 只能关注问题 $X$、先前压缩的内容 $C_j(j<i)$ 和当前思路 $S_i$。
• 在生成期间，$S_i$ 只能关注问题 $X$、先前压缩的内容 $C_j(j \le i)$。