ToW: Thoughts of Words Improve Reasoning in Large Language Models

TL;DR

This work introduces Thoughts of Words (ToW), a training-time data augmentation that attaches fine-grained thoughts to next-word predictions to improve reasoning and reduce factual hallucinations. By distilling 70K ToW annotations from larger models and continually pre-training on ToW-augmented data, the approach yields up to 7-9% reasoning gains and up to 10% hallucination mitigation across diverse reasoning benchmarks, without task-specific biases. The method comprises a two-stage pipeline: thoughts generation and consistency check, with words categorized into trivial, exact match, soft consistent, and unpredictable to guide reasoning. ToW is presented as task-agnostic and broadly applicable, though it acknowledges limitations in training data scale, potential biases, and control of ToW generation; future work aims to scale and broaden applications while improving reliability.

Abstract

We introduce thoughts of words (ToW), a novel training-time data-augmentation method for next-word prediction. ToW views next-word prediction as a core reasoning task and injects fine-grained thoughts explaining what the next word should be and how it is related to the previous contexts in pre-training texts. Our formulation addresses two fundamental drawbacks of existing next-word prediction learning schemes: they induce factual hallucination and are inefficient for models to learn the implicit reasoning processes in raw texts. While there are many ways to acquire such thoughts of words, we explore the first step of acquiring ToW annotations through distilling from larger models. After continual pre-training with only 70K ToW annotations, we effectively improve models' reasoning performances by 7% to 9% on average and reduce model hallucination by up to 10%. At the same time, ToW is entirely agnostic to tasks and applications, introducing no additional biases on labels or semantics.

Figures (9)

• Figure 1: Existing next-word prediction schemes suffer from factual and reasoning hallucinations. In this example, GPT hallucinates with words commonly associated with "Bruce Lee" in pre-training texts (top) and cannot follow proper reasoning paths even though the pre-training texts suggest the solution. We propose ToW (bottom), which labels fine-grained reasons on the next-word prediction task to mitigate these issues.
• Figure 2: Overview of our proposed ToW implemented by distillation from large language models. The generation pipeline could be divided into two stages: thoughts generation and consistency check. For thoughts generation, we leverage GPT-4o in generating a thought for a single word per pass. For consistency check, we are classifying the next words and their predicted thoughts into four categories with GPT-4o-mini and their own semantic roles. Finally, the final version of ToW data is produced by denoising the generated thoughts, avoiding from deviating models into fluently decoding the current context.
• Figure 3: Ablation on different data compositions. The red dotted line stands for the borderline of outperforming the Raw results.
• Figure 4: The comparison of ToW vs. Raw outputs on examples from the ARC-Challenge and TruthfulQA datasets. ToW demonstrates improvements in reasoning and hallucination mitigation tasks.
• Figure 5: On average, incorrect model predictions are accompanied by longer outputs (in tokens). This is particularly true for ToW-NoDeN across all datasets. ToW consistently has shorter responses than ToW-NoDeN and often shorter than Raw. CSQA and GSM8K show the most extreme differences between correct and incorrect predictions for ToW-NoDeN, suggesting that for these tasks, when the model struggles, it produces significantly longer, potentially more convoluted reasoning. Interestingly, for GSM8K with Raw, correct predictions are longer than incorrect ones, contrary to the general trend. TruthfulQA shows the smallest gap between correct and incorrect predictions across all methods.