Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Looking beyond the next token

Abitha Thankaraj, Yiding Jiang, J. Zico Kolter, Yonatan Bisk

TL;DR

This work tackles the mismatch between next-token prediction during training and human planning by arguing that data arrangement can capture long-horizon dependencies without changing model architecture. It introduces Trelawney, a data-centric augmentation that inserts future context using lookahead tokens, coupled with a mixed-data training objective and <T>-generation at inference to induce planning and controllable generation. Across Star Graph, algorithmic reasoning (CLRS-Text SCC), and Tiny Stories benchmarks, Trelawney enhances planning, long-term goal generation, and controllable storytelling while preserving standard language modeling performance. The method presents a scalable, architecture-free path to richer long-horizon capabilities and hints at new directions beyond traditional language modeling, including potential reinforcement-learning integrations and uncertainty-guided control.

Abstract

The structure of causal language model training assumes that each token can be accurately predicted from the previous context. This contrasts with humans' natural writing and reasoning process, where goals are typically known before the exact argument or phrasings. While this mismatch has been well studied in the literature, the working assumption has been that architectural changes are needed to address this mismatch. We argue that rearranging and processing the training data sequences can allow models to more accurately imitate the true data-generating process, and does not require any other changes to the architecture or training infrastructure. We demonstrate that this technique, Trelawney, and the inference algorithms derived from it allow us to improve performance on several key benchmarks that span planning, algorithmic reasoning, and story generation tasks. Finally, our method naturally enables the generation of long-term goals at no additional cost. We investigate how using the model's goal-generation capability can further improve planning and reasoning. Additionally, we believe Trelawney could potentially open doors to new capabilities beyond the current language modeling paradigm.

Looking beyond the next token

TL;DR

This work tackles the mismatch between next-token prediction during training and human planning by arguing that data arrangement can capture long-horizon dependencies without changing model architecture. It introduces Trelawney, a data-centric augmentation that inserts future context using lookahead tokens, coupled with a mixed-data training objective and <T>-generation at inference to induce planning and controllable generation. Across Star Graph, algorithmic reasoning (CLRS-Text SCC), and Tiny Stories benchmarks, Trelawney enhances planning, long-term goal generation, and controllable storytelling while preserving standard language modeling performance. The method presents a scalable, architecture-free path to richer long-horizon capabilities and hints at new directions beyond traditional language modeling, including potential reinforcement-learning integrations and uncertainty-guided control.

Abstract

The structure of causal language model training assumes that each token can be accurately predicted from the previous context. This contrasts with humans' natural writing and reasoning process, where goals are typically known before the exact argument or phrasings. While this mismatch has been well studied in the literature, the working assumption has been that architectural changes are needed to address this mismatch. We argue that rearranging and processing the training data sequences can allow models to more accurately imitate the true data-generating process, and does not require any other changes to the architecture or training infrastructure. We demonstrate that this technique, Trelawney, and the inference algorithms derived from it allow us to improve performance on several key benchmarks that span planning, algorithmic reasoning, and story generation tasks. Finally, our method naturally enables the generation of long-term goals at no additional cost. We investigate how using the model's goal-generation capability can further improve planning and reasoning. Additionally, we believe Trelawney could potentially open doors to new capabilities beyond the current language modeling paradigm.

Paper Structure

This paper contains 30 sections, 8 equations, 7 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Pitfalls of teacher forcing
  4. Trelawney
  5. Augmentation schemas
  6. Dataset Construction and Training Objective
  7. Inference
  8. Experiments
  9. Star Graph
  10. Choice of ${\mathbf{z}}$.
  11. Algorithmic Reasoning
  12. Natural Language Planning
  13. Related work
  14. Discussion
  15. Appendix
  16. ...and 15 more sections

Figures (7)

  • Figure 1: Trelawney. (Top) We modify the training data by inserting tokens from the future, delimited with special tokens $\texttt{<T>}$ and $\texttt{</T>}$, into the modified sequences. By first predicting tokens in the future, the model is encouraged to learn the tokens pertaining to what it will generate in the future (i.e., F'G'), and the path leading to the future (i.e., CDE) as well as the actual future (i.e., FG) will be easier to predict. (Bottom) Illustration of Trelawney's effect during generation. The top is the prompt, and the middle is different future tokens. The generations are coherent and read naturally.
  • Figure 2: In the star graph, there are key "hard nodes" that indicate the moment of branching, after which the path and goal become clear. Above is a visualization of the construction of ${\mathcal{D}}'$ for the star graph with corresponding linearizations.
  • Figure 3: Our evaluation of story generation demonstrates greatly improved performance when goal-conditioned, without hurting the performance of unconditional generation.
  • Figure 4: Results on Star Graph.
  • Figure 5: Accuracies - Strongly connected components Trelawney-Rule-Based
  • ...and 2 more figures