Code Pretraining Improves Entity Tracking Abilities of Language Models

TL;DR

The paper investigates whether pretraining language models on code improves their ability to track discourse entities as texts unfold. By contrasting base models with counterparts augmented with code, as well as variants trained on math data and subjected to alignment tuning, the authors conduct a systematic, multi-family analysis using a boxes-based entity-tracking task with constrained decoding. The key finding is that code pretraining yields clear, consistent improvements across model families and sizes, while additional math training and alignment tuning show limited or mixed benefits, with base models often benefiting more from alignment and the best results arising from combining code with instruction tuning. These results support the view that structured data such as code enhances reasoning and state-tracking capabilities in LLMs and inform finetuning strategies for improved entity-tracking performance in real-world scenarios.

Abstract

Recent work has provided indirect evidence that pretraining language models on code improves the ability of models to track state changes of discourse entities expressed in natural language. In this work, we systematically test this claim by comparing pairs of language models on their entity tracking performance. Critically, the pairs consist of base models and models trained on top of these base models with additional code data. We extend this analysis to additionally examine the effect of math training, another highly structured data type, and alignment tuning, an important step for enhancing the usability of models. We find clear evidence that models additionally trained on large amounts of code outperform the base models. On the other hand, we find no consistent benefit of additional math training or alignment tuning across various model families.

Figures (4)

• Figure 1: Entity tracking results for DeepSeek, Gemma, and Llama 2 models. Error bars indicate 95% confidence intervals, and the black dashed lines show the performance of the random baseline.
• Figure 2: Entity tracking results for models trained with additional math data. See Table \ref{['tab:models']} for the model names of the base and math models. Error bars indicate 95% confidence intervals, and the black dashed lines show the performance of the random baseline.
• Figure 3: Entity tracking results for alignment-tuned DeepSeek, Gemma and Llama 2 models. The top panels show models without additional code training, whereas the bottom panels show models that have been trained on additional amounts of code before alignment tuning. Error bars indicate 95% confidence intervals, and the black dashed lines show the performance of the random baseline.
• Figure 4: Regular expression used for constrained decoding of entity states.