Guiding Language Models Reasoning with Planning Tokens

Thank you for your detailed review. Below is our response to the weakness section:

1. Compare to prefix tuning: Please refer to the general response. We have tried the prefix tuning baseline in the early stage of the project; they perform significantly worse than our current baselines: N/A (full fine-tuning/LORA), and General (adding the same prefix tokens in front of each step). We will include these results in our main table.
2. Title: While our proposed method has great potential on applying to various kinds of reasoning, our empirical results are currently focused on the math word problem datasets. We will restrict our title to math reasoning by changing it to: Guiding Language Models Math Reasoning with Planning Tokens
3. Contextual information of each step: We actually have already taken the contextual information into consideration, by using the contextualized embedding of each reasoning step produced by the base LLM. More specifically, we concatenate the question and reasoning steps, and encode the whole sequence with the base LLM. We then average the contextual embeddings of all tokens in a reasoning step. This makes the embedding of each reasoning step conditioned on the previous ones. We will clarify this important detail in the paper.
4. VQ-VAE v.s. SQ-VAE: We actually tried VQ-VAE first, and found it is very unstable to train on the step embeddings, so we switched to a soft version of it. We agree that this soft version of VQ-VAE is not a significantly novel contribution. Instead, this is an implementation choice for our proposed method.
5. Hierarchical generation: Thanks for the suggestion. We tried to adopt a sparse attention scheme in the earlier stage of the project to enforce a stronger dependency constraint between the planning tokens and the reasoning step that follows, similarly to what you are suggesting. More specifically, we modified the attention mask to make the generated step only depend on the corresponding planning tokens and the planning tokens to only depend on the previous planning tokens and the question, however the performance dropped quite dramatically, showing that the flexibility induced by full attention is important for final performance. Also, note that hierarchical generation usually comes at the cost of more difficult parallelization, given that two LLMs (planner and generator) might need to be kept in memory. All in all, these are interesting avenues to be explored, but the simplicity and effectiveness of our approach remain still attractive in our opinion.
6. Analysis of planning tokens: We kindly ask the reviewer to provide some more details on which kind of analysis is missing. In section 3.3, we first provide an analysis of the error rate versus the reasoning length. Then we provide a detailed error taxonomy. Finally, we provide a probing-based analysis of the inference planning token types on page 8, titled “Distinguishability of the Induced Clusterings”.

Below is our response to the question section:

1. Clustering: We only do clustering on the training set. The testing set is not used in our learning process.
2. Testing set subsampling: Because at testing time, the model will need to generate a long sequence of text as the solution, which makes testing slow. Considering our time and computing constraints, we decided to restrict the size of our testing sets to speed up our experiments. Thank you for pointing this out. We will clarify this experimental choice in the paper.
3. Ablation on all datasets: We would love to provide them on all datasets, but we weren’t able to do so because of our time and compute constraints, since these experiments are expensive to run.

Did you try prefix tuning with LoRa? In this setting, for prefix tuning the tunable tokens are randomly initialized and do not represent an explicit cluster. For SQ-VAE, the planned tokens are induced by your method. This is the only difference. LoRA is kept the same.
In this way, the number of tunable parameters should be the same.

Thank you for your positive review!

The MATH dataset is significantly more difficult compared to the other two datasets. It remains challenging for very strong LLMs, such as GPT-4.

• Data source: The MATH dataset comes from the AMC math competitions. In contrast, GSM8K comes from middle school and high school math exams; AQUA comes from multiple choice GRE math tests.
• Answer format: The answers of MATH data are free-form text, e.g., Latex expressions such as \frac{1-\sqrt{2}}{2}, which makes it hard for the model to get every bit of the answer correctly. In contrast, the answers of GSM8K are single int numbers; the answers of AQUA are single characters, indicating one of the choices.

When we summarize the error taxonomy, we take all three datasets into consideration. However, most of the mistakes in the MATH dataset fall in the wrong logic category, as the problems are not that obvious to approach, and the model seems to be confused about what the next step should be. Thanks for pointing this out. We will include an error taxonomy of MATH in the paper.

We thank the reviewer for pointing out the related works that we missed; we will adjust the text accordingly.

We emphasize that we are not claiming that tuning augmented tokens is a contribution of this work. Many existing works on prompt/prefix tuning [1-3] served as inspiration for our method. The main contribution of our work lies in a particular application of soft tokens for significantly improving the reasoning ability of modern LLMs.

Many differences distinguish our work from [2] and [3]:

