Next-Token Prediction Should be Ambiguity-Sensitive: A Meta-Learning Perspective

TL;DR

This paper argues that next-token prediction should account for ambiguity, introducing a meta-learning benchmark and a method to improve model performance in high-ambiguity situations, inspired by cognitive science insights.

Contribution

It presents MetaHMM, a synthetic benchmark for ambiguity in sequence prediction, and proposes a method to enhance models' handling of ambiguous contexts based on cognitive theories.

Findings

Transformers struggle with high-ambiguity predictions.

The proposed method improves performance in ambiguous contexts.

Preliminary results show better capacity allocation and scalable inference.

Abstract

The rapid adaptation ability of auto-regressive foundation models is often attributed to the diversity of their pre-training data. This is because, from a Bayesian standpoint, minimizing prediction error in such settings requires integrating over all plausible latent hypotheses consistent with observations. While this behavior is desirable in principle, it often proves too ambitious in practice: under high ambiguity, the number of plausible latent alternatives makes Bayes-optimal prediction computationally intractable. Cognitive science has long recognized this limitation, suggesting that under such conditions, heuristics or information-seeking strategies are preferable to exhaustive inference. Translating this insight to next-token prediction, we hypothesize that low- and high-ambiguity predictions pose different computational demands, making ambiguity-agnostic next-token prediction a detrimental inductive bias. To test this, we introduce MetaHMM, a synthetic sequence meta-learning benchmark with rich compositional structure and a tractable Bayesian oracle. We show that Transformers indeed struggle with high-ambiguity predictions across model sizes. Motivated by cognitive theories, we propose a method to convert pre-trained models into Monte Carlo predictors that decouple task inference from token prediction. Preliminary results show substantial gains in ambiguous contexts through improved capacity allocation and test-time scalable inference, though challenges remain.