Guiding Language Model Reasoning with Planning Tokens

Authors: xinyi  wang
caccia  lucas
oleksiy  ostapenko
xingdi  yuan
william yang wang
alessandro  sordoni
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-21 14:26:21

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.

Main paper: 
 introduction: 
 The great potential of solving complex reasoning problems, including world knowledge reasoning (Hendrycks et al., 2020; Suzgun et al., 2022) , logical reasoning (Pan et al., 2023) , and math reasoning (Cobbe et al., 2021; Hendrycks et al., 2021b) , using pre-trained large language models (LLMs) (Touvron et al., 2023a; b; Brown et al., 2020) has drawn much attention recently. A popular and effective paradigm of reasoning with LMs is chain-ofthought (CoT) reasoning (Wei et al., 2022; Wang et al., 2022) . In CoT the LMs are required to generate both the reasoning steps and answer for a given problem. Recent research has recognized the importance of CoTs for complex reasoning problems. Multiple works focus on augmenting high-quality alternative CoTs in training data. For example, Yue et al. (2023) fine-tune LLMs on multiple math datasets with CoT and program-of-thought (PoT) solutions. Yuan et al. (2023) applies rejection sampling on the LLM samples. Other works elicit reasonings from exogenous resources, such as more capable LLMs, i.e. GPT-4 (Mukherjee et al., 2023; Luo et al., 2023) . Although these methods are shown to be highly effective, the fundamental problem of how to effectively guide an LM to generate more useful CoT is not directly tackled by these data-augmentation-based methods.We aim to propose a lightweight training method to fundamentally improve the effectiveness of the generated CoTs. While current methods usually directly generate the complete CoT solution, we hypothesize that a hierarchical generation of the CoT steps will benefit the overall quality of the solution. More specifically, we introduce planning tokens at the beginning of each CoT step, which are special tokens that encode a high-level solution plan of the current CoT step. We consider three ways of inferring the planning tokens: heuristic assignment, clustering reasoning states, and latent variable learning with VAE (Kingma & Welling, 2013) . At training time, we first assign planning tokens to each CoT step and Figure 1 : An overview of our method: Given an input question x and its CoT steps r 1 , ..., r i , we train a planning token inference model to infer the planning tokens t i . Subsequently, we fine-tune an LM on data augmented with the inferred planning tokens. Given a new question x ′ , the trained LM can infer planning tokens t ′i , CoT steps r ′i , and final answer y ′ . Fire and snowflake denote trainable and frozen models, respectively. Note that x, r i , t i , y can all include multiple tokens, we depict them as single blocks/variables for simplicity.then do regular supervised fine-tuning. At inference time, the LM will be able to generate the corresponding planning tokens first, before generating each CoT step. In practice, the planning tokens' embeddings add negligible additional learnable parameters to the LLMs (< 0.001%). This makes our method integrate well with parameter-efficient fine-tuning methods like LoRA (Hu et al., 2021) .While the idea of adding new tokens to the generative LM's vocabulary and then training the associated embeddings has been explored before (Li & Liang, 2021; Lester et al., 2021) , the function and effect of our proposed planning tokens are significantly different from the previous works. Our planning tokens are designed to increase and guide the reasoning ability of LM fine-tuned with other supervised fine-tuning methods, instead of acting as a parameter-efficient fine-tuning method (Li & Liang, 2021; Lester et al., 2021) on its own. The fact that planning tokens are specialized within one task and can be generated and inferred by LMs like regular tokens also sets us apart from the previous soft token tuning based methods. For example, Qin & Eisner (2021) treat the soft tokens as special indicators for specific relations, and Goyal et al. (2023) fixed the position of the pause token to right after the prompt and serves only as extended computation before generation.We perform experiments on three math word problem (MWP) datasets (GSM8K (Cobbe et al., 2021) , AQUA (Ling et al., 2017) , and MATH (Hendrycks et al., 2021a )), and one multihop QA dataset (StrategyQA (Geva et al., 2021) ). The experiment results in Section 3.1 show that by adding planning tokens at training time, we are able to improve upon the baseline without planning tokens by 3.3% accuracy points on average over three pre-trained language models (Phi 1.5 (Gunasekar et al., 2023) , 7B Llama 2 (Touvron et al., 2023b) , and 13B Llama 2) and three MWP datasets. We also show that the gains of our method do not solely come from the augmented computational space induced by increasing the length of the reasoning with planning tokens but also from the specialization induced by the planning tokens. In Section 3.2, from a detailed error analysis of the generated CoT solutions, we found that our method performs notably better on longer CoT solutions. The error type analysis also suggests that our method can better refer to the previous information and increase the computing capacity of the model. We also investigate how LLMs utilize the planning tokens by looking into the attention weights, which suggest that planning tokens can carry information throughout the whole sequence. 
