Learning to Insert [PAUSE] Tokens for Better Reasoning

TL;DR

This paper introduces Dynamic Inserting Tokens Training (DIT), a novel method that strategically inserts [PAUSE] tokens at low-confidence positions in sequences to improve reasoning in large language models, achieving significant accuracy gains.

Contribution

The paper proposes a dynamic, confidence-based token insertion method that outperforms traditional fine-tuning and previous approaches in reasoning tasks.

Findings

Up to 4.7% accuracy improvement on GSM8K

Up to 3.23% accuracy improvement on AQUA-RAT

Up to 3.4% pass@1 improvement on MBPP

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 Dynamic Inserting Tokens Training (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.