Learning to Insert [ PAUSE ] Tokens for Better Reasoning (2025)

Abstract: To enhance reasoning capabilities, previous works have explored incorporating special-purpose tokens into the training process. These strategies strengthen the learning mechanism of transformer-based large language models (LLMs). Building on prior research, in which inserting dummy tokens consecutively just before reasoning steps can enhance effectiveness, we introduce a novel approach termed \textbf{D} ynamic \textbf{I} nserting Tokens \textbf{T} raining \textbf{(DIT)} . Our method identifies positions within sequences where model confidence is lowest according to token log-likelihood. Strategically inserting [PAUSE] tokens on these positions bolsters the model’s predictive capabilities for subsequent tokens. Experimental results across diverse datasets and models, from the 2.7B model to the 8B model, demonstrate that DIT consistently outperforms traditional fine-tuning and previous token insertion methods. With this simple yet effective method, we achieve accuracy gains of up to 4.7%p on GSM8K, 3.23%p on AQUA-RAT, and pass@1 improvements of up to 3.4%p on MBPP datasets. Our work shows a model-based, dynamic approach rather than a heuristic one, thereby broadening the scope of research in reasoning.

## introduction the primary mechanism in the text generation of recent large language models (llms) with the decoder-based transformer architecture @xcite involves autoregressive calculation. these models evaluate the probability of the best subsequent token, determined by the sequence of preceding tokens @xcite @xcite . this next token prediction [pause] token inserting methodology. we suggest that learning with [pause] token insertion using this method will result in a probability boost for low-probability tokens during inference. to effectively adapt llms for specific tasks, additional reasoning steps generally support adaptation, using approaches such as the chain of thought @xcite , which can be trained as a subsequent objective @xcite . while effective, training decoder-only llms with ntp loss based on previous input tokens does not always yield optimal re-sults @xcite . for instance, llms often fail on surprisingly trivial problems in mathematical reasoning tasks @xcite . to address such limitations, recent studies have explored mechanisms for adding dynamic computational budgets for transformer architectures. for instance, the universal transformer @xcite adapts model depth based on input complexity, while speculative decoding @xcite offloads demanding reasoning to larger expert models. recently, @xcite have shown that inserting dummy tokens during pre-training and decoding can act as a mechanism for introducing additional computational steps. building on prior studies @xcite @xcite , these dummy tokens strategically delay the output of decoder-only models, creating additional computational pathways in the attention layers and improving performance. however, this method relies on the stochastic inclusion of the dummy tokens during pre-training, followed by their fixed placement (depending on hand-crafted rules) during fine-tuning. rather than relying on heuristic methods, we propose a model-based approach for dummy token insertion, focusing solely on fine-tuning. our novel method identifies points with low token loglikelihoods, interpreted as regions of model uncertainty, and subsequently inserts dummy tokens (denoted as @xcite ) immediately preceding these positions. this approach is based on previous works that set token log probability as an indicator of the model's confidence score @xcite @xcite . we conduct experiments on two reasoning tasks: mathematical reasoning and code reasoning, using three datasets (gsm8k @xcite , aqua-rat @xcite , and mbpp @xcite ). our results demonstrate the broad applicability and effectiveness of the proposed method, yielding consistent performance gains. specifically, our method shows performance improvements over existing methods, with gains of up to 4.7%p on gsm8k, 3.23%p on aqua-rat in accuracy, and a 3.4%p increase in pass@1 on mbpp. we demonstrate that by incorporating [pause] tokens, the model can focus more on challenging tokens with higher loss, resulting in a higher training signal for harder tokens, improving training and inference. by using dummy tokens as a contributor to the loss term for succeeding tokens, we highlight the importance of learning difficult tokens, supported by an in-depth analysis of token probability. furthermore, our experimental analysis reveals how [pause] tokens enhance the probability robustness of outliers, solidifying their role in improving overall model performance. to summarize, our contributions are as follows: ## [pause] token goyal et al. ( 2024 ) introduce a method for pretraining models with dummy tokens (denoted as @xmath0. this approach involves using (k + m ) vectors instead of just the k existing vectors when generating the (k + 1) th token by adding @xmath1. in their study, the authors introduce a method for @xmath2:n . they insert @xmath3, throughout the sequence to get @xmath4, ỹ1:n+m p t . the model is trained using the next token prediction (ntp) loss while disregarding the loss associated with predicting @xmath5. the positions where the next token is @xmath6: ỹk+1 = <@xmath7. the pre-training loss for a decoder-only language model f θ is formulated as follows: ## ablations in this section, we conduct an ablation study on the number of [pause] tokens and the consecutive insertion lengths, which we consider important in our experimental setup. using the gsm8k dataset, which shows the most significant performance improvement over the baseline setting sft, we perform the ablation study on the dit model based on the phi-2 model. the number of [pause] token as detailed in appendix c, we insert a total of five [pause] tokens during training. a critical consideration is 1 3 5 7 9 53 54 55 56 num of total [pause] tokens num of consecutive [pause] tokens accuracy(%) figure 5: impact of varying the number of [pause] tokens on model performance. the graph shows that inserting five [pause] tokens achieves optimal performance improvement for the phi-2 model on the gsm8k dataset. the results, depicted in the blue graph in figure 5 , indicate that our choice of five [pause] tokens is optimal, yielding the highest performance improvement for the phi-2 model base using the gsm8k dataset. the graph also shows no consistent pattern suggesting that a higher or lower number of [pause] tokens is inherently better. notably, across all configurations, dynamic insertion of [pause] tokens (dit) consistently achieves superior scores compared to other methods. the role of the [pause] token is pivotal: strategically inserting [pause] tokens in front of tokens with a likelihood below a certain threshold can aid in learning challenging tokens. however, as evidenced by the performance degradation of aaw and ran in ## related works token insertion recent decoder-based language models @xcite @xcite generate the upcoming token based on the probabilities of previous tokens. this characteristic has been extensively explored in various studies, which have proposed incorporating special tokens to modify the token log-likelihood values. for instance, some have introduced controllable codes in language models @xcite , while others have explored the removal of tokens during processing in a vision transformer @xcite . in addition, the concept of memory transformers has emerged, where dummy tokens are prefixed to token vectors to enable additional trainable memory for storing representations @xcite . despite these advancements, there remains an opportunity for further research on the strategic insertion and removal of tokens in transformer models. model confidence training of llms using the transformer's decoder ensures that during inference, generated tokens are based on the highest log probability tokens or methods like beam search to find optimal probabilities within a range @xcite . hence, when the model generates tokens, a low log probability value for a specific token indicates that the probability differences among token candidates are relatively small. this means that the model is more likely to generate tokens with higher uncertainty. prior works indicate that this low probability considered as low confidence of the model can indicates a lack of knowledge in model @xcite @xcite . our study, which aims to bolster model confidence on a token-by-token basis, builds upon these foundational insights. ## conclusions previous research @xcite has underscored the effectiveness of the dummy token in enhancing performance by creating additional computational pathways in attention layers and delaying output generation. building on this, our study strategically employs [pause] tokens to improve reasoning tasks. we introduced a novel approach where [pause] tokens are dynamically inserted based on token log-likelihood, rather than following fixed rules or relying on randomness. this method targets low-confidence areas, enhancing model stability during training and inference. our findings demonstrate that inserting [pause] tokens at low log-likelihood positions significantly boosts the model's ability to handle challenging tokens with higher losses and improve the probability robustness of outlier tokens during inference.