• Latent token type: All previous approaches assume the ground-truth information about the token type to be given both for training and, crucially, testing. For example, [3] assumes knowledge of the relation type. In our work, we explore a more general setting, in which the token type is a latent variable and we train the reasoning model to directly predict the planning token at inference time. Therefore, the model needs to infer the planning token type from the previous reasoning steps. In one of our baseline settings, we use Arithmetic planning tokens, which is a particular case of token type inference informed by prior knowledge: the arithmetic operation can be understood as a relation type that links inputs and reasoning steps. Empirical results (Table 1) show that the planning tokens inferred by K-Means and SQ-VAE always provide better performance than the Arithmetic planning tokens, which indicate the importance of a learning-based token type inference algorithm.
• Fixed token location: Our planning tokens do not have a fixed location in the input prompt as in [1-3]. Instead, they will be freely generated by the model during the inference time, instead of being fed as input prompts. We only provide those inferred planning tokens in the training data, and they are added at multiple places of the sequence, instead of only in the front or in the middle.
• Integrability: Our method is supposed to be integrated into other training/fine-tuning paradigms (e.g. full fine-tuning or LORA) to improve LLM’s math reasoning capacity, instead of using on its own.
• Harder task: In this work, we test the proposed methods on a set of math reasoning tasks, which are arguably more difficult than the relation completion tasks used in [3]. Our method is designed to specifically enhance the reasoning abilities of modern decoder-only LLMs. This is in contrast to [3], which is built for relation completion tasks in the context of masked LMs. As shown in the general response, simple prompt/prefix tuning systems underperform in our setting.

To address the reviewer's statement on novelty, we reiterate our contributions which distinguish us from the related work raised in the review:

• As shown in our general response, traditional prompt/prefix tuning techniques do not work well on hard math reasoning tasks, our method redesigned the old paradigm and turned it into a strong auxiliary technique to improve the math reasoning capabilities of other fine-tuning paradigms.
• Planning tokens improve the intra-step consistency of the generated solution, and increase the model’s reasoning capacity by specializing it to different reasoning types.
• Our method is able to infer the type of latent planning tokens instead of using fixed, predefined types.

[1] Lester, Brian, Rami Al-Rfou, and Noah Constant. "The Power of Scale for Parameter-Efficient Prompt Tuning." Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. 2021.

[2] Li, Xiang Lisa, and Percy Liang. "Prefix-Tuning: Optimizing Continuous Prompts for Generation." Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers). 2021.

[3] Qin, Guanghui, and Jason Eisner. "Learning How to Ask: Querying LMs with Mixtures of Soft Prompts." Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT). 2021.

Thank you for your positive review. We have tried prefix-tuning baselines in the early stage of the project; they were significantly worse than our current baselines: N/A (full fine-tuning/LORA), and General (adding the same prefix tokens in front of each step). Please refer also to the general response. Thanks for pointing this out. We will update these results in our result section and add discussion.

We kindly ask the reviewer to further clarify which prompt-based methods they referred to as our potential baselines. We are happy to add them!

[1] Lester, Brian, Rami Al-Rfou, and Noah Constant. "The Power of Scale for Parameter-Efficient Prompt Tuning." Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. 2021.

[2] Li, Xiang Lisa, and Percy Liang. "Prefix-Tuning: Optimizing Continuous Prompts for Generation." Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers). 2021.

Dear reviewers, we have uploaded the revised version of our paper. We apologize for the wait as the Llama 13B model baselines are very slow to run. Below is the list of modifications (marked in blue in the paper draft):

1. Our title is changed to Guiding Language Models Math Reasoning with Planning Tokens.
2. We added a more in-depth discussion on the Bayesian motivation behind our method, which we hope will make our proposed planning token method better motivated.
3. We clarified that we use contextualized embedding of a step that encodes the information from the question and all previous steps.
4. We clarified the mathematical guarantees and motivations of using a soft-quantized VAE to infer the planning token types. We also clarified that we didn’t use VQ-VAE because it’s hard to train on our data.
5. We have added prompt-tuning and prefix-tuning baselines to the main table, which are 10+% lower than our original baselines.
6. We revised the error taxonomy section by manually verifying the GPT4-produced evaluations and added error taxonomy on both MATH and GSM8K.
7. We clarified that we subsample the testing set for time efficiency.
8. We discussed the soft-prompt tuning works in related work.

We want to thank you for taking the time to review our paper. We know it’s a busy moment, but we want to remind the reviewers that the discussion period is ending soon. We are happy to address any remaining concerns before the discussion period ends.