ICLR 2025

Published as a conference paper at ICLR 2025
THE BELIEF STATE TRANSFORMER
Edward S. Hu1,2
Kwangjun Ahn1
Qinghua Liu1
Haoran Xu1,3
Manan Tomar1,4
Ada Langford1
Dinesh Jayaraman2
Alex Lamb1
John Langford1
1Microsoft Research
2University of Pennsylvania
3UT Austin
4 University of Alberta
ABSTRACT
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
1
INTRODUCTION
Transformer models (Vaswani et al., 2017) have created a revolution in language modeling (Achiam
et al., 2023) with the capability to generate language with many emergent properties at large scale.
Examining these models for flaws in the pursuit of further progress, it’s notable that they struggle
with planning-heavy problems (Bubeck et al., 2023; Momennejad et al., 2024). How can we modify
the architecture, objectives, and algorithms to create a model more capable of reaching goals?
To make progress, we propose the new Belief State Transformer (BST) architecture and objective
in S...