Guiding Language Model Reasoning with Planning Tokens

TL;DR

This paper introduces a hierarchical planning token approach to guide large language models in reasoning tasks, improving accuracy with minimal additional parameters across multiple datasets.

Contribution

It proposes a novel hierarchical generation scheme using planning tokens to structurally guide reasoning, requiring negligible additional parameters.

Findings

Significant accuracy improvements on math and QA datasets

Effective across multiple large language models

Minimal increase in trainable parameters

Abstract

Large language models (LLMs) have recently attracted considerable interest for their ability to perform complex reasoning tasks, such as chain-of-thought (CoT) reasoning. However, most of the existing approaches to enhance this ability rely heavily on data-driven methods, while neglecting the structural aspects of the model's reasoning capacity. To encourage a more structural generation of CoT steps, we propose a hierarchical generation scheme: we let the LM generate a planning token at the start of each reasoning step, intuitively serving as a high-level plan of the current step, and add their embeddings to the model parameters. Our approach requires a negligible increase in trainable parameters (0.001%) and can be applied through either full fine-tuning or a more parameter-efficient scheme. We demonstrate our method's effectiveness by applying it to three different LLMs, showing notable accuracy improvements across three math word problem datasets and one multihop QA dataset with respect to standard fine-tuning baselines.