Prophecy Transformer

Autoregressive transformers trained on next-token prediction suffer from a key weakness; early decisions can restrict the model to suboptimal continuations. Furthermore, they do not learn a meta-model of how their mistakes correlate with ground truth, nor how to recover. These issues are critical in the context of reasoning, where such approaches often become trapped in irreconcilable dead-ends without a way to update their predictions based on these failures. To address this, we introduce Prophecy Transformer, an architecture which conditions its generations on a nonparametric latent variable encoding the model's internal activations produced during an initial rollout following its own policy. By attending to this latent variable from the start of generation, our model gets a view into its possible future internal state which reveals information that would only become apparent after committing to certain decisions. Tested on algorithmic reasoning tasks such as 3-SAT, graph coloring, TSP, and countdown, our method significantly outperforms autoregressive models due to its prospective and self-corrective capabilities. Across tasks, Prophecy Transformer acquires self-referential representations encoding the quality of its predictions and the ability to update the model's decisions based on its predicted future state, unlocking large gains in reasoning ability.