method: 
 Setup We assume that we have a dataset D composed of triples {(x, r, y)}, where x is a problem to solve, r is the ground-truth reasoning associated with the problem and y is the final answer. The assumption that the ground-truth reasoning is included in the dataset might be circumvented by either generating CoTs with rejection sampling from a base LM (Yuan et al., 2023) or by creating reasoning paths from a larger model (Luo et al., 2023; Mukherjee et al., 2023) . We fine-tune a base LLM to predict the concatenation of reasoning and answers given the inputs, i.e.:arg max θ ∑ D log p(r, y|x; θ),where we assume that x, r, y are properly tokenized and θ is the trainable parameter of the pre-trained LLM. To provide more high-level guidance to the LLM at fine-tuning time with reasoning data, we assume that each reasoning r is the realization of some higherlevel planning process, which is latent, thus unobserved in the dataset. To model this assumption, we define by reasoning step a sub-sequence of tokens appearing in the reasoning. We create reasoning steps by splitting the reasoning with a delimiter token, which might be dataset-specific, but for simplicity, we fix it to be \n for all datasets.Let r = r 1 , . . . , r S , where S is the number of steps in the reasoning. We assume that each reasoning step r i is generated conditional on the previous reasoning steps and a discrete latent planning variable t i taking values into 1 . . . P. Given that the planning variables value between 1 and P, they can be considered as condensed summaries of each reasoning step and can therefore provide useful guidance while producing the reasoning r. We also assume the answer y is generated with a special planning variable t y which has always a value equal to P for each example. Without loss of generality, we assume that only t y can take value P, i.e. p(t i = P) = 0 for each t i , i ̸ = y.If the planning variables were indeed observed, the complete data D c would be composed of tuples (x, t 1 , r 1 , . . . , t S , r S , t y , y). Let us assume for the moment that the complete dataset is available to us. In order for the dataset to be processed by an LLM, we verbalize each planning variable t i with a (or possibly multiple) planning token(s). Here we abuse the notation t i for this (set of) planning token(s) for simplicity. In practice, this means that we extend the vocabulary of the LLM tokenizer with P (or more) new tokens of our choice, and modify accordingly the input and output embedding matrix of the LLM. The embeddings of the new tokens will be additional training parameters for the LLM, which add a negligible parameter overhead over the base model.With this modification in mind, the complete data D c can be readily used to train an LLM by maximizing the complete data log-likelihood:arg max θ + ∑ D clog p(t 1 , r 1 , . . . , t S , r S , t y , y|x; θ + ), by θ + we denote the original LLM parameters comprising the embeddings of the planning tokens. From the equation above, we can see that the LLM will be trained to first generate the high-level planning token for each reasoning step, and then generate the detailed reasoning step. This can help LLMs to generate more consistent and controlled reasoning chains.We consider three different ways of inferring the planning variable from the given CoT training data: Arithmetic, K-Means, and SQ-VAE. In practice, this inference step corresponds to assigning the instantiated planing tokens to the reasoning steps prior to fine-tuning the LM and learning the embeddings of the planing tokens.Arithmetic For math word problems, it is natural to consider the basic arithmetic operation contained in each reasoning step r i as the plan token t i similar to Zhang et al. (2023) ; Qin & Eisner (2021) . This yields four planning tokens (+,-, ×, ÷), which can be easily extracted from each r i . If more than one operator is found, then we use as many planning tokens as the number of operations found, e.g. <+><*>.This approach has several drawbacks. First, some operations, such as doubling a quantity, can be implemented with more than one operator (e.g. + and ×). Thus, this arithmetic partitioning can fail to capture the similarity between these reasoning steps. Second, for some more complex datasets such as MATH, where the reasoning goes beyond basic operators, such heuristics might not be simple to design if not provided by the dataset. Finally and most importantly, such a heuristic does not apply to non-math datasets like multihop QA datasets. To relax these shortcomings, we turn to inference planning tokens from the embeddings of CoT steps.We first gather reasoning steps r i for all examples in our dataset, then apply an inference algorithm on the representations ϕ(r i ), ∀r i ∈ D. We obtain ϕ(r i ) by using an LLM to encode the whole text sequence of an example (x, y), and then averaging the last hidden layer outputs of the corresponding tokens in r i . In this way, we aggregate both the question and previous steps' information in the contextualized representation ϕ(r i ).We run an off-the-shelf K-Means implementations on the set of representations ϕ(r i ) for all reasoning steps in the dataset D. Each planning variable t i is then assigned the index of the centroid closest to r i .In this approach, our aim is to learn a non-linear transformation of each representation ϕ(r i ) that better captures the meaningful directions of variation of the representations for the dataset D under consideration. Variational Autoencoders (VAEs) (Kingma & Welling, 2014) offer a probabilistic approach to learning such non-linear latent representations of data. In their original formulations, VAEs learn a Gaussian-distributed latent space of representations. To induce a discrete structure in the latent space, we follow Miao et al. (2017) , and use a "Gaussian-softmax" parameterization, which soft-quantizes the latent representations before reconstructing the input data. In particular, our VAE maximizes the following objective:L(enc, dec, Φ) = E enc(z i |ϕ(r i )) [||dec(ϕ(r i )| z) -ϕ(r i )|| 2 ] -KL(enc(z i |ϕ(r i )), N (0, I)), z = Φq, q = softmax(z i )where enc and dec are the encoder and the decoder neural networks that transform the steps representation into the non-linear manifold and back into the original representation space respectively (the encoder parameterizes mean and log-variance of a Gaussian distribution); q is the soft-quantization of the Gaussian-distributed representation z i , which corresponds to the distribution over clusters. Φ is a learnable matrix with P -1 rows (the last value is reserved for the answer planning variable) representing the cluster centroids. Therefore the model tries to reconstruct the input representation given a soft-combination of cluster centroids. With this model, after training, we obtain the assignment for the planning variables as t i = arg max softmax(enc(ϕ(r i ))), where enc(ϕ(r i )) denotes the mean of the Gaussian distribution parameterized by enc.Each planning variable may be verbalized by one or more planning tokens. Moreover, we discussed two orthogonal intuitive effects that might be achieved by augmenting the dataset with planning tokens. The first is to augment the computational capacity of the model by providing additional "scratch" space to predict the next reasoning step; the second is the specialization induced by information-bearing planning tokens. To be able to disentangle both effects in our model, we introduce two hyper-parameters in our approach: n prefix and n special. Each planning variable can be verbalized using a variable number of both generic prefix planning tokens and special planning tokens. A planning token annotated example with n prefix = n special = 3 is shown below: (Ling et al., 2017) contains 100K samples of mathematical problems, along with sequences of human-readable mathematical expressions in natural language. StrategyQA contains 3K multi-hop questions annotated with decomposed single-hop questions, which we used as the Chain-of-thought (CoT) path of the question.Base LLMs Our empirical analysis uses several decoder-only architectures of varying sizes. We use the 7B and 13B variants of Llama2 (Touvron et al., 2023b) , both trained over 2 trillion tokens from publicly accessible data sources. We also experiment with Phi-1.5 (Gunasekar et al., 2023) , a 1.3B parameter model trained on a mixture of textbook-quality code data, and additional synthetically generated textbook and exercise data.Baselines First, we compare against the vanilla baselines of full-fine-tuning (Full-FT) for Phi-1.5 and LoRA fine-tuning for Llama-2. Then, we compare to a baseline that combines plain soft token tuning (General) with the base fine-tuning method (Full-FT/LoRA). For this baseline, we prepend the same prefix tokens to each CoT step. i.e. P = 1.Note that the General baseline can be viewed as prompt-tuning (Lester et al., 2021) + the base fine-tuning method, when the prefix tokens are only prepended to the first CoT step.To make prompt-tuning more comparable to our planning token methods, we enhance the original prompt-tuning by also prepending prefix tokens to other CoT steps.We also tried the prompt-tuning (Lester et al., 2021) and prefix-tuning (Li & Liang, 2021) only parameter efficient tuning baseline. We found that they are not suitable for complex CoT reasoning tasks in our setting, as their capacity is limited when using a reasonable number (similar to our method) of tunable tokens. The inference time and memory complexity become unreasonable when we try to scale the trainable parameters to be similar to LoRA (3k tunable tokens). Detailed results can be found in the Appendix.Our Methods Our methods are denoted with the three different planning types: Arithmetic, K-Means and SQ-VAE. Table 2: Impact of varying the number of clusters (left) and varying the number of planning tokens (right) (n prefix+n special). We set n prefix = n special. For general, we match computation by adding prefix tokens. 
results: 
 We present our main results in Table 1 . Generally, we observe that for all three datasets considered and all the model sizes, the best-performing approach leverages planning tokens. We note that, across scales, Full-FT + General and LoRA + General improves over vanilla fine-tuning (Full-FT or LoRA), echoing our understanding from Chi et al. (2023) and Feng et al. (2023) that adding additional tokens before each reasoning step increase the compute capacity of the LM and results in better performance. For the three datasets, Arithmetic does not consistently outperform General pointing to the fact that hand-designed heuristics might not be optimal and the gain observed over Full-FT/LoRA may due to additional allowed computation space. However, the other two embedding-based planning type inference methods, K-Means and SQ-VAE, consistently outperform both General and Arithmetic, pointing to the importance of using machine-learned planning tokens specialization. Over the 7B version of Llama2, all methods with planning tokens provide consistent gains over the baselines of up to 1.7 points, a relative increase of 6%. Our results over the Llama2 13B tell a similar story, where again the best-performing method across all datasets uses planning tokens. We show a 1.5 gain in average accuracy (4.5 % relative gain). Among these two approaches, SQ-VAE seems to be a more reliable approach, as is it always the bestperforming method on average. This can be understood as the non-linear latent inference method works slightly better than the linear one. More expressive latent inference methods might be needed to obtain a larger performance improvement.We show ablation studies with respect to the number of clusters P and the number of planning tokens used in total. A larger P means we classify reasoning steps into more fine-grained planning types. Note that P = 1 means we always use the same planning tokens for all steps which is the same as the General baseline. From Table 2 (left), We observe that the performance of both K-Means and SQ-VAE first goes up and then goes down when the number of clusters increases. This likely reflects that our planning token inference models are not able to produce too fine-grained reasoning type. From Table 2 (right), we observe that it is not necessary to assign too many tokens to a planning type. We suspect that too many planning tokens will result in a longer sequence which negatively affects the language model's reasoning ability.Errors by Reasoning Length In Figure 2 , we show the accuracy by the complexity of the test example as measured by reasoning length on AQUA and GSM8K. For the two datasets, the introduction of planning tokens improves long reasoning performance, thus hinting that the planning tokens provide better control over long reasoning chains. For GSM8K, planning tokens seem to underperform the baseline for short reasoning chains. While both Arithmetic and SQ-VAE improve on long reasonings, SQ-VAE improves more on average.In future work, we can assign position-aware planning tokens for different steps, to enhance the long reasoning improvement.Error Taxonomy After inspecting the generated chain-of-thoughts solutions from all three datasets, we propose to classify the errors into the following five categories:• Misunderstanding of question: generated solution does not answer the given question. • Computation errors: errors in evaluating equations.• Inaccurate extraction of question information: inaccurate references to the information provided in the question.• Wrong application of math knowledge: equations, concepts, or theorems used in the generated solution are not suitable for the question.• Wrong logic/unreasonable step: generation does not follow common logic.We prompt GPT4 to obtain the error type of each generated example, and then manually verify whether the predicted error type matches the explanation. In Figure 3 , the SQ-VAE error frequencies significantly lower than the N/A baseline are in bold font. As we can see, across three datasets, our method consistently improves on the computation errors and the inaccurate extraction of question information. The lower computation errors indicate that our inferred planning tokens either increase the computation capacity of the model or contain compute-related information. Our method being able to make more faithful reference to the information in the questions indicates planning tokens are helpful for longrange control of generations, which echoes the observation and assumptions in the previous sections. The differences in other error categories are either insignificant or inconsistent over datasets. Note that GSM8K and AQUA have a similar distribution over the error types, while MATH has significantly more misunderstanding of questions and wrong application of math knowledge. This is understandable as MATH is significantly more difficult as it is from American Mathematics Competitions (AMC). These questions are less obvious to approach and require applications of non-trivial math theorems.To better understand how LLMs make use of the planning tokens, we inspect the attention by first computing an overall average of the attention weights assigned to the planning tokens versus the normal tokens. As shown in Table 3 , we found that K-Means and SQ-VAE planning tokens are assigned much higher attention weights than General and Arithmetic planning tokens, when comparing to normal tokens. This implies that the inferred planning tokens might be more helpful for the generation, which is consistent with the main results performance in Table 1 .While the raw attention weight itself might be a debatable way of understanding the token importance, the attention pattern still serves as a valid way of understanding how the Transformer works. Similar to Olsson et al. (2022) , we identify attention heads that have strong patterns corresponding to the planning tokens as shown in Figure 4 , and deduct how language models make use of the planning tokens from the patterns.More specifically, we compute the average attention weight difference between planning tokens and normal tokens from different attention heads, and visualize the attention heads with the largest difference, as they are likely to strongly correspond to the utilization of planning tokens. We visualize two such attention heads on a CoT generated with SQ-VAE planning tokens in Figure 4 . The attention head shown on the left strongly attend to the specialized planning tokens throughout the CoT sequence. The attention head shown on the right shows that the general planning tokens are only attended by themselves, which might imply that the general planning tokens only serve as indicators of the beginning of a CoT step, instead of carrying information to the whole sequence.We investigate whether SQ-VAE learns better planning types than K-Means via a probing task (Alain & Bengio, 2017) . Concretely, we collect a dataset consisting of CoT steps and their planning token labels obtained from K-Means and SQ-VAE. We train a simple model (either a logistic regression or a shallow neural network) to learn the mapping from a sentence to its planning token label. We hypothesize that the better planning type should facilitate easier learning dynamics for simple models because the class boundaries should be easily accessible from the data. As a case study, we study the 5-cluster setting on the GSM8K dataset. Given a planning type, we randomly sample 2000 sentences from each of the five categories. We leverage a pre-trained text encoder2 to convert sentences into vector representations. Taking the text encodings as input, the probe network (logistic regression, 2-layer MLP, or 6-layer MLP) performs a 5-way classification.We show test accuracy curves in Figure 5 . We observe a clear gap between SQ-VAE (solid lines) and K-Means (dashed lines) in the curves. This reflects our main results reported above, where SQ-VAE produces a better overall downstream task performance. 
