# The Belief State Transformer

arXiv: https://arxiv.org/abs/2410.23506 · HTML: https://mlsanity.com/paper/2410.23506/

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 31 references that are in the map.
- First posted on arXiv 30 Oct 2024.

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Deep contextualized word representations](https://mlsanity.com/paper/1802.05365/) (2018)
- [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://mlsanity.com/paper/2401.10774/) (2024)
- [TinyStories: How Small Can Language Models Be and Still Speak Coherent English?](https://mlsanity.com/paper/2305.07759/) (2023)
- [Better & Faster Large Language Models via Multi-token Prediction](https://mlsanity.com/paper/2404.19737/) (2024)
- [Efficient Training of Language Models to Fill in the Middle](https://mlsanity.com/paper/2207.14255/) (2022)
- [Liquid Time-constant Networks](https://mlsanity.com/paper/2006.04439/) (2020)
- [The pitfalls of next-token prediction](https://mlsanity.com/paper/2403.06963/) (2024)
- [Teaching Arithmetic to Small Transformers](https://mlsanity.com/paper/2307.03381/) (2023)
- [Failures of Gradient-Based Deep Learning](https://mlsanity.com/paper/1703.07950/) (2017)
- [Transformers, parallel computation, and logarithmic depth](https://mlsanity.com/paper/2402.09268/) (2024)
- [Evaluating Cognitive Maps and Planning in Large Language Models with CogEval](https://mlsanity.com/paper/2309.15129/) (2023)
- [Transformers represent belief state geometry in their residual stream](https://mlsanity.com/paper/2405.15943/) (2024)
- [Non-Monotonic Sequential Text Generation](https://mlsanity.com/paper/1902.02192/) (2019)
- [Insertion-based Decoding with automatically Inferred Generation Order](https://mlsanity.com/paper/1902.01370/) (2019)
- [Regularizing Neural Machine Translation by Target-bidirectional Agreement](https://mlsanity.com/paper/1808.04064/) (2018)
- [Meet in the Middle: A New Pre-training Paradigm](https://mlsanity.com/paper/2303.07295/) (2023)
- [Arrows of Time for Large Language Models](https://mlsanity.com/paper/2401.17505/) (2024)
- [On the non-universality of deep learning: quantifying the cost of symmetry](https://mlsanity.com/paper/2208.03113/) (2022)
- [The Mystery of the Pathological Path-star Task for Language Models](https://mlsanity.com/paper/2410.13779/) (2024)
- [LSTM vs. GRU vs. Bidirectional RNN for script generation](https://mlsanity.com/paper/1908.04332/) (2019)

## Built on by (most cited first, 12 of 12)

- [Roll the dice & look before you leap: Going beyond the creative limits of next-token prediction](https://mlsanity.com/paper/2504.15266/) (2025)
- [Qsharp: Provably Optimal Distributional RL for LLM Post-Training](https://mlsanity.com/paper/2502.20548/) (2025)
- [Beyond Multi-Token Prediction: Pretraining LLMs with Future Summaries](https://mlsanity.com/paper/2510.14751/) (2025)
- [Solve the Loop: Attractor Models for Language and Reasoning](https://mlsanity.com/paper/2605.12466/) (2026)
- [Looking beyond the next token](https://mlsanity.com/paper/2504.11336/) (2025)
- [The Topological Trouble With Transformers](https://mlsanity.com/paper/2604.17121/) (2026)
- [The State-Prediction Separation Hypothesis](https://mlsanity.com/paper/2607.01218/) (2026)
- [Explore-Execute Chain: Towards an Efficient Structured Reasoning Paradigm](https://mlsanity.com/paper/2509.23946/) (2025)
- [What Makes the Preferred Thinking Direction for LLMs in Multiple-choice Questions?](https://mlsanity.com/paper/2502.18435/) (2025)
- [Understanding the Emergence of Seemingly Useless Features in Next-Token Predictors](https://mlsanity.com/paper/2603.14087/) (2026)
- [Belief-Guided Decision Making with Uncertainty Gating in the Game of Go](https://mlsanity.com/paper/2607.26946/) (2026)
- [Observe Before You Alert: Adaptive Driver Alerting with Vision-Language Models](https://mlsanity.com/paper/2609.08130/) (2026)

## Abstract (arXiv)

We introduce the "Belief State Transformer", a next-token predictor that takes both a prefix and suffix as inputs, with a novel objective of predicting both the next token for the prefix and the previous token for the suffix. The Belief State Transformer effectively learns to solve challenging problems that conventional forward-only transformers struggle with, in a domain-independent fashion. Key to this success is learning a compact belief state that captures all relevant information necessary for accurate predictions. Empirical ablations show that each component of the model is essential in difficult scenarios where standard Transformers fall short. For the task of story writing with known prefixes and suffixes, our approach outperforms the Fill-in-the-Middle method for reaching known goals and demonstrates improved performance even when the goals are unknown. Altogether, the Belief State Transformer enables more efficient goal-conditioned decoding, better test-time inference, and high-quality text representations on small scale problems. Website: https://edwhu.github.io/bst-website