mlsanityData to 14 Oct 2026

The Belief State Transformer

arXiv 2410.23506 ↗First posted 30 Oct 2024cs.LG cs.AI cs.CLcited by 12

Not on the map. Why?

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

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

25 of 31 references.

References · 25

  1. 2017Attention Is All You Needcited by 60,361
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,399
  3. 2023GPT-4 Technical Reportcited by 18,396
  4. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,718
  5. 2023Sparks of Artificial General Intelligence: Early experiments with GPT-4cited by 2,592
  6. 2018Deep contextualized word representationscited by 2,290
  7. 2024Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Headscited by 843
  8. 2023TinyStories: How Small Can Language Models Be and Still Speak Coherent English?cited by 449
  9. 2024Better & Faster Large Language Models via Multi-token Predictioncited by 360
  10. 2022Efficient Training of Language Models to Fill in the Middlecited by 255
  11. 2020Liquid Time-constant Networkscited by 161
  12. 2024The pitfalls of next-token predictioncited by 160
13 more
  1. 2023Teaching Arithmetic to Small Transformerscited by 123
  2. 2017Failures of Gradient-Based Deep Learningcited by 80
  3. 2024Transformers, parallel computation, and logarithmic depthcited by 79
  4. 2023Evaluating Cognitive Maps and Planning in Large Language Models with CogEvalcited by 77
  5. 2024Transformers represent belief state geometry in their residual streamcited by 77
  6. 2019Non-Monotonic Sequential Text Generationcited by 55
  7. 2019Insertion-based Decoding with automatically Inferred Generation Ordercited by 43
  8. 2018Regularizing Neural Machine Translation by Target-bidirectional Agreementcited by 40
  9. 2023Meet in the Middle: A New Pre-training Paradigmcited by 22
  10. 2024Arrows of Time for Large Language Modelscited by 20
  11. 2022On the non-universality of deep learning: quantifying the cost of symmetrycited by 18
  12. 2024The Mystery of the Pathological Path-star Task for Language Modelscited by 8
  13. 2019LSTM vs. GRU vs. Bidirectional RNN for script generationcited by 7

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 6; 2026: 6.

12 papers.

  1. 2025Roll the dice & look before you leap: Going beyond the creative limits of next-token predictioncited by 29
  2. 2025Qsharp: Provably Optimal Distributional RL for LLM Post-Trainingcited by 15
  3. 2025Beyond Multi-Token Prediction: Pretraining LLMs with Future Summariescited by 15
  4. 2026Solve the Loop: Attractor Models for Language and Reasoningcited by 13
  5. 2025Looking beyond the next tokencited by 5
  6. 2026The Topological Trouble With Transformerscited by 3
  7. 2026The State-Prediction Separation Hypothesiscited by 3
  8. 2025Explore-Execute Chain: Towards an Efficient Structured Reasoning Paradigmcited by 2
  9. 2025What Makes the Preferred Thinking Direction for LLMs in Multiple-choice Questions?cited by 1
  10. 2026Understanding the Emergence of Seemingly Useless Features in Next-Token Predictorscited by 0
  11. 2026Belief-Guided Decision Making with Uncertainty Gating in the Game of Gocited by 0
  12. 2026Observe Before You Alert: Adaptive Driver Alerting with Vision-Language Modelscited by 0