related work: 
 Trainable new tokens The idea of adding new tokens with trainable embeddings on the input side of a generative LM has been proposed before (Li & Liang, 2021; Lester et al., 2021) . The most common way of adding new tokens is to insert them at a fix position in the prompt given to LMs. By only training the embeddings of these new tokens, these methods act as parameter-efficient fine-tuning techniques to adapt LMs to specific task (Qin & Eisner, 2021; Li & Liang, 2021; Lester et al., 2021) . Our planning tokens are not intended to serve as a parameter-efficient fine-tuning method. Instead, our method creates a small parameter overhead to the base fine-tuning method and serves as guidance to LM's reasoning process.Qin & Eisner (2021) uses a mixture of soft tokens to model a relation, which tries to specialize within the relation extraction task. However, we do not explicitly compute the mixture probability across all token types as in Qin & Eisner (2021) , which can be untractable for CoT generation. Instead, at each step, we let the LM choose one planning token type.Another line of work prepend newly added tokens as memory to transformers Burtsev et al. (2020); Bulatov et al. (2022) ; Darcet et al. (2023) , which echos our understanding that increasing sequence length can increase the capacity of the Transformer. Goyal et al. (2023) proposes to append a new pause token after the prompt so that the LM can have more computation budget before starting the actual generation. This shares the same idea with the token as a memory line of work. On the contrary, Mu et al. (2023) proposes to append a new gist token after the prompt so that the following generation will only attend to this gist token for compression purposes. Our General baseline can be viewed as an enhanced version of Goyal et al. (2023) . Instead of only adding new tokens after the prompt, we also add new tokens before each CoT step, which further increases the computation capacity.To the best of our knowledge, no existing work proposes to add new tokens that can be generated by the LMs at inference time. The learned specialization of planning tokens within one task, instead of assigning a token for each specific task, is also novel.Math Reasoning Recently, LM-based math reasoning models have shown to be highly effective (Azerbayev et al., 2023; Yang et al., 2023) . Recent studies on complex math reasoning problems usually adopt a CoT-based approach (Zhang et al., 2023; Li et al., 2023) that finetunes/prompts LLMs to generate reasoning steps before giving the final answer. Most fine-tuning works are highly dependent on the augmented CoT solutions generated by a strong pre-trained LM (Yue et al., 2023; Yuan et al., 2023; Mukherjee et al., 2023; Luo et al., 2023) . Our planning token method can fundamentally improve the fine-tuning performance regardless of the present of augmented CoT data. A possible extension of our work is to combine our method with the CoT data augmentation method to boost the performance even more. Our method is especially related to Zhang et al. (2023) . They perform CoT fine-tuning of GPT2 by first predicting the math operation of each reasoning step at generation time, which is less efficient than our end-to-end method. 
