The pitfalls of next-token prediction (2403.06963v2)

Abstract: Can a mere next-token predictor faithfully model human intelligence? We crystallize this emerging concern and correct popular misconceptions surrounding it, and advocate a simple multi-token objective. As a starting point, we argue that the two often-conflated phases of next-token prediction -- autoregressive inference and teacher-forced training -- must be treated distinctly. The popular criticism that errors can compound during autoregressive inference, crucially assumes that teacher-forcing has learned an accurate next-token predictor. This assumption sidesteps a more deep-rooted problem we expose: in certain classes of tasks, teacher-forcing can simply fail to learn an accurate next-token predictor in the first place. We describe a general mechanism of how teacher-forcing can fail, and design a minimal planning task where both the Transformer and the Mamba architecture empirically fail in that manner -- remarkably, despite the task being straightforward to learn. Finally, we provide preliminary evidence that this failure can be resolved using a simple modification that predicts multiple tokens in advance. We hope this finding can ground future debates and inspire explorations beyond the next-token prediction paradigm. We make our code available under https://github.com/gregorbachmann/Next-Token-Failures

• The paper demonstrates that teacher-forcing can impair complex planning by failing in lookahead tasks during next-token prediction.
• It introduces a minimal path-finding example that exposes issues like the Clever Hans cheat and Indecipherable Token failure.
• Empirical results on models like Transformers show that predicting multiple future tokens may mitigate these limitations.

Exploring the Limits of Next-Token Prediction in LLMs

The Distinct Phases of Next-Token Prediction

In recent years, next-token prediction (NTP) has become a central paradigm in training Generative LLMs (LMs), notably driving the success of models such as GPT-3. NTP is essentially the task of predicting the probability of the next token in a sequence, given all previous tokens. This procedure underpins both the training phase, through a mechanism known as teacher-forcing, and the inference phase, via autoregressive modeling. Despite the widespread adoption and success of this method, concerns linger regarding its efficacy for tasks requiring complex planning or reasoning.

At its core, the chain rule of probability assures us that any sequence generation task can be decomposed into a series of NTP tasks. However, a deep-rooted skepticism exists, revolving around the potential for errors to compound during autoregressive inference, thus questioning the model's capability for intricate planning tasks.

Unveiling a More Profound Issue

A deeper, yet less explored issue lies not in the autoregressive inference, but in the training phase itself—teacher-forcing. The consensus presumes that teacher-forcing effectively teaches the model to accurately predict the next token, making any shortcomings a matter of execution rather than learning. However, we posit that in certain "lookahead tasks," where predicting later tokens relies on previously imagined tokens not yet generated, teacher-forcing could fundamentally fail to grasp the required complex planning mechanisms.

The Path-Star Example: A Case of Inherent Failure

To crystallize this concern, we introduce a minimal task exemplifying the failure of NTP due to teacher-forcing. We conceptualize a directed graph-based path-finding problem, highlighting two endemic issues in teacher-forcing: the Clever Hans cheat and the Indecipherable Token failure. The former describes the model's reliance on spurious correlations present due to the exposure of partial ground truth during training, an effect that simplifies task learning but destroys the model's ability to generalize. The latter reflects the reduced supervision for critical parts of the task stemming from the model's reliance on Clever Hans shortcuts, significantly hindering the model's learning capabilities for these crucial elements.

Empirical Validation and Beyond

Our empirical investigations across different architectures (Transformers and Mamba) conclusively demonstrate the presence of these failure modes, even in the setting where the task is conceptually straightforward. Remarkably, alternative training objectives proposed to circumvent the standard teacher-forcing paradigm—specifically, those that prompt the model to predict several future tokens simultaneously—show potential in overcoming these failures, albeit in limited settings.

Reflecting on Next-Token Prediction's Future

The failure of NTP in even simple scenarios raises pertinent questions regarding its efficacy for more complex, real-world tasks, such as creative writing or advanced reasoning. This insight beckons the exploration of alternatives to NTP, prompting both theoretical introspection and empirical investigations to better understand and enhance models' planning and generalization capabilities.

As we venture into this exploration, the lessons from the path-star example and the notion of teacherless training present a promising avenue, suggesting a shift towards training paradigms that inherently encourage models to learn complex planning without the pitfalls associated with NTP. It is an open invitation for the community to explore the mechanisms of LLM training, pushing the boundaries of what these remarkable systems can achieve.