Learning to Plan Long-Term for Language Modeling

Abstract: Modern LLMs predict the next token in the sequence by considering the past text through a powerful function such as attention. However, LLMs have no explicit mechanism that allows them to spend computation time for planning long-distance future text, leading to a suboptimal token prediction. In this paper, we propose a planner that predicts a latent plan for many sentences into the future. By sampling multiple plans at once, we condition the LLM on an accurate approximation of the distribution of text continuations, which leads to better next token prediction accuracy. In effect, this allows trading computation time for prediction accuracy.