conclusion: 
 We proposed using planning tokens to gain better control over reasoning with LLMs. Our planning tokens increase the performance over baselines at a minimal parameter cost. We studied our method's efficacy across three different base models and multiple datasets. First, we see a straightforward extension of our framework which could use multiple planning variables at every reasoning step, each planning variable is responsible for deliberating in a different "view" of that particular reasoning step (Wang et al., 2023) . Second, future work should go beyond our heuristic inference procedures and learn the inference network, such as to maximize the marginal log-likelihood of the observed data: we could then interpret the overall model as a Sequential VAE (Goyal et al., 2017) . Finally, it is meaningful to continue the exploration towards interpretability and explainability of the planning tokens (Khashabi et al., 2021) . We believe such research could shed light on prompt searching/optimization studies performed both by humans (Zamfirescu-Pereira et al., 2023) and machines (Shin et al., 2020; Sordoni et al., 2023) . 
ethics statement: 
 We do not foresee a significant societal impact resulting from our proposed method. In this work, we propose to equip LLMs with planning tokens to obtain better reasoning capabilities. We find that the planning tokens discovered by embedding-based methods (less interpretable, e.g., K-Means and SQ-VAE) consistently outperform those from heuristicsbased methods (more interpretable, e.g., Arithmetic). Therefore, an obvious direction for future work is to improve the interpretability of such methods. In addition, while we aim to improve an LLM's reasoning ability, it is possible that the training of the planning tokens could be affected by misinformation and hallucinations introduced by the base LLM. Therefore, we believe that caution needs to be exercised if one extends our work to a setting where humans are directly involved, e.g., in an educational setting. 
 
 Appendix: Question: Every day, Wendi feeds each of her chickens three cups of mixed chicken feed, containing seeds, mealworms and vegetables to help keep them healthy. She gives the chickens their feed in three separate meals. In the morning, she gives her flock of chickens 15 cups of feed. In the afternoon, she gives her chickens another 25 cups of feed. Due to excessive computational requirements, in the experiments using Llama2 as the base model, we rely on low-rank adapters (LoRAs) (Hu et al., 2021) to fine-tune the base LLM. We apply LoRAs to the projection parameters of each MLP block in the base LLM. For all methods, we use a rank of 16, and a dropout of 0.05. To further reduce the memory usage, we load 13B Llama2 in 8 bits at training time. For Phi-1.5, we perform full-model fine-tuning.For our variants with planning tokens, we add the embeddings of the planning tokens to the learnable parameters of the base model, which increases less than 0.001% of the total trainable parameters. We train all models for 10 epochs, with a learning rate of 2e-5 using the AdaFactor optimizer (Shazeer & Stern, 2018) for full fine-tuning, and a learning rate of 2e-4 using the AdamW optimizer (Loshchilov & Hutter, 2017) for parameter efficient tuning.In Table 4 , We show that the prompting (Lester et al., 2021) /prefix (Li & Liang, 2021) tuningbased fine-tuning methods are not suitable for complex reasoning problems, compared to other parameter-efficient fine-tuning methods like LORA. The only way to increase the capacity of these methods is to add more tunable tokens to the prompt while a large number of new tokens will significantly increase the inference time and space complexity. Adding a comparable amount of new tokens to Table 1 yields significantly lower performance than our method combined with LORA.In Figure 7 , we show the frequency of each planning token that appears in the GSM8K test set. Annotation means the planning type predicted by the SQ-VAE. Generation means the planning token generated by the fine-tuned language model. As we can see, the marginal distribution of the planning variable approximately matches between SQ-VAE and the language model, which means fine-tuning the language model with planning tokens can effectively learn to infer the correct planning type from the previous steps. We find the attention head with the largest difference between planning tokens and normal tokens as shown in Figure 8 as the attention head that strongest corresponds to the planning token mechanism. Here is the prompt we used for error type classification:We open source the code at https://github.com/WANGXinyiLinda/planning tokens.We use mpnet-baseSong et al. (2020), a default model in Sentence-TransformersReimers & Gurevych (2019). 


Decision: This paper proposes an interesting and intuitive method to improve language models' multi-step reasoning ability by introducing "planning tokens" before each reasoning step. While the experiments focus only on math reasoning, which limits generalizability, the promising results and potential for extensions to other reasoning tasks make this a solid contribution. Please make clarifications as recommended by the reviewers and perform any additional analysis of clustering, which would further strengthen the paper. Overall, the reviewers recommend acceptance given the interesting idea, solid experiments, and clear presentation.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: 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?
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: 1. How does it construct clusters? Are the clusters also used in the test dataset?

2. Could you show more analysis for clusters?
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

Comments: comment: Sounds good, looking forward to that!
Thanks for your reply! We had a less clear intuition about how planning tokens can help other reasoning types. Maybe for multi-hop QA one can infer a planning variable for each hop. We are working on adding a new dataset, StrategyQA [1], which has a training set with multi-hop questions decomposed into single-hop questions, to see the effect of our proposed method. As math world problems are most aligned with the overall direction we want to pursue in our research, our current paper focuses on them. But it's certainly interesting to see if extensions on other types of reasoning tasks are possible. We will keep you posted on the progress of the new dataset if we can finish before the discussion period.

[1] Geva, M., Khashabi, D., Segal, E., Khot, T., Roth, D., & Berant, J. (2021). Did aristotle use a laptop? a question answering benchmark with implicit reasoning strategies. Transactions of the Association for Computational Linguistics, 9, 346-361.
Thank you for the response! Is there anything you'd like to say regarding the first point in "reasons to reject"?