Do language models plan ahead for future tokens?

TL;DR

This paper investigates whether transformer language models plan ahead during inference by testing hypotheses about how future information is prepared and stored, using novel training schemes and synthetic data experiments.

Contribution

It introduces myopic training to test if transformers pre-cache future features or use breadcrumbs, providing evidence for pre-caching and insights into model planning mechanisms.

Findings

Evidence supports pre-caching in synthetic data experiments.

Pre-caching increases with model scale in language modeling.

Results suggest transformers may use pre-caching or breadcrumbs for future token prediction.

Abstract

Do transformers "think ahead" during inference at a given position? It is known transformers prepare information in the hidden states of the forward pass at time step $t$ that is then used in future forward passes $t+\tau$. We posit two explanations for this phenomenon: pre-caching, in which off-diagonal gradient terms present during training result in the model computing features at $t$ irrelevant to the present inference task but useful for the future, and breadcrumbs, in which features most relevant to time step $t$ are already the same as those that would most benefit inference at time $t+\tau$. We test these hypotheses by training language models without propagating gradients to past timesteps, a scheme we formalize as myopic training. In a constructed synthetic data setting, we find clear evidence for pre-caching. In the autoregressive language modeling setting, our experiments are more suggestive of the breadcrumbs hypothesis, though pre-caching increases with model scale.