Guiding Language Model Reasoning with Planning Tokens

Summary

This paper introduces adding planning tokens in front of each chain-of-thought (CoT) step to guide LLMs on math reasoning. Three different ways to infer the planning tokens are proposed, in which a clustering method and a VAE-based method have shown effective on several math reasoning datasets. The paper is clearly written and the experiments are solid.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

The idea of using planning tokens to guide CoT is intuitive, and the proposed K-Means and SQ-VAE -based planning token inference methods have shown a good practice of this idea.

Reasons to reject

The experiments are limited to finetuning LLMs on a small scope of math reasoning datasets. In the proposed method, the planning tokens should be inferred from a training set using either clustering or generative methods. The scalability of the method may be a problem when training an LLM with multiple data resources of different distributions, as it may be exhaustive to obtain the planning tokens for each dataset.

Summary

This work describes a method for fine-tuning language models to take advantage of additional "planning tokens" injected ahead of reasoning steps when solving problems with multi-step intermediate reasoning. The authors inject a number of additional embeddings into models' input+output embedding matrices, then augment reasoning data by prefixing each reasoning step in a dataset with some number of planning tokens. The authors evaluate several means of assigning planning token identities when augmenting the training data: using a hand-designed heuristic (the arithmetic operations present in a step), k-means clustering of the step's average hidden states, or quantized variational autoencoder encodings of steps' average hidden states. The LM is then fine-tuned on the resulting data, optionally using a LoRA parameterization. The authors test their method on three math datasets and observe consistent gains over vanilla fine-tuning. Additional analysis reveals that models do indeed learn to use the planning tokens, and in fact pay more attention to planning token identities when they are generated using k-means or a VAE, indicating that planning token identity is exploited by tuned models.

Rating

8

Confidence

4

Ethics flag

1

Reasons to accept

- The method is intuitively appealing. - The authors back it up with both consistent improvements in downstream performance numbers and interesting analysis. - The paper's cleanly written and very easy to follow. - Points to neat directions for further work.

Reasons to reject

I don't see any issues that would merit rejection. The one minor weakness of the work is only evaluating on math reasoning, but that's a pretty common thing to do these days, so it really isn't that much of a problem.

Questions to authors

Typos: The caption for Table 2 contains "We set *n_special = n_special*", presumably this is not intentional. Suggestions: In §2.1, page 4, the SQ-VAE equations reference a *softmax* function that goes from latent vectors to the $q$ distribution - this includes a projection from the cluster vector dimension down to P bins, right? Might be good to clarify that this is a little more than the normal "softmax" function.

Summary

This paper proposes a Planning token to help generate reasoning steps more effectively in Math Reasoning tasks. The Planning token can be created using three methods (Arithmath, K-Means, SQ-VAE). Among them, the SQ-VAE method, which learns representations more intricately, exhibits the highest performance. While the proposed solution is interesting, it is rather difficult to understand the overall process of the proposed solution.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

1. The proposed solution proposes a reasoning path by refining step units. 2. It constructs prompts for error taxonomy and efficiently proceeds with GPT-4 utilization without human labor.

Reasons to reject

1. Overall, the proposed solution is difficult to understand. 2. It is necessary to show the effect of special tokens in evaluation. 3. It is necessary to compare the baseline solution using pause tokens. Refer to the following paper. [1] Goyal et al., Think before you speak: Training Language Models With Pause Tokens. CoRR abs/2310.02226 (2023)

Questions to authors

1. How does it construct clusters? Are the clusters also used in the test dataset? 2. Could you show more analysis for clusters?

Summary

This paper proposes a method to improve LLMs' chain-of-thought reasoning ability on math problems, by adding trainable "planning tokens" at the beginning of each reasoning step. The method is well-motivated, with consistent performance gains demonstrated by solid experimental results. The presentation is clear, which makes the paper easy to follow. However, the experiments involve only one task (math) and three datasets, which can limit the generalizability of the findings. Additionally, the analysis of attention weights may not directly support the causal claim made by the authors.

Rating

7

Confidence

4

Ethics flag

1

Reasons to accept

1. The tackled problem (math reasoning) is important and timely. 2. The method is well-motivated and thoroughly investigated with three ways of implementation, all of which are reasonable. The additional computation cost is minimal when integrated with existing parameter-efficient finetuning approaches (e.g. LoRA). 3. The experimental design is sound, with reasonable baselines and various base LMs. 4. The results show clear and consistent performance gains of the proposed method over baselines across all three datasets. 5. The authors conduct a comprehensive and thoughtful set of analyses to reveal further insights into the strengths and weaknesses of the proposed method. 6. The paper is well-written and easy to follow.

Reasons to reject

1. The experiments only involve one task (math reasoning) and three datasets, whereas existing works along this line often show applicability to more than one type of reasoning task, such as multi-hop QA, planning, logic inference, etc, in addition to math. 2. The performance gaps are sometimes small for certain settings in Table 1. It would be good to report statistical significance. 3. In Section 3.2 (Analysis), the authors report the performance of both Arithmetic and SQ-VAE in Figure 2, but the text only discusses the latter. The sentence “SQ-VAE outperforms the baseline especially for examples requiring longer reasonings” makes it sound like SQ-VAE is the strongest, but actually Arithmetic performs better on quite a few groups (# reasoning steps =4, 5,8, 9). I’m curious about how to interpret these mixed results. 4. In the analysis of attention on planning tokens, I doubt if the observation “planning tokens are assigned much higher attention weights than normal tokens” can support the claim “planning tokens are in general more helpful for the generation”. There has been a long debate on whether attention weights can be interpreted as token importance faithfully ([[1]](https://arxiv.org/abs/1902.10186), [[2]](https://aclanthology.org/D19-1002/), [[3]](https://aclanthology.org/2020.acl-main.432/), among others). A better alternative can be Attention Flow or Attention Rollout ([[4]](https://aclanthology.org/2020.acl-main.385/)) instead of raw attention weights.

Questions to authors

1. The notation "t_i" is overloaded with both "planning variable" and "planning token" and is inconsistent across the paper. In Figure 1, t_i refers to a planning token. In Section 2.1, t_i refers to a planning variable, and "each planning variable can be 'verbalized' by one or more planning tokens", which sounds confusing. 2. Potential typo in Table 2: "we set n_special = n_special" -> should be "n_prefix = n_special"? 3. Confusing sentence in the Conclusion -- “we might want to explore position-wise specialization of the planning tokens, i.e. use a different variant of the planning token w.r.t. the position of the step in the reasoning”: what does "variant" mean? 4. Which model's outputs were used in the error taxonomy analysis?