Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning
目录
1. 引言
2. 相关研究
2.1 Chain-of-Thought 推理
2.2 LLM 推理中的潜在空间
3. 方法
3.1 Token Assorted 方法
3.2 使用离散潜在标记进行推理
4. 实验
4.1 评测任务
4.2 主要实验结果
5. 未来研究方向
1. 引言
大规模语言模型（LLMs）在推理和规划任务中的表现已显著提升，尤其是在 链式思维（Chain-of-Thought, CoT）方法的辅助下。然而，现有的 CoT 方法依赖于大量文本标记（Text Tokens），这些标记虽然增强了模型的逻辑推理能力，但也导致了以下问题：
- 输入过长，计算开销大：大量的文本标记主要用于维持语言连贯性，而非核心推理信息。
- 训练成本高：由于需要处理完整的推理轨迹（Reasoning Traces），导致计算资源消耗显著增加。
- 泛化能力受限：传统方法难以高效适应新任务。
本文提出了一种新的混合表示方法——Token Assorted，即在推理过程中混合离散潜在标记（Latent Tokens）和文本标记。通过 向量量化变分自动编码器（VQ-VAE）抽掉（abstract away）推理轨迹的一部分，从而在减少输入长度的同时保持模型的推理能力。
研究目标包括：
减少推理轨迹的长度，降低计算成本。
提升 LLMs 的推理能力，在逻辑推理和数学推理任务上取得更好的表现。
探索混合标记训练方法，让 LLMs 快速适应新的潜在标记。
2. 相关研究
2.1 Chain-of-Thought 推理
CoT 方法通过将复杂问题拆解成多个推理步骤，提高 LLMs 在数学问题、逻辑推理和规划任务中的表现。然而，该方法的主要问题在于输入长度过长，影响模型的计算效率。
2.2 LLM 推理中的潜在空间
近年来，研究者开始探索 LLMs 在 潜在空间（Latent Space）中的推理能力。例如：
- 连续潜在空间推理：使用隐藏状态（hidden state）作为输入嵌入，使模型能够并行考虑多个推理路径
- 插入暂停标记（Pause Tokens）：在文本中插入“可学习暂停标记”，让模型在生成最终答案前进行额外计算。
- 编码器-解码器结构的潜在标记：使用 Transformer 结构在嵌入空间进行推理，但在现实数据集上的表现尚未超越 CoT。
本文方法不同于以上使用连续潜在空间的工作，直接使用离散潜在标记，从而提高推理效率，同时保持可读性。
3. 方法
3.1 Token Assorted 方法
Token Assorted 方法通过混合离散潜在标记（discrete latent token）和文本标记来压缩推理轨迹。具体步骤如下：
1）使用 VQ-VAE 训练潜在标记：
- 通过 VQ-VAE 将推理步骤（思维链）转换为离散潜在标记 Z = (z_1, z_2, ..., z_t)。
- 控制压缩率 r（即文本标记与潜在标记的比例），例如 r=16 代表每 16 个文本标记压缩为 1 个潜在标记。
2）替换部分文本标记：
设推理轨迹（思维链）标记 为 C = (c_1, c_2, ..., c_t)，我们用离散潜在标记 Z 替换其中前 m 个标记：
其中，P 代表问题提示（Prompt）标记，S 代表最终答案（Solution）标记，⊕ 代表拼接。
3.2 使用离散潜在标记进行推理
部分替换：替换 m 个标记，保留部分文本标记
从左到右替换：替换前 m 个标记。
随机混合不同替换比例（Randomized Mixing Strategy）：
- 以往的研究采用 多阶段课程学习（multistage curriculum training） 的方式，逐步增加 m 的值，直到达到预设的上限。然而，这种训练方法 计算复杂度高、效率低，需要精细的 超参数调优 来保证效果。
- 本文采用了一种 简单的单阶段训练策略，对每个样本随机设定 m 的值，而非逐步递增。这种方法不仅提高了训练效率，还提升了最终模型的性能。
4. 实验
4.1 评测任务
合成数据集：
- Keys-Finding Maze（路径规划任务）
- ProntoQA（逻辑推理任务）
- ProsQA（复杂逻辑推理任务）
数学推理任务：
- 训练数据集：MetaMathQA, Dart-MATH
- 测试数据集：
- GSM8K（基础数学）
- Math（高级数学）
- OlympiadBench-Math（奥数）
- Fresh-Gaokao-Math-2023（高考数学）
4.2 主要实验结果
Token Assorted 方法比 CoT在 Keys-Finding Maze 上 提升 19.8%，推理轨迹缩短 71%。
在 ProntoQA 和 ProsQA 上分别提升 1.2% 和 18.7%。
数学推理任务：大部分数据集上比表现更好。
标记数目： 在所有任务上使用比 CoT 更少的标记
表 4.4：我们的潜在标记随机替换策略明显优于其他选择：
- All-Replace：所有文本 CoT 标记都被潜在标记一次性替换
- Curriculum-Replace：在训练过程中逐渐用潜在标记替换整个 CoT 子序列的文本标记
- Poisson-Replace：以 0.5 的概率替换单个文本标记块
5. 未来研究方向
进一步优化编码器（VQ-VAE）结构，提高潜在标记的表示能力。
在更大规模数据集（如 Codeforces 竞赛数据）上测试 Token Assorted 方法的适用性。
探索 Token Assorted 在多模态任务（如视觉推理）的应用。
论文地址：https://arxiv.org/abs/2502.03275
进 Q 学术交流群：922230617