CoLT: Reasoning with Chain of Latent Tool Calls

TL;DR

CoLT introduces a framework where reasoning is performed through explicit tool calls, combining the strengths of latent and token-based reasoning to improve efficiency and accuracy in large language models.

Contribution

CoLT is a novel framework that implements latent reasoning as explicit tool calls, avoiding model augmentation and enabling efficient, accurate reasoning.

Findings

Achieves higher accuracy than baseline latent models.

Produces shorter reasoning chains.

Compatible with reinforcement learning and various decoders.

Abstract

Chain-of-Thought (CoT) is a critical technique in enhancing the reasoning ability of Large Language Models (LLMs), and latent reasoning methods have been proposed to accelerate the inefficient token-level reasoning chain. We notice that existing latent reasoning methods generally require model structure augmentation and exhaustive training, limiting their broader applicability. In this paper, we propose CoLT, a novel framework that implements latent reasoning as ``tool calls''. Instead of reasoning entirely in the latent space, CoLT generates seed tokens that contain information of a reasoning step. When a latent tool call is triggered, a smaller external model will take the hidden states of seed tokens as its input, and unpack the seed tokens back to a full reasoning step. In this way, we can ensure that the main model reasons in the explicit token space, preserving its ability while improving efficiency. Experimental results on four mathematical datasets demonstrate that CoLT achieves higher accuracy and shorter reasoning length than baseline latent models, and is compatible with reinforcement learning algorithms and different decoder structures.