Token Assorted: Mixing Latent and Text Tokens for Improved Language Model Reasoning

大型语言模型（LLMs）在使用链式思维（CoT）数据训练时，表现出出色的推理和规划能力，其中每一步的思考过程通过文本标记明确列出。然而，这导致输入过长，许多词语用于支撑文本连贯性而非核心推理信息，处理这些输入消耗了大量计算资源。在这项工作中，我们提出了一种推理过程的混合表示方法，通过VQ-VAE生成的潜在离散标记部分抽象掉初始推理步骤，显著缩短了推理轨迹的长度。我们在两种场景中探讨了潜在轨迹抽象的应用：1) 从头开始训练模型以解决钥匙寻找迷宫问题；2) 使用包含未见过的潜在标记的扩展词汇表对LLMs进行微调，以解决逻辑和数学推理问题。为了促进有效学习，我们引入了一种简单的训练程序，随机混合潜在标记和文本标记，从而实现对新潜在标记的快速适应。我们的方法在各种基准测试中始终优于基线方法。

论文试图解决大型语言模型（LLM）在处理链式思维（CoT）数据时输入过长的问题，这些冗长的输入包含许多支持文本连贯性而非核心推理信息的词语，导致计算资源消耗巨大。这是一个在提高LLM效率和性能方面的新挑战。

关键思路是引入了一种混合表示方法来抽象初始推理步骤，通过使用VQ-VAE生成的潜在离散令牌部分替代文本令牌，从而显著缩短推理痕迹的长度。这种方法不仅减少了输入长度，还保持了推理过程的关键信息。相比现有的研究，该方法创新地结合了潜在变量和文本令牌，为更高效的推理提供了一种新途径。

论文的亮点包括：1) 在两个不同场景下验证了潜在追踪抽象的有效性；2) 提出了一种简单但有效的训练程序，能够快速适应新的潜在令牌；3) 实验设计涵盖了从头训练和微调现有LLM的任务，证明了该方法的广泛适用性。此外，论文展示了其方法在多个基准测试中的一致优越表现。虽然没有提到是否开源代码，但实验使用的Keys-Finding Maze问题和逻辑数学推理问题的数据集值得关注。

最近在这个领域相关的研究包括《Improving Neural Comprehension with Chain-of-Thought Reasoning》和《Latent Variable Models for Neural Machine Translation》等论文。这些研究探索了如何通过改进模型结构或引入新的训练方法来增强LLM的推理能力。本论文的独特之处在于它将潜在变量应用于推理过程中，这与上述工作形成了补充关系。