Guiding Language Model Reasoning with Planning Tokens
Authors: Xinyi Wang
Lucas Caccia
Oleksiy Ostapenko
Xingdi Yuan
William Yang Wang
Alessandro Sordoni
Venue: COLM (2024)
Keywords: chain-of-thought reasoning, fine-tuning
TL;DR: We propose to add specialized planning tokens in front of each chain-of-thought step to guide and improve language models' math reasoning ability.

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.

Published as a conference paper at COLM 2024
Guiding Language Model Reasoning with Planning Tokens
Xinyi Wang'; Lucas
Caccia?, Lucas William Yang Wang', Alessandro Sordoni 2.3,4
Oleksiy Ostapenko 2,34, Xingdi Yuan?,
University of California, Santa Barbara Microsoft Research
Mila Quebec AI Institute 'Université de Montréal Sinyi_wang@ucsb.co williamêcs, wwsb. edu, yea,c.
Abstract
Large language models (LLMs) have recently attracted considerable inter- est 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 gener- ation 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 more parameter-efficien 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!
1 Introduction
The great potential of solving complex reasoning problems, including world knowledge reasoning (Hendrycks 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-of- thought (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 given problem. Recent research has recognized the importance of Co'Is 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-thougg. (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- (Mukherjee 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 cata-augmmentation-aasss 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 lokens the beginning of each CoT step, which are special tokens that encode 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
"This work was done during an internship at Microsoft Research, Montreal.
We open source the code at https:/ //s: thub. com//ANNXinyil
tokens.

---

Published as a conference paper at COLM 2024
0. Training Data
Planning Token Inference
Latent Inference
2. Training
3. Generation
y
11
Languagi Model
Language Model
An overview of our method: Given an input question x and its CoT steps r! ,.., ri, we train a planning token inference model to infer the planning tokens f'. Subsequently, fine-tune an LM on data augmented with the inferred planning tokens. Given new question x', the trained LM can infer planning tokens 1", CoT steps r", and final answer y'. Fire and snowflake denote trainable and frozen models, respectively. Note that x, r", t can all include multiple tokens, we depict them as single blocks/ blocks/variables variables for simplicity.
Figure 1:
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 parfieter-efficceet 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 -arfiiter-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 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.
2 Method
Setup We assume that we have a dataset D composed of triples {(xr,, where x is a problem to solve, is the ground-truth reasoning associated with the problem and y is the
2

---

Published as a conference paper at COLM 2024
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 inpuls, i.e.:
arg max
Σ [logp(r,y)x;0), D
where we assume that x,y, are properly tokenized and 0 is the trainable parameter of the pre-trained LLM. To provide more high-level guidance the LLM at fine-tuning time with reasoning data, we assume that each reasoning r is the realization of some higher- level planning process, which is latent, thus unobserved in the dataset. To model this assumption, we define by reasoning step 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.
Letr = :ttr=r),,.., r,,,,,,,,, where S is the number of steps in the reasoning. We assume that each reasoning step r² is generated conditional on the previous reasoning steps and discrete latent planning variable 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 special planning variable & which has always value equal to P for each example. Without loss of generality, we assume that only f can take value P, i.e. p(t' = P) = 0 for each +
If the planning variables were indeed observed, the complete data D would be composed of tuples (x, rt,,,,,,..,, ,r.,,,,,! S,,,,,, ,,, 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 r' with a (or possibly multiple) planning token(s). Here we abuse the notation t for this (set of) planning token(s) for simplicity. In practice, this means that 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 Dc can be readily used to train an LLM by maximizing the complete data ligelikel:hood:
arg max &
Σ D
Elogp(t', ,r,,..,, ,5, 5 y x)x;,
by 8+ 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.
2.1 Planning Tokens Inference
We consider three different ways of inferring the planning variable from the given CoT train- ing 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, is natural to consider the basic arithmetic operation contained in each reasoning step ri the plan token I' similar to Zhang et al. (2023); Qin & Fisner (2021). This yields four planning tokens (+,-, x, x)+), which can be easily extracted from each r 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 x). Thus, this arithmetic
3

---

Published as a conference paper at COLM 2024
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, turn to inference planning tokens from the embeddings of CoT steps.
We first gather reasoning steps r' for all examples in our dataset, then apply an inference algorithm on the representations ((r'), Vr E !D We obtain ((r') 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'. In this way, we aggregate both the question and previous steps' information in the contextualized representation & (rr).
K-Means We run an off-the-shelf K-Means implementations the set of representations ((r') for all reasoning steps in the dataset D. Each planning variable I is then assigned the index of the centroid closest to r".
SQ-VAE In this approach, our aim is to learn non-linear transformation of each represen- tation p(r') that better captures the meaningful directions of variation of the representations for the dataset under consideration. Variational Autoencoders (VAEs) (Kingma & Welling, 2014) offer probabilistic approach to learning such non-linear latent representations of data. In their original formulations, VAEs learn a Gaussian-distributes representations. To induce discrete structure in the latent space, we follow Miao et al. (2017), and use "Gaausian-softmax" parameterization, which soft-quantizes the latent representations before reconstructing the input data. In particular, our VAE maximizes the following objective:
latent space of
L(enc, dec, 4) E Fene(z!le(r) [lleec (o ( [[2] (e
= 44, 4 softmax(z')
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 Daussian-distributee representation z', which corresponds to the distribution over clusters. & is a learnable matrix with P - 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 fl arg max softmax(ened ((r))) where enc((((') denotes the mean of the Gaussian distribution parameterized by enc.
2.2 Planning Tokens Parametrization
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 parameersss our approach: n prerffx and n_special. Each planning variable can be verbalized using variable number of both generic prefix planning tokens and special planning tokens. planning token annotated example with n-prefix n_special = is shown below:
Question: Chenny 10 years Alyana years younger than Chenny. How
she is years older than Alyana?
sprefix, 0><prefix__ 1><prefix 2><kmeans1 0><kmeens1 1><kmeanst 2> Alyena is 10 << <<10-4=>>>6 years old.
<preffx_0=<srrefix [psprefix_1><prefft 1=<preffx_2=<kmeanss wtw.222-mmmnnnn__-emeeeeeee 1attmmeans__11_kkeeee 1><kmeans3.2 2> So, Anne is 6 2 <<c- 2-8 years old.
Dorefix_D<<prefiie tiispprefix_sspreftie 2=<prefix_2x<ensseerr wtw.:annweeeexeaannseeeeeee
answer
4

---

Published as a conference paper at COLM 2024
Model
Method
#clusters
#trainable
GSMSK
MATH
AQUA
S-QA
Avg
Phi 1.5 (1.3B)
Full-FT
0
100%
12.5
1.3
27.2
13.5
General
100%
15.4
2.0
35.4
17.6
Arithmetic
4
100%
15.0
2.3
33.1
16.8
+ K-Means
5
100%
14.5
2.7
36.5
17.7
+ SQ-VAE
100% 5
15.8
3.3
34.3
17.8
Llama2 (7B)
LoRA
0
0.343%
38.2
6.5
36.6
58.4
34.9
Pause
1
0.344%
37.2
6.7
36.2
58.0
34.5
General
1
0.344%
38.5
6.7
37.8
59.1
35.5
Arithmetic
4
0.344%
39.5
38.2
+ K-Means
5
0.344%
39.1
6.7
40.5
62.2
37.1
+ SQ-VAE
0.344% 5
40.0
7.0
41.3
62.8
37.8
Llama2
LoRA
0
0.279%
44.6
7.2
41.3
69.5
40.7
(13B)
General
1
0.280%
47.9
7.9
42.5
70.3
42.2
Arithmetic
4
0.280%
41.9
4.6
35.8
K-Means
0.280%
49.6
44.1
71.5
43.4
+ SQ-VAE
5
0.280%
50.6
8.5
43.9
72.4
43.9
Table 1: Testing accuracy of fine-tuned language models on different math word datasets.
We set the number of planning tokens (n-preefxttn-specid)) to 6.
Here <kmeans1 0>, <kmeans1 1>, and <kmeans1 2> are the three planning tokens induced by the first K-Means cluster.
3 Experiments
Datasets We conduct experiments on four datasets. The Grade School Math dataset (GSM8K) (Cobbe et al., 2021) contains 8.5K examples of linguistically diverse grade school math world problems. The MATH dataset (Hendrycks et al., 2021a) is a collection of 12.5K challenging competition mathematics problems formatted in latex notation. The AQUA- RAT dataset (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 used the Chain-of-thought (CoT) path of the question.
Base LLMs Our empirical analysis uses several decoder-only architectures 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), 1.3B parameter model trained on mixture of textbook-quality code data, and additional synthetically generated textbook and exercise data.
Baselines First, we compare against the vanilla baselines of full-fnne-tning (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 this baseline, we prepend the same prefix tokens to each CoT step. i.e. P 1
Dall-FT/LLRRA. For
Note that the General baseline can be viewed as prompt-tuning (Lester 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 not suitable for complex CoT reasoning tasks our setting, as their capacity 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.

---

Published as a conference paper at COLM 2024
Plan Type
Number of Plan Tokens
Plan Type
Number of Clusters P
2
6
10
3
5
7
10
General
37.9
38.5
38.3
General
38.5
Arithmetic
39.5
38.9
38.5
K-Means
37.9
39.1
39.9
36.7
K-Means
38.9
39.1
38.1
SQ-VAE
39.6
40.0
39.4
38.9
SQ-VAE
40.0
38.8
38.2
Table 2: Impact of varying the number of clusters (left) and varying the number of plan- ning tokens (right) (n prefix+n special). We set prefix = n special. For general, we match computation by adding prefix tokens.
3.1 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/ Lool-FTLLoRA may due to additional allowed planning type inference
computation space. However, the other two embedding 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 the baselines of up to 1.7 points, a relative increase of 6%. Our results over the Llama2 13B tell 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 more reliable approach, as is it always the best- performing 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 larger performance improvement.
Ablation 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 (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 longer sequence which negatively affects the language model's 's reasoning ability.
3.2 Analysis
Errors by Reasoning Length In Figure 2, we show the accuracy by the complexity of the test example measured by reasoning length 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-thouu thoughts solutions from all three datasets, we propose classify the errors into the following five categories:
Misunderstanding of question: generated solution does not answer the given question.

---

Published as a conference paper at COLM 2024
Aqua
Pianing type
Planing type
5.6
0.6
N/A
Arithmetic SQ-VAE
0.4
Arithmetic SQ-VAE
NIA
6.5
Accuracy
-0.23
Accuracy
0.3
0.2
0.2
0.12
0.1
0.1
0.0
4
6
7
9
10 11 12
0-3
4-8
9-20
Number reasoning steps (ground truth)
Number reasoning steps (ground truth)
Figure 2: Accuracy on GSM8K (left) and Aqua (right) on test examples by their number of ground-truth reasoning steps. SQ-VAE consistently increases performance for test examples that require more steps reasoning to be solved.
Error
Misunderstanding of question
Computation errors
Insccurate extraction of question information
Wrong application of knowledge
Wrong logic
type
Example
Question:
H..How many cups of does she to give chickenss the final meal of day...
...7 3301x
Question: He spends the hour driving at speed of 30mph...
The number feet plane ground is the target of geometric sequence...
...$2ab m 12 m 2"3
3371x...
...then each chicken gets 40/20
30mph he traveled 2"30=<x2-30-00000
=xx40/20-2x32 cups feed per chicken. The answer 2s:
miles...
Dataset
SQ-VAE
SQ-VAE
SQ-VAE
SQ-VAE
SQ-VAE
GSMBK
4.15
4.1%
8.0%
6.1%
34.5%
33.0%
0.2%
0.5%
16.2%
14.9%
MATH
21.0%
23.5%
8.1%
5.0%
27.3%
24.7%
20.3%
18.2%
23.1%
23.0%
AQUA
8.3%
8.7%
12.0%
9.4%
25.8%
23.2%
2.8%
2.0%
19.5%
20.1%
Figure Examples of different error types generated by 7B Llama 2. The frequency of each error type generated by the LoRA baseline and our LoRA + SQ-VAE planning method on test sets are predicted by GPT4 and then manually verified.
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 the generated solution 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-VAÉ error frequencies significantly lower than the N/A baseline are in bold font. As we 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 related information. Our method being able to make more faithful reference to the information in the questions indicates planning tokens are helpful for long- range 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 is from American Mathematics Competitions (AMC). These questions are less obvious to approach and require applications non-trivial math theorems.
Attention on planning tokens 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 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.
7

---

Published as a conference paper at COLM 2024
General
Arithmetic
K-means
SQ-VAE
Normal token
0.0045
0.0039
0.0039
0.0043
Planning token
0.0009
0.0012
0.0071
0.0070
Table 3: Averaged attention on planning tokens v.s. normal tokens across all layers, heads, and previous tokens on the test set of GSM8K, with Llama 2 (7B) model.
0.15
0.10
0.35
Question: Chenny is 10 years old. Alyana is years younger than Chenny. How old is Anne she years older than Alyans? 28 Alyana
<<10-4-6008 years <prefix_0>> 0+ sprefix_12 <prefix_2 2> spretix 2> <vee4_ <vae4_ <vae4_22 2> So, Anne is 6 2--0- years old.
sprefix d> <prefix_ <prefix sprefix eved <vaef <vae1.
<pretix 0= sprefix 2> sprefix 2> spretix 2> canswer 0a <answer 1> <answer 2 The answer 8
Figure 4: Two attention heads activated by planning tokens, with a test example generated by Llama (7B) LoRA fine-tuned with SQ-VAE planning tokens. Question tokens are omitted in the attention heat map but included in attention computation. A lighter color means a larger attention weight. Best viewed zoomed in.
This implies that the inferred planning tokens might 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 Transformer works. Similar to Olsson et al. (2022), we identify attention heads that have strong patterns corresponding to the planning tokens shown in Figure and deduct how language models make use of the planning tokens from the patterns.
valid way of understanding how the
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 strongly correspond to the utilization of planning tokens. We visualize two such attention heads on CoT generated with SQ-VAE planning tokens in Figure 4. The attention head shown on the left strongly attend 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 Col step, instead of carrying information to the whole sequence.
Distinguishability of the Induced Clusterings We investigate whether SQ-VAE learns better planning types than K-Means via a probing task (Alain & Bengio, 2017). Concretely, collect dataset consisting of CoT steps and their planning token labels obtained from K-Means and SQ-VAE. We train a simple model (either logistic regression or shallow neural network) to learn the mapping from sentence to its planning token label. We
8

---

Published as a conference paper at COLM 2024
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 case study, we study the 5-cluster setting the GSM8K dataset. Given a planning type, we ran- domly sample 2000 sentences from each of the five categories. We leverage a pre-trained text encoder to convert sentences into vector representations. Taking the text encodings as input, the probe network (logis- tic regression, 2-layer MLP, or 6-layer MLP) performs 5-way classification.
test accuracy
0.5
probe
logistic regression
2:ayer mip
6layer
0.3
parriigg typa
VA-
20
60
100
k training examples
We show test accuracy curves in Figure 5. We ob- serve 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 pro- duces better overall downstream task performance.
Figure 5: Testing accuracy of the probes on the sentence classification task.
4 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 parcmeter-efficient fine-tuning techniques to adapt LMs to specific task (Qin & Fisner, 2021; Li & Liang, 2021; Lester et al., 2021). Our planning tokens are not intended to serve as parameter-efficieet fine-tuning method. Instead, our method creates small parameter overhead to the base fine-tuning method and serves as guidance to LM's reasoning process. Qin & Eisner (2021) uses mixture of soft tokens to model relation, which tries to specialize within the relation extraction task. However, we do not explicitly compute the mixture probability across all token types in Qin & Eisner (2021), which 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 2023). Recent studies on complex math reasoning problems usually adopt CoT-based approach (Zhang et al., 2023; Li et al., 2023) that fine- tunes/ /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 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. À possible extension of our work is to combine our method with the CoT data augmentation method boost the performance even more. Our method is especially related to Zhang et al. (2023). They perform CoT fine-
We use mpnet-base Song et al. (2020), default model in Sentence- Transformers Reimers & Gurevych (2019).

---

Published as a conference paper at COLM 2024
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.
5 Conclusion
We proposed using planning tokens to gain better control over reasoning with LLMs. Our planning tokens increase the performance over baselines minimal parameter cost. We studied our method's s efficacy across three different base models and multiple datasets. First, we see straightforward extension of our framework which could use multiple planning variables at every reasoning step, each planning variable is responsible for deliberating in 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 maximize the marginal log-likelihood of the observed data: we could then interpret the overall model 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/opoi /optimization studies performed both by humans Canfirescu-Pereing et al., 2023) and machines (Shin et al., 2020; Sordoni et al., 2023).
searchingooptiizztionn
Ethics Statement
We do not foresee 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 heuristics- based 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, 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 directly involved, e.g., in educational setting.
References
Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes. ArXiv, abs/1610.016644, 2017.
Zhangir Azerbayev, Hailey Schoelkopf, Keiran Paster, Marco Dos Santos, Stephen McAleer, Albert Q Jiang, Jia Deng, Stella Biderman, and Sean Welleck. Llemma: An open language model for mathematics. arXiv preprint arXio:2310.10631, 2023
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems, 35:11099-11091, 2022.
Mikhail Burtsev, Yuri Kuratov, Anton Peganov, and Grigory V Sapunov. Memory trans- former. Xiv preprint ar arXit:2006.11527 2020.
Ta-Chung Chi, Han Fan, Alexander I Rudnicky, and Peter J Ramadge. Transformer working memory enables regular language reasoning and natural language length ex- trapolation. ar Xio preprint ar Xit:2305.03376 2023.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichrr Nakano, et al. Training verifiers to solve math word problems. Xix preprint ar Xit:21100.1463. 2021.
10

---

Published as a conference paper at COLM 2024
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. arviv preprint arXio:2309.16588, 2023.
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: A theoretical perspective. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreviiie net/ forum?3dd-qHrADgAdlY
Mor Geva, Daniel Khashabi, Elad Segal, Tushar Khot, Dan Roth, and Jonathan Berant. Did aristotle use a laptop? a question answering benchmark with implicit reasoning strategies. Transactions of the Association for Computational Linguistics, 9:346-361, 2021.
Anirudh Goyal, Alessandro Sordoni, Marc-Alexande Côté, Nan Rosemary Ke, and Yoshua Bengio. Z-forcing: Training stochastic recurrent networks, 2017.
Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. Xrv preprint arXiv:2310.02226 2023.
Suriya Gunasekar, Yi Zhang, Jyoti Aneja, Caio César Teodoro Mendes, Allie Del Giorno, Sivakanth Gopi, Mojan Javaheripi, Piero Kauffmann, Gustavo de Rosa, Olli Saarikivi, et al. Textbooks are all you need. Xiv preprint wxXi:22366.16444 2023.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding, In International Conference on Learning Representations, 2020.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the MATH dataset. CoRR, abs abs/2103.03874 2021a. URL
wwwpaawwwww...raa////aa///// 03874.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arxiv preprint ar #rXiv:2103.03874 2021b.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. 2021. URL
Daniel Khashabi, Shane Lyu, Sewon Min, Lianhui Qin, Kyle Richardson, Sean Welleck, Hannaneh Hajishirzi, Tushar Khot, Ashish Sabharwal, Sameer Singh, et al. Prompt waywardness: The curious case of discretized interpretation of continuous prompts. arxiv preprint ar 1rXiv:2112.08348, 2021.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.6114 2013.
arXio preprint
Diederik P. Kingma and Max Welling, Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for efrfmeter-efficcent prompt tuning. In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen- tau Yih (eds.), Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 3045-3059, Ónline and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10. 10.18653/ 21/6633/v12222emmm/ www.:222mmmmmmmmmmmmm.... URL org 2021. wwnee-nnnaaleeeeee mamp.mmain...m 243.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli (eds.), Proceedings of the 59th Anmual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing (Volume 1: Long Papers), pp. 4582-4597 Online, August 2021. Association for Computational Linguistics. doï: 353. URL https: hwtp///aaaaaaaaoaaaaoaaoao org/ www.lllll......a..aaaaaaaaaallla 353.
11

---

Published as a conference paper at COLM 2024
Yifei Li, Zeqi Lin, Shizhuo Zhang, Qiang Fu, Bei Chen, Jian-Guang Lou, and Weizhu Chen. Making language models better reasoners with step-aware verifier. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. pp. 5315-5333, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.18653/ /120023.cll-oogg.. 291. URL https: http:///aaaaaahoooooooo org/2223. acl acl-long 291.
Wang Ling, Dani Yogatama, Chris Dyer, and Phil Blunsom. Program induction by rationale generation: Learning to solve and explain algebraic word problems. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Pp. 158- 158-167, Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 10.18653/ 11/117-1015 URL https: /atess://lllattt wcc.aanoooog..org//11----
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arxiv preprint axXit:1711005001, 2017.
Haipeng Luo, Qingfeng Sun, Can Xu, Pu Zhao, Jianguang Lou, Chongyang Tao, Xiubo Geng, Qingwei Lin, Shifeng Chen, and Dongmei Zhang. Wizardmath: Empowering mathematical reasoning for large language models via reinforced evol-instruct. arv preprint Xrit:2308.00555 2023.
Yishu Miao, Edward Grefenstette, and Phil Blunsom. Discovering discrete latent topics with neural variational inference. In International conference on machine learning, pp. 2410-2419. PMLR, 2017.
Jesse Mu, Xiang Lisa Li, and Noah Goodman. Learning to compress prompts with gist tokens. arXiv preprint arXiv:2304.08467 2023.
Subhabrata Mukherjee, Arindam Mitra, Ganesh Jawahar, Sahaj Agarwal, Hamid Palangi, and Ahmed Awadallah. Orca: Progressive learning from complex explanation traces of gpt-4, 2023.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint ar Xxfv:2209.11895 2022.
Liangming Pan, Alon Albalak, Xinyi Wang, and William Wang. Logic-LM: Empowering large language models with symbolic solvers for faithful logical reasoning. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 3824, Singapore, December 2023. Association for Computational Linguistics. doi: 10.18653/ ht///v1/211/ 2212223..fndee 19223findings-emn// 248. URL https: //aclanthology. org/2023 findings-emnle 2.nnl..28..... 248.
Guanghui Qin and Jason Eisner. Learning how to ask: Querying LMs with mixtures of soft prompts. In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani- Tur, Iz Beltagy, Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou (eds.), Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 5203 5212, Online, June 2021. Association for Computational Linguistics. doi: (0188653/1//// www.:222...aaalaecmeaaaaaaa hww.a-mainnneeee0 URL https: /ttps://aclo https://aclanttoology.oo wwwpthhologg.orr/2eeee 2021 naacl-main 410.
Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bert- networks. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 11 2019. URL https: /ttpss:///xvvvoor// abs/1908. 10084.
Noam Shazeer and Mitchell Stern. Adafactor:
Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596-4604. PMLR, 2018
Taylor Shin, Yasaman Razeghi, Robert L Logan IV, Eric Wallace, and Sameer Singh. Au- toprompt: Eliciting knowledge from language models with automatically generated prompts. arXiv preprint ar 209:20000.15980, 2020.
12

---

Published as a conference paper at COLM 2024
Kaitao Song, Xu Tan, Tao Qin, Jianfeng Lu, and Tie Liu. Mpnet: Masked and permuted pre-training for language understanding. Advances in Neural Information Processing Systems, 33:16857-16867, 2020.
Alessandro Sordoni, Xingdi Yuan, Mare-Alexandre Côté, Matheus Pereira, Adam Trischler, Ziang Xiao, Arian Hosseini, Friederike Niedtner, and Nicolas Le Roux. Deep language networks: Joint prompt training of stacked lims using variational inference. arXio preprint arXio:2306.12500) 2023.
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, et al. Challenging big- bench tasks and whether chain-of-thought can solve them. arviv preprint arXiv:2210009221, 2022.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint ar arXio:2302.13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arviv preprint arXiv:2307 009088, 2023b.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consistency improves chain of thought reasoning in language models. ar Xiu preprint ar Xrii:2203.111777 2022.
Zhenhailong Wang, Shaoguang Mao, Wenshan Wu, Tao Ge, Furu Wei, and Heng Ji. Un- leashing cognitive synergy in large language models: A task-solving agent through multi-persona self-collaboration arviv preprint ar htt::2307.05300 2023.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Chain-of-thoughh prompting elicits reasoning in large language models. Adoances in Neural Information Processing Systems, 35:24824-24833 -48224, 2022.
Kaiyu Yang, Aidan M Swope, Alex Gu, Rahul Chalamala, Peiyang Song, Shixing Yu, Saad Godil, Ryan Prenger, and Anima Anandkumar. Leandojo: Theorem proving with retrieval- augmented language models. Xiv preprint ar Xrx::2336..666 2023.
Zheng Yuan, Hongyi Yuan, Chengpeng Li, Guanting Dong, Chuanqi Tan, and Chang Zhou. Scaling relationship learning mathematical reasoning with large language models. 8 2023. URL http: httt::://axxv..oo/ orw. hts//bbs//2300 030./0088.........
Xiang Yue, Xingwei Qu, Ge Zhang, Yao Fu, Wenhao Huang, Huan Sun, Yu Su, and Wenhu Chen. Mammoth: Building math generalist models through hybrid instruction tuning. ar Xiv preprint arXi::3009.0553 2023.
JD Zamffrescu-Pereiiaa, Richmond Y Wong, Bjoemn Hartmann, and Qian Yang. Why johnny can't prompt: how non-ai experts try (and fail) to design Ilm prompts. In Proceedings of the 2023 CHI Conference on Human Factors in Computing Systems, pp. 1-21, 2023.
Mengxue Zhang, Zichao Wang, Zhichao Yang, Weiqi Feng, and Andrew Lan. Interpretable math word problem solution generation via step-by-step planning. 6 2023. URL http: //arxiv.org/ablll wtt.:oor//aa//./..00/0// 00784.
13

---

Published as a conference paper at COLM 2024
A Appendix
generation example
Question:
Every day, Wendi each her chickens cups mixed chicken feed, containing
seeds, mealworms and vegetables to help keop 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 afteroon,, she gives her chickens another cups feed. How many cups feed does she to give chickens in the final meal the day the size of Wendiss ilockis 20 chickens?
Vanilla supervised fine-tuning
Fine-tuning with arithmetic planning tokens
Every day, Wendi gives her chickens 15 cups feed in the morning + cups feed afternoon <<15-25-40=>400
sprefic> <e> Wendi gives her flock 15 cups of feed in the morning and another 25 cups
aftemoon, total of 15+25 <<15++5=40=><40 cups offe..
cups of feed.
sprefiio <t> Wendi has 20 chickens, then she needs 20º3 200000000000000 cups feed to feed flock.
feeds them 40 cups of feed every day, then each chicken gats 40/20
If chickens
sprefix> <.> Wendi has elready given her flock 40 cups feed, then needs to give
<x40/20-2>a2 cups feed per chicken.
60-40 <<80-
Drifting away
40-20>0-20 more cups of feed.
The answer is:
from the
reasoning flow
sprefoc <anewer> The enswer is:
Figure 6: An example of question from the GSM8K dataset, along with two chain-of- thoughts solutions generated by LLMs. Left: The reasoning chain generated by vanilla supervised fine-tuning LLM drifts away from correct reasoning flow and leads to wrong answer. Right: we use planning tokens to guide an LLM's generation at every reasoning step to encourage the correct reasoning flow. For simplicity, we show the variant of our method which utilizes arithmetic planning tokens: <prefix>, <+>, <t>, <->, <answer> are the planning tokens added to the original vocabulary.
Training details
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.
Comparison with prompt tuning
In Table 4, We show that the prompting (Lester al., 2021) /prefix (Li & Liang, 2021) tuning- based fine-tuning methods are not suitable for complex reasoning problems, compared to other parameter-efficien 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.
Planning token distribution
In Figure 7, we show the frequency each planning token that appears in the GSM8K test set. Annotation means the planning type predicted by the SQ-VAÉ. 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- min-tunnng the language model with planning tokens can effectively learn to infer the correct planning type from the previous steps.
14

---

Published as a conference paper at COLM 2024
LM
Method
GSM8K
MATH
AQUA
Avg
Llama2 (7B)
Prefix
8.9
3.6
32.9
15.1
Prompt
15.2
5.3
26.8
15.8
LORA
38.2
6.5
36.6
27.1
Llama2
Prefix
16.0
3.2
30.3
16.5
(13B)
Prompt
27.8
6.4
26.0
20.1
LORA
44.6
7.2
41.3
31.0
Table 4: Comparison between LORA and soft prompt tuning based method. We similar number of tokens for Prefix tuning and Prompt tuning as our methods in Table 1 to keep the inference time complexity the same.
Annotation
13001382
765 805
577
240
<yc_1>>
<vac_3>
Figure 7: The marginal distribution (count) of the SQ-VAE planning tokens over the GSM8K test set.
Planning token attention head
We find the attention head with the largest difference between planning tokens and normal tokens as shown in Figure 8 the attention head that strongest corresponds the planning token mechanism.
Head corresponding to planning tokens
13342
20000
20000
Planning tokens
Normal tokens
Figure Heat maps of averaged attention on all attention heads of Llama (7B) fine-tuned with SQ-VAE planning tokens on GSM8K.
GPT4 prompting details for error type classification
Here is the prompt used for error type classification:
15

---

Published as a conference paper at COLM 2024
Error Classification Prompts
Your goal is to identify the error type in reasoning steps of an Al agent in response
to question.
You are given the corresponding question, the reasoning steps generated by an Al model
the ground truth reasoning steps in the format:
#urQuesni::: [question]
####Geerrteed Answer: [reasoning]
#WIHGT reasoning: [ground truth reasoning]
Your task is to identify the error type in Al's reasoning There are 5 different error
types: (1) Misunderstanding of question (2) Computation errors. (3) Inaccurate extraction of question information, (4) Wrong application of math knowledge. (5)
Wrong logic/ /urueasoaabbb step
Here are some examples of each of the error types:
1 M Misunderstanding of question: ask about total number of cups in the final meal, answer
number cups each chicken get.: Question: Every day, Wendi each of her chickens three
of mixed chicken feed containing seeds. 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 feed
mealworms
chickens in the day if the size of Wendi's 's flock is 20 chickens? Every day. Wendi gives her chickens cups
How many cups
need to give
final
of
afternoon <<15=25=40>>-00 cups of feed. If she has 20 chickens and she feeds them 40 cups of feed every day, then each chicken gets 40/20 <<<0/20=2>>2 cups of feed per chicken. The
morning + cups feed
Computation errors:
3371x: Question: What
smallest positive
integer SNS such that the value $7 + (30 (times N)5 is not prime number? We have + (30 (times N) = 3371xS We write $3371x$ as $3300+7lx5. in order
3000 + \dots + (30) (times = 7 3301x
factor We can try adding one to $715 and subtracting it from $3300$ we add 2. we get $735, but 5735
multiple of 33
prime If add we get $835. but $835 is prime If we add 16, we get $995. but $993 is prime If get $1955,
get $75$ but $75$ prime
add
prime is prime If we add 128. we get $323$ but $3235 is prime If we add 256, we get 55815 but 5581S is prime $11275, but $11275 is prime If add 2048. we get $13355, but $15445 prime we add $13355 is prime If we add 4096, we get $1544$ add 16384, get $24015, but
32, we get $1315,
we
$1955
we add 512. we get $8035, but $8035
prime we
add 1024.
8192, get $17635, but $1763$
$2401$ is prime If we add 32768. we get $2777S, but $2777$ is prime Therefore. smallest integer
boxed (2777)
is:
16

---

Published as a conference paper at COLM 2024
Error Classification Prompts (Cont'd) d
Inaccurate extraction question information:
Question: John drives for 3 hours
of 60 mph and then around because he realizes he forgot something very important at home. He tries to get home in hours spends the first hours in standstill traffic. He spends the half-hour driving at speed of 30mph, before being able to drive the remaining mph.
from those 4 hours? He was driving at 60mph for 3 hours so he traveled 60+3=<<60+<==800
going
of
miles. He was driving for hours drave 4-2-<c4-2-2200
spent hours
standing still traveled >>30<<<<<30-60>>>>> 80mph
hours. He drove hours 30mph he miles <- should be half hour He drove 2 hours at miles
hour He drove 180 miles at 60mph, 60 miles at 30mph and 160 miles at 80mph for total of 100+660100<<<1100 100-180--0-10--1111 60-110-300>>3300 miles. <- probably because The answer is: 300
traveled <x80=<<2+00=60>> 10:000100>>>000
should
make sense the model change it to
Wrong application of math knowledge: it's
geometric sequence: Question: An airplane climbs 100 feet during the first second after takeoff. In each previous
succeeding second climbs second. How many seconds daes it take for the plane to reach an altitude of above ground is the target of geometric sequence with first term 0 and common ratio We are asked
more
during
12,000
height?
number
plane
12000 The first find the number of terms of this sequence which have a sum ratio
and the 2 more than the first the third is 2 more than the second and so on This means is "2"0) (0) (dot(22) (d)))c))c((() (4) (8) (8))cdot
that the second
that product of all cdot) (256) cdot (512) cdot [1024) cdot (2048) cdot (16) cdot (32) cdot (64) cdoot
(32768) cdot (65536) (131072)) (2211444)))))))) (262144) cdot (524288) ( (1048576)5, and the sum of all these terms cddtts + 1048576 = =2)ccot(220) + cdoots 2'(63) to the right hand side to get $2(1 + (cdots + 2"(44) = 2e (2'
(4096) cdot (8192)
(16384))ccor (1336)) ( cdot
$262144 262145
cdott |cdots (63)5
cedots + (2
left
equal (64) 15, so add $15
lrright) (65) - 15 Finally.
right hand side get (65) -1 == 2-2 - (65) = 1$, so it takes 5oooed (65) seconds for the plane to reach an altitude of 12,000 feet. The answer is:
subtract
Wrong logic/ unreasoable step: does not make sense to use G! -1: Question: how Sam insist sitting next to each other? (Two seatings are equivalent if one is rotation of the number
many ways
seat people at round table John
other .) arrangements in which John and Sam are next to each other is $41$ Therefore.. there are 561-41=\\ooed (75)$ different ways seat the group. The answer is: Question: If $2ab 125, evaluate 2bb"2. We have $2ab so $8a 2b" 2"3 (cdot (8a "2)b"2 2"3 /ddo 2b"2 (boxed ( The answer is $88 ~25
seating arrangements
number seating
2"3 (edot b*2s,
output follows: #P: [only index error type 1-5) #101 Explanation [short and precise explanation].
Here is example analyze:
#uestionon
#WWlGeeerttedd
#T reasoning:
Your response:
17


Reviews:
[{"review_id": "hOFsgJnylH", "reviewer": "Reviewer_652L", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive review. We would like to respond to the reviewer’s comments and questions below:\n\nResponse to reasons to reject:\n\n1. **Statistical significance**: Thanks for bringing this up. The fine-tuning runs are pretty resource and time-consuming for us so we did not have multiple run results for our main table. To make up for this, we conduct multiple analyses to support our results.\n\n2. **Reasoning step analysis**: We attribute the mixed performance to improved long reasoning capabilities primarily due to the introduction of specialized new tokens between each reasoning step. While both methods improve on long reasonings, SQ-VAE improves more on average. We will discuss more about this in the revision.\n\n3. **Attention interpretation**: Thanks for pointing us to the related work. While the raw attention weight might be a debatable way of understanding the token importance, it still serves as a valid way of understanding how the Transformer works (e.g. in recent work like the induction heads [1]). Similar to [1], we also 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. We will rewrite the statement that attention weight itself implies token importance.\n\nResponse to questions:\n\n1. **Abuse of notations**: We indeed abuse the notation $t^i$ for both the planning variable and planning token for simplicity, as stated in section 2.  Here $t^i$ can either represent one planning variable, or multiple planning tokens corresponding to one planning variable. In the last sentence of the caption of Figure 1, we explain that $t^i$ can be multiple tokens. We will rewrite the description of $t^i$ to make it clearer.\n\n2. **Typos**: Thank you for pointing this out. It should be n_prefix = n_special.\n\n3. **The confusing sentence in the Conclusion**: Here by “variant”, we mean we can assign different planning tokens for different steps, even if they are classified to the same planning variable, to differentiate the position. We will rewrite this sentence to make it clearer.\n\n4. **Model generating error taxonomy**: It’s Llama 2 7B. We will include this information in the revision.\n\n[1] Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., ... & Olah, C. (2022). In-context learning and induction heads. arXiv preprint arXiv:2209.11895.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The tackled problem (math reasoning) is important and timely.\n2. 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).\n3. The experimental design is sound, with reasonable baselines and various base LMs.\n4. The results show clear and consistent performance gains of the proposed method over baselines across all three datasets.\n5. The authors conduct a comprehensive and thoughtful set of analyses to reveal further insights into the strengths and weaknesses of the proposed method. \n6. 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.\n2. The performance gaps are sometimes small for certain settings in Table 1. It would be good to report statistical significance.\n3. 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.\n4. 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.\n2. Potential typo in Table 2: \"we set n_special = n_special\" -> should be \"n_prefix = n_special\"?\n3. 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?\n4. Which model's outputs were used in the error taxonomy analysis?", "extra_scores": {}}, {"review_id": "toD2ATND81", "reviewer": "Reviewer_y4dA", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive review. We are happy to see the reviewer finds our work interesting. We would like to respond to the reviewer’s comments and questions below:\n\nResponse to reasons to reject:\n\n1. Writing: We are happy to include any specific modifications that can improve clarity.\n\n2. Effect of special tokens: Our experiments are conducted to prove the effect of special tokens. In Table 1, we show the effect of different types of special tokens by performance on math reasoning datasets, with and without them. In Figure 2, we show the effect of special tokens on long reasoning. In Figure 3, we show the effect of special tokens on different error types. In Table 3 and Figure 4, we show the effect of special tokens on attention weights.\n\n3. Pause token baseline: Our General baseline is an upgraded version of the pause token, by essentially adding pause tokens at more positions. To better illustrate this, the accuracy of using the pause token with Llama 2 7B on GSM8K, MATH, and AQUA is 37.2, 6.7, and 36.2. For comparison, the general baseline is 38.5, 6.7, and 37.8, and our SQ-VAE method is 40.0, 7.0, and 41.3. We will add this to the main table for clarity.\n\nResponse to questions:\n\n1. How clusters are constructed: please refer to section 2.1, where we introduce three ways to infer planning token types and construct clusters, including a heuristic way, K-means clustering, and a VAE-base way. The clusters are constructed using training sets. At test time, the cluster/planning type is inferred by the language model.\n\n2. Cluster analysis: we show a probing-based analysis of clusters in the experiment section. It would be great if the reviewer could elaborate on the possible additional cluster analysis.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The proposed solution proposes a reasoning path by refining step units.\n\n2. 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.\n\n2. It is necessary to show the effect of special tokens in evaluation.\n\n3. It is necessary to compare the baseline solution using pause tokens. Refer to the following paper.\n\n[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?\n\n2. Could you show more analysis for clusters?", "extra_scores": {}}, {"review_id": "WMx5PSGIFW", "reviewer": "Reviewer_H5H2", "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.\n\nThe 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.\n\nThe 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.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive review. We would like to respond to the reviewer’s questions below:\n\nTypos: It should be n_prefix = n_special instead of n_special = n_special. Thanks for pointing this out.\n\nSoftmax function: Yes, there should be a projection matrix from the neural network hidden size to P bins. We will clarify this in the revision. Thanks for pointing this out.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The method is intuitively appealing.\n- The authors back it up with both consistent improvements in downstream performance numbers and interesting analysis.\n- The paper's cleanly written and very easy to follow.\n- Points to neat directions for further work.", "reasons_to_reject": "I don't see any issues that would merit rejection.\n\nThe 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:\nThe caption for Table 2 contains \"We set *n_special = n_special*\", presumably this is not intentional.\n\nSuggestions:\nIn §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.", "extra_scores": {}}, {"review_id": "VV5m8pZvEc", "reviewer": "Reviewer_mxHr", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive review. We would like to respond to the reviewer’s comments below:\n\nScalability: In this paper, we conduct experiments on small-scale data due to the computation limit. We are interested in scaling our method to the mixture of many different reasoning data and improving the general reasoning ability of language models in the future. Scaling up the number of clusters/planning types and applying our method to the mixture of all datasets might be a solution to this problem.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "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.", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept