The Belief State Transformer
Authors: Edward S. Hu
Kwangjun Ahn
Qinghua Liu
Haoran Xu
Manan Tomar
Ada Langford
Dinesh Jayaraman
Alex Lamb
John Langford
Venue: ICLR 2025 Poster (2025)
Keywords: representation learning
transformers
next-token prediction
reasoning
planning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We show that transformers trained with belief state representations solve tasks that conventional next-token predictors (e.g. GPT) cannot.

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

Published as a conference paper at ICLR 2025
THE BELIEF STATE TRANSFORMER
Edward S. Hul
Kwangjun Ahn'
Qinghua Liu'
Haoran Xu
Manan Tomar
Ada Langford'
Dinesh Jayaraman
Alex Lamb'
John Langford'
Microsoft Research
University of Pennsylvania
ST Austin
Univeesiity of Alberta
ABSTRACT
We introduce the "Belief State Transformer", next-token predictor that takes both prefix and suffix 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 inddennnnddeeenentt fashion. Key to this success is learning 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 Midd-le 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. httph/dwwwuuguu...... github. hobbsttwwbbbittt
INTRODUCTION
Transformer models (Vaswani et al. 2017) have created revolution in language modeling (Achiam al., 2023) with the capability to generate language with many emergent properties large scale, Examining these models for flaws in the pursuit of further progress, it's notable that they struggle with glanning-heaye problems (Bubeck et al., 2023) Momennejad 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 /Seeciion 2 Informally, belief state is sufficient amount of information from the past to predict the outcomes of all experiments in the future, which can be expressed as either a distribution over underlying world states or distribution over future outcomes. The Belief State Transformer is similar to standard decoder-only Transformer (e.g.. GPT2), except that it has encodings that run both forward and backward. Both of these encodings are fed into output heads which predict not only the next token after the prefix but also the previous token before the suffix as shown in Figure
In Section 3 we then study in depth how the Belief State Transformer performs a known-hard problem, the star graph (Bachmann & Nagarajan 2024) which is an elegantly simple sequential prediction problem known to confound next token prediction approaches. It's easy show that transformers can represent star graph solutions using known results (e.g., (Sanford et al.) 2024 Frydenlund 2024)), ) so the problem here is one of optimization. In particular, problems can be embedded within star graph problems, with parity known as difficult for gradient- based optimizers. Surprisingly, despite throwing away the backward encoder for inference, the BST solves even relatively difficult instances of star graphs with experiments detailed in [Section 3.4 Analyzing this discovery, the BST benefits from extra gradients, enabling avoidance the parity-by- gradient problem systematically. We also show that data augmentation approaches and ablations of the BST cannot solve the star graph problem systematically.
discover that parity
Building on this discovery, Section 4 Transformers recover the full belief state in a compact representation for the output head, In contrast, forward-only transformer and even modifications which predict every future token do not. This result implies that the Belief State Transformer learns maximal information from sequence- there no other objectiveeeennesentationn which pulls more relevant information into compact belief state.
proves this is general phenomenon: ideal Belief State

---

Published as a conference paper at ICLR 2025
X2
Xy-1
Xy
A
D
Output Head
W
... Y
Z
Forward Encoder
Backward Encoder
Xyek-1
Figure I: The Belief State Transformer has two encoders, one running forward and one backward with output head for the next forward token and another for the previous backward token.
The Belief State Transformer creates new allowances which we further explore in Section 5 In particular, it's easy to specify goal token (or tokens) explicitly to generate conditoend sequence. We compare this to the fill-nnthhe-midde approach (Bavarian et al. 2022). finding that the Belief State Transformer succeeds more effectively on Tinystories dataset. The Belief State approach also enables inference planning because rollouts occur with the Next head and semi-indepp comi-independent evaluation rollout quality can enabled with the Prev head. Going further, the Belief State Transformer superior to other transformer-based approaches.
embedder, with the relevant embeddings significantly
Altogether, show that the Belief State Transformer extracts more information (in theory and in practice) from a set of sequences, enabling Transformer models to perform well in new regimes. Performance at larger scale is of course important question for further consideration.
THE BELIEF STATE TRANSFORMER
This section introduces the Belief State Transformer. We start by introducing the architecture and the training objective then discuss how to utilize the model for inference.
2.
ARCHITECTURE AND ORJECTIVE
Let F1.7 be shorthand for the sequence I1, 31,.4.T7 TT. First, we set up the following networks:
Forward encoder
F(I1:))
Encodes prefix
Backward encoder
B ((x+++:T)
Encodes suffix
Next decoder
âtt1 ~ Tn( F(Z1.), ((2(4x4))
Predicts next token
Prev decoder
4 T F(rr:)) ((x++):7) Predicts previous token
The forward encoder aggregates the prefix into a latent F(Z).)) and the backward encoder aggregates the suffix into latent B( We GPT2-style encoders throughout our experiments, including baselines. The output heads T, and Tp then predict their respective tokens. In our experiments, the parameters of T,, n and Tp are tied with only the last layer differing. See Figure for an illustration.
The Belief State Transformer objective is the straightforward sum of the objectives of forward and backward Transformers conditioned on the prefix and suffix.
E log (2 F(x1a), B(2++127) log (p F(Z1a), B(I(+ker)
(2)
obvious alternative (called "Fill the Middle" (Bavarian et al. 2022)) when both prefix and suffix available is simply putting them together and then using forward encoder. Information theoretically, Fill in the Middle works, but we'll see that the Belief State Transformer has several advantages: it causes the system coalesce compact belief state which has many benefits explored here. This approach also extracts ( gradients from sequences enabling a gradient based optimizer to solve new problems discussed in the next section. See Appendixx El code and scaling rules.
Training on all prefix-suffix pairs is surprisingly efficient. First, we cache all forward forr = {Vi e [0 T] F(x1x)) and backwards bi:T+1 (Vi [1.7 T + 1] B(rar)) latents. Then the loss
2

---

Published as a conference paper at ICLR 2025
Equation (2) is computed over training examples (fr, bj) and their labels (2141, Ty for valid i with - - 1. We first compute the gradients of the output decoder, and then add them up to compute the gradients of the encoders. Since there (n) gradients over the output head per position of the forward or backward encoders this optimization saves a large amount of memory and compute.
2.2 BELIEF STATE INFERENCE
During inference time, the forward model T, (F( B(0)) is given a prefix 21 and an empty suffix G. Autoregressive sampling (ARS) is always used, where we sample the next token å from the next token decoder, add it to the prefix, and repeat. Note that since B(0) can precomputed. this approach requires no more parameters at inference time than standard forward-only Transformer. Later in Section
we study more complex inference approaches.
3 TESTING PLANNING ABILITIES WITH STAR GRAPHS
Next-tode Next-token
Data Aug.
FIM
Multi-token
Belief State
100
Success Rate
80
60
40
20
G(2,5)
G(2,20)
G(5,5)
Figure 2: The Belief State Transformer outperforms baselines in all star graph navigation tasks.
Bachmann & Nagarajan (004) propose the star graph problem as an elegantly simple yet challenging task for forward Transformers to solve. In Figure 2 reproduce their results while adding data augmentation baseline and the new Belief State Transformer results. Notably, the Belief State Transformer performs exceptionally well without relying domain-specifi. adaptations. In the following sections, we explain the star graph problem, present a new theory to account for these results, provide detailed discussion of our experiments, and ablate key design choices.
3.1 THE STAR GRAPH PROBLEM
A star graph (depicted in Figure 3) G(d, is graph with d paths of length emanating out from start node. construct graph, nodes n, sampled uniformly from training example is formatted as a sequence containing the edge list. the start and end nodes, and path of length from start to end: [E n1, ny n1. 122, n.,... Despite its simplicity, modern next token prediction models fail solve it.
Start
Path
N). A
Easy Tokens
Diffiuutt Token
This task captures a core challenge found in practical planning tasks like story writing, where creating coherent narrative requires the author to keep the story's resolution and backstory mind while progressing through each plot point.
Figure 3: Illustration of the star graph problem from Bachmann & Nagarajan (2024).
3.2
WHY DO FORWARD-ONLY APPROACHES FAIL?
As shown by (Bachmann & Nagarajan 2024 Appendix F.2) through extensive experiments, next- token predictors quickly learn "flawed cheat" strategy: soon after training begins, forward-only transformers learn arbitrarily select neighbor of the current node since, aside from the start node,

---

Published as a conference paper at ICLR 2025
each node has only one outgoing edge. This flawed strategy leads the model to choose neighbor of the start node without accounting for the designated goal during inference, limiting the test accuracy to 1/d, where d is the number of neighbors of the start node. The key issue with this flawed cheat is that once the model learns it, finding the correct solution becomes exceedingly difficult. To better understand this challenge,
provide formal evidence of its difficulty for gradient optimization.
Theorem 1 (Informal). Once the "flawed cheat" strategy is perfected, learning the correct path is at least as difficult as learning full parity functions.
A full proof is provided in Appendix B The full parity problem is a notoriously hard problem for gradient-based optimizers (e.g.. Shalev-Shwartz et al. 2017 Abbe & Sandon 2023). In fact, it is conjectured that learning the full parity function requires exponentially many samples and computations in the input dimension (e.g.. [bbee Boix-Adsera 2022).
The analysis in this section indicates that once the flawed cheat is perfected, the limited supervision available for the task leads significant intractability.
3.3 How DOES THE BELIEF STATE TRANSFORMER SUCCEED?
The Belief State Transformer trained to predict the previous token for every suffix, which prevents the problem from collapsing into the parity problem described in Theorem I, At a high level, for our approach to reduce to the parity problem, a large number of gradients must approach zero. However, the Belief State Transformer ensures that the information necessary to construct path suffixes (e.g.. n2:1) is present in the input to the output head. From this information, predicting n2 is straightforward, allowing the model solve the problem effectively.
This is an instance of a more general phenomenon: the Belief State Transformer naturally converges toward extracting compact, complete belief state as discussed further in Section 4.1
3.4 EXPERIMENTAL RESULTS ON STAR GRAPH
Here, we provide details on the Stargraph results in Figure 2] graphs: G(2,5), 5), G(5, 5), G(2, 20). In each experiment, we choose one graph topology and generate many example graph sequences, with all methods receiving the same amount of data and every baseline receiving least as much computation as the Belief State Transformer uses. For evaluation, the models are conditioned on the edge list, start, goal. and current path, with the task of next node prediction: p(n, &, n1, ni, 11x-1). We report the path accuracy, which is the percentage of correct path generations during the test time, over 10,000 evaluation graphs.
We run experiments on three types of
Empty suffix at inference. The Belief State Transformer trains an additional encoder B for suffixes. However, as discussed in [Section 2.2 we remove the dependency on B during inference time by pre-computing the backward latent bg B(0) with the end token and proceed with auto-regressive sampling of the Next decoder T. F(E, n1, Per, Di,... n2), .n)). bp). ). The model is still able to produce goal-conditioned behavior since the goal node nr is present in the prefix input to the forward encoder.
Baselines.
We select baselines that can be applied to a broad class of problems.
This baseline follows the conventional strategy of training Transformer with the next token prediction objective and teacher forcing, i.e., the model is trained by feeding the correct previous token as input.
Forward-only next-token prediction:
Data augmentation: common strategy improve performance to employ some form of data augmentation, although this requires some domain expertise to perform (Lee et al.) [022)). This baseline augments the training data by replacing the goal with subgoals to potentially improve the learning of goal-conditioned behaviors. Specifically, the goal node, usually the terminal node in the path, is replaced with an intermediate node in the path. Then, a Transformer is trained with the next-token objective on this augmented dataset.
tiil-ie-teemmdddee: The FIM approach (Bavarian et al. 2022) moves a span of text from the middle to the end, and then trains the transformer with next token prediction. This can be seen generalization of the data augmentation approach, where subsequences of intermediate nodes are used in lieu of a single intermediate node.
4

---

Published as a conference paper at ICLR 2025
Multi-token Prediction: The multi-token approach refers to Transformer trained to predict the entire path I p(n, E, RT) in one forward pass without providing access tokens from partial path (hence the name "Teacherless" in Bachmann & Nagarajan 2024
Our code and baselines adopt the GPT2 Transformer architecture, using standard hyperparameter settings for number of layers, embedding dimension, etc. All models are trained on the same dataset with the same training budget. See Appendix for additional training information and model setup.
3.5 RESULTS
As in Figure 2] the Belief State Transformer successfully learns to solve all the graphs. Since the Belief State Transformer uses an empty suffix, the parameter counts at inference time are very similar with minor variations driven by variations in output heads.
The baselines have varying degrees of success. The forward only baseline achieves at most a 1/d success rate,
learns output valid paths at random during inference time. Similarly, the data augmentation and FIM baselines perform poorly. We suspect that despite access to additional information like intermediate sub-goals or paths, the gradients which the model is exposed to are still inadequate to encourage the development of appropriate representations due to the easy availability of shortcut solutions.
The multi-token baseline, while successful in some smaller graphs like G(2,5), fails to solve more complicated graphs with more arms (G(5, 5)) with longer path lengths (G(2, 20)), reproducing prior results of Bachmann & Nagarajan (2024).
3.6 BELIEF STATE TRANSFORMER VARIANTS
Next, we ablate the Belief State Transformer to characterize its performance. The Belief w/o Prev ablation removes the previous token decoder and its objective from the training. The Belief w/o Backward ablation removes the backward encoder B from the training and inference process.
Belief w/o Prev
Belief w/o Backward
Belief State
100
Success Rate
80
60
20
G(2,5)
G(2,20)
G(5.5)
Figure 4: Ablations of the Belief State Transformer the star graph. Both the belief state objective and the backward encoder are crucial.
The results in Figure show that both the belief state objective and backward encoder are important components of the Belief State Transformer, and removing either drops performance back down to randomly guessing amongst valid paths. Prediction of the Prev token forces the Transformer to learn to represent long term dependencies which ends up being useful for goal-conditioned navigation.
4 BELIEF STATE TRANSFORMER ANALYSIS
In this section, show that the Belief State Transformer discovers compact belief state, and that the forward-only and multi-token approaches do no First, we formally define a belief state.
'The fact that the forward-only approach does not produce compact helief state is well known. We formalize here contrast with the Belief State Transformer.

---

Published as a conference paper at ICLR 2025
Definition 1 (Belief State). For any probability distribution set of sequences P( for any partial sequence s = 21:t, a vector Us is a belief state for there exists randomized function funciioo g(v,) which can sample from the conditional distribution P(22+1:((21.)
By definition, a belief state captures all available information relevant for predicting the future tokens. Once the belief state is learned, there is additional useful information be everyh-en necessary for future predictions is already encoded within it.
4.1 BELIEF STATE DISCOVERY
We first show that successfully optimizing a Belief State Transformer results compact belief state. Theorem 2. Let D P(xLT) represent any given probability distribution over set of sequences. Consider an ideal Belief State Transformer that satisfies the following conditions for all prefixes Tı:t and suffiess 2+tk+1:T
(2041 F(ara); B(I++++1T)) P Fia, (etk+IT)))))
T (T F(a B B(E++R+TT) = P 2i Pe4R+1
(4)
Then, for any partial sequence ana supported by D, the forward encoding of the ideal Belief State Transformer, F( is a belief state for Dia-
Proof. Let 8t Z1:t denote the prefix. To prove that ft == F(st) a belief state for st. we must show that, given Je, one can sample from the conditional distribution P(2++1.7 (s).
The key observation that the conditional distribution can be decomposed as follows:
P( 8e) = P(xr 8t) ( 8t, (T) ... PP(x++1 8t, T+22:)).
For an ideal Belief State Transformer, this decomposition can be rewritten as:
(xT fe, B(0) (TT-1 fe, B(xT)) ((TT)))..TTTTA. (3 ft, B(Zr+2:))
Thus, using the forward encoding ft, one can generate the remaining sequence 21+1:7 by sampling in reverse ceder-first sampling zT, then ø7-1 conditioned ær, and on, until 2 Each step involves using the Prev decoder and updating the backward encoder with the newly generated token.
Since the forward encoding ft enables sampling from the conditional distribution P( s:) in this way, follows that fe is belief state for 8t.
4.2 NOXT-TOKEN OR MUL.TITOKKNN PREDICTION DOES NOT GUARANTEE BELIEF STATES
Theorem establishes that an ideal Belief State Transformer learns correct belief states. contrast, the following two theorems demonstrate a fundamental limitation of standard next-token predictors and their variants when viewed from the belief state perspective.
Theorem 3. Consider standard next-token predictor with forward encoder F and output head T such that
T(F(I1:)) = P( x1:)).
There exists distribution P(x1:) over sequences and next-token predictor the form Equation (5) such that the inpul to the output head is not a belief state.
The proof is provided in Section B.2 Next, analyze the case of multi-token training (see Section 3.4), where the model is asked to predict multiple future tokens at once.
Theorem 4. Consider the multi-token setting where, for predicting H tokens into the future, the model is of the form
(F(((z1.)) =((e4 (11)) for j = 1,2,...,
(6)
There exists a distribution D P(.1:) over sequences and a multi-token model satisfying Equa sion such that the input the output head is not belief state.
The proof is provided in Section B.3 In summary, the analysis in this section highlights the inherent limitations of next-token predictors: the input to the output head fails capture the belief state.

---

Published as a conference paper at ICLR 2025
5 EXPERIMENTING WITH THE BELIEF STATE TRANSFORMER
Given model jointly trained on prefixes and suffixes, the representation is useful in new ways which are not available to simple forward Transformers. Here we detail two different forms of search based on forward and backward probabilities (respectively) as well as belief state embedding extractions.
Setup. We use TinyStories (Eldan & Li 2023). dataset consisting of synthetic short stories. TinyStories aims represent key challenges text generation while keeping training tractable for small to medium scale models. We tokenize the dataset into a vocabulary space of size 1000, and discard stories greater than 256 tokens long resulting dataset consisting of 2.7 million stories.
During evaluation, the models generate text using prefix-suffix snippets from am evaluation set of 100 unseen stories. Given stories from two competing models, GPT4 then asked output an analysis of each story examining multiple factors (e.g. grammar, flow, cohesiveness, creativity) before outputting a final recommendation. We follow best practices evaluating with multiple trials and shuffling choice order. We report the winrate and confidence interval (CI) for each model. See Appenddx D.2 for more details and examples of the GPT4 judge outputs and scoring.
5.1 GOAL-CONDITIONED TEXT GENERATION
In the goal-conditioned setting, the user provides the model with a prefix and suffix, and the model infills the text in between. See below for example of the prefix and suffix. We describe goal- conditioned planning procedure with the Belief State Transformer for text generation in ilAllorihhm
Method. The algorithm performs roll-outs. In each roll-out (starting at line 3), a candi- date trajectory is generated that differs from the greedy trajectory but maintains high prob- ability. The key features of the process are as follows:
Algorithm Goal-conditioned Planning
Require: prompt Z1:t, goal Te+k:Ts horizon & rollouts
1: Priority Queue Q <- (1,
2:
Candidates C f- 0
Roll-outs start here
Greedy generation: Starting with se- quence popped from the priority queue Q, a candidate trajectory is extended up to a max- imum length k using argmax greedy selection (line 6). The completed trajectory is appended to the set of candidates C' (line 11).
3:
for j = 1,..., do
4:
(priority sequence 8) + pop(()
5:
while s k
6:
Greedy generation
Fmax &- arg max, Tr.(z) (x: F(s): B(x++x:T)
7:
Priority queue update for future generation value T Tiax= =Tn(fmms F(*), B(I++k:T)
8:
for # Zmax do
Priority queue update for future generation: Simultaneously, the priority queue Q is up- dated to track alternative candidate sequences.
Q+ (r. Tn (())
9:
x)
Tiaa
10:
E- 8 ++ Imax
11:
Appending C with greedy generation
This is done by appending non-greedy tokens to the queue, with their priority set by the rela- rive suboptimality of each token. Specifically, lines and add alternative tokens with pri- ority equal the current sequence priority multiplied by the ratio of the alternative token's probability to that of the greedy token. Since this ratio < 1, priorities decreases with suboptimality. Also, this ensures that partial sequences of different lengths remain comparable. as suboptimality is independent of sequence length. This encourages branching ambiguous points.
C CU {s x146T]
12:
Scoring
Output: arg max=1 of Equation (7)
3. Scoring: Once candidate trajectories are generated, they are scored by evaluating the consistency of generated tokens I++1:4+6-1 with the goal T++k:T using the next-head probability:
T
II (x (1)
i=
We then return the highest-scoring trajectory based Equation (7)
Baseline. We select the Fill-intthemmidlle (FIM) (Bavarian et al. 2022) approach as a natural goal-conditioned baseline. FIM trains single forward-only transformer where the input is the
7

---

Published as a conference paper at ICLR 2025
concatenated suffix and prefix, and the objective is next token prediction. Because the forward-only approach is radically less efficient, we sample suffixes uniformly at random during training. During inference, beam search (Graves) 2012) is performed to search for high sesh--rbbbilltyy sequences.
For fair comparisons, we train the FIM and Belief State Transformer with similar amount of resources in terms of wall-clock and GPU usage, and set the architecture configuration number parameters between models similar (80M for Belief, 85M for FIM). During inference time, we configure the search methods that the computational budget is similar, i.e. setting the number of beams and depth to the values we use in Algorithm See Appendix D.I for more details.
that the
Results. As in Figure GPT4 prefers the Belief State transformer 64% % of the time with binomial confidence interval of [54%, 74%]. Qualitatively, the stories of Belief State Trans- former are superior to FIM in several ways. The Belief State outputs more frequently display clear narrative structure, with beginning, middle, and end, while the FIM model tends to gen- erate shorter stories with simpler structure and repetitions. The FIM model also frequently failed to connect its completions to the suffix, resulting in abrupt and grammatically incorrect transitions. See below for example.
Wiralee (%)
80
20
Figure 5: Storytelling winrates.
Prefix: Once a time, there was king. He war big and strong king who raled over his kingdom. One day, he wanted to take nice long buth, so he filled the bathau with warm water
Belief: He got and started to wash. The king was very busy, but he was so basy that he didn't 't notice the water was foot. looked down and saw burning. He was very scared and he quickly his foot the water. The king very worried. He
getting too hot Suddenly, felt sharp pain
FIM: and lots of bubbles. stepped into the buthtub and started to play with the bubbles. Suddenly. the king the bubbles getting higger and bigger. asked the king. "What's 's happening"" The king replied, "The bubble is getting bigger and bigger. It's getting
noticed that
Suffix: quickly grobbed cloth and begon cleon decided to take another sook in the bathtub. He put a lot of bubbies in the water to make it nice and bubbly, He relazed avazy.
The king
from cleaning up the that he
again felt
In this example, the Belief State Transformer's story has clear narrative structure with setup, conflict and resolution. It is also able correctly connect the prefix, generated text, and suffix. On the other hand, the FIM model has repetitions ("bigger and bigger"), confusing dialogue where a previously unmentioned person asks the king a question, and abrupt and grammatically incorrect transition to the suffix ("It's getting quickly grabbed"). See Sce|Appeniices D.2 and D.3 for more examples as well as the grade reports and evaluation outputs of the GPT4 judge.
5.2 UNCONDITIONAL TEXT GENERATION
Method, Next, we investigate the unconditional setting, when the goal is unknown. We reuse the logic in Algorithm with two following changes: setting the goal to the empty input zje == {0}, and scoring the sequences with the previous head T, over fixed amount of suffix tokens k. We choose to use the previous head rather than next head as semi-independent evaluator next next-generated tokens to reduce the bias associated with self-evaluation. By scoring over the last & tokens, this selects for trajectories whose endings are more likely.
II Tp(:a (((11-1)) Baati1,))((s-s)) B(x++1a)
(8)
Results. As seen in Figure 6, the Belief State transformer outperforms the FIM model. The Belief State Transformer outputs consistently have correct grammar, whereas the FIM models often abruptly end in the middle sentences.
40
Winrate (%)
30
This grammatical flaw is consistently picked up by the GPT4 judge, and results in the FIM's lower winrate. See below for an example.
20
10
FIM
Belief
8
Figure 6: Unconditioned winrates.

---

Published as a conference paper at ICLR 2025
dog. protect you The felt sufe with girl. was very kind and the mole soon come to thut her He was
Prefix:
scared
happy have new friend. She wanted lavorite things. She took him to the park and showed him all the different things she could see. The mole was so excited see different things. The little girl was happy to see the mole having fun.
Bellef:
so gentle and kind. The little
the mole all her
ran
played
kids. He even ball toy
load He looked up and saw big. dog. The dog was barking and growling at the mole. The mole was seared that he ran away as fast as be could. The dog chased after mole. the mole was too fast. The dog was getting closer and clos
always there to protect
day,
was walking through woods
Suffix:
5.3 ABLATIONS
Finally, we conduct few ablations to verify our design choices. First, the Belief Beamsearch ablation uses beamsearch to gen- erate the candidate set rather than the priority queue scheme in Alglorrtt 1 Next, in the unconditioned setting, the Belief Next Score ablation uses the next head Tn to score the suffix instead of Tp-
Winrate (%)
40
FIM
Bearnsearch Tie
Figure 7 shows the performance drops in both cases. The Belief Beamsearch outputs is more prone to repetition than
Winnate (%)
30
Algorithm I] Next, the Belief Next Score ablation frequently outputs stories that end abruptly and incorrectly, similar the
20
FIM baseline in the unconditional setting.
FIM
Next
Tie
5.4 UNDERSTANDING BELIEF STATES
Figure Ablations.
We investigated the representations learned by the Belief State Transformer to provide insight into what the states capture and how they are learned. First, we verified that the representations are indeed belief states, by training small MLP networks from the first hidden state on the G(2,5) Stargraph to predict the future token on specific timestep. We found that the Belief State Transformer has the information predict all future tokens, while the information contained within the state learned using the Forward-On objective is incomplete (Figure 8) top-left).
An analysis of how the probe accuracy improved over the course of training revealed a clear trend. The belief state captures information about the tokens at the end of the sequence at the beginning of training, while information about the earlier tokens is learned later (Figure 8) top-right).
The same probing technique can understand how well the Belief State successfully captures in- formation about the graph description. We found that the Belief State Transformer captures more information about the graph description than the Forward-Only model (Figure 8 bottom). Capturing the graph description completely is sufficient but not necessary for learning the belief state, since all the information the graph description might not be used for predicting the future tokens.
OTHER RELATED WORK
Many previous works explore non-lefi-to-right approaches for sequence modelling. ELMo (Peters) al. 2018) and BERT (Devlin al.) 2018) were prominent approaches that trained language representations with past and future information. Gu et al. (2019) model the generation order as latent variables and performs beam search over the generation order itself produce the final text. This is beneficial because in general, the best ordering for decoding is task task-dependent. On the other hand, Welleck al. (2019) explore generating text in a binary tree order, but their model struggles to outperform traditional left-to-right autoregressive generation in tasks like language modeling, machine translation, sentence completion, and word reordering. In Nguyen et al. (2024), they propose training two separate transformers and implement several strategies to "meet in the middle" during decoding. In contrast, our approach involves jointly training a transformer capable of both forward and backward decoding, which has important implications for creating compact belief states. A general study of the quality of backward prediction is done [Papadopoulos et al. (2024) discovering that backward prediction is possible but generally slightly worse than forwards

---

Published as a conference paper at ICLR 2025
Future Token Prediction Accuracy
1.00
Future Token Prediction Accuracy
1.0
0.5
-
2
-
0.0
0
1
2
3
0
0-56
4
0.00
w/o context
Fonward-Only
Belief State
20
25
30
35
Iterations (x 1000)
Graph Description Prediction Accuracy
1.0
0.5
0.0
6
7
12
13 Belief State
14
15
16
w/o context
Next token
Figure 8: Probing results star graphs show that Belief State Transformer's learned representation captures more information about future tokens (top-left). Top right: this belief state develops earlier in training for tokens which are further in the future. Bottom: besides future tokens, the inputted graph descriptions are also captured in the embeddings.
prediction. Somewhat further afield, combined forward/backwwar decoding approaches with RNNs have proved useful (Mou al. 2015 Serdyuk et al. [2018] Mangal et al. 2019) and similarly for neural machine translation (Liu et al. Zhang et al. (018))
The concept of predicting multiple tokens has been explored in other works. Gloeckle al. (0224) Cai et al.] (2024) aim to increase decoding speed by employing models that predict multiple tokens rather than just the next token. DeepSeek-V V3 (Liu et al., (024) pretrains the LLM to predict the next two tokens, and shows the objective is beneficial in the large scale setting. Such approaches align with theoretical insights that predicting multiple tokens fosters the creation of more robust
representations of future contexts.
Shai et al. (2024) show that transformers create non-compact representation of belief states within their residual stream. Compact belief states are accessible in state-space models (Hasani et al. 2020
Gu & Dao 2023). although this comes with different trade-offs compared to transformer based approaches. For example, the Mamba training process is known to fail on star graphs (Bachmann & Nagarajan 2024).
CONCLUSION
The Belief State Transformer advances goal-conditioned next-token predictors, effectively addressing the limitations of traditional forward-only transformers. The ability of our model to learn a compact belief state provides a maximal solution: since the belief state contains all the information useful in predicting the future, more complex objective can elicit more information. Through experiments with both star graphs and story writing tasks, we have demonstrated the necessity of each component of our model, particularly in challenging scenarios where standard transformers struggle.
Looking ahead, while our results demonstrate the superior performance of the Belief State Transformer in small scale story writing, exploring its application other goal-conditioned tasks would be valuable. Our current experiments serve as proof-of-coneee Further investigation into the scalability of our approach to larger practical scenarios is essential.
10

---

Published as a conference paper at ICLR 2025
8 ACKNOWLEDGEMENTS
This work was supported in part by the NSF CAREER 2239301, NSF 2331783, and DARPA TIAMAT HR001 12490421 grants. The authors would like thank the MSR AI Frontiers team for fruitful discussions and support, and the ICLR reviewers for their constructive feedback.
REFERENCES
Emmanuel Abbe and Enric Boix-Adsera. On the non-universality of deep learning: quantifying the cost of symmetry. Advances in Neural Information Processing Systems, 11117788-17201, 2022.
Emmanuel Abbe and Colin Sandon. Polynomial-time universality and limitations of deep learning. Communications on Pure and Applied Mathematics, 76(11):3493 3549, 2023. doi: https:// doi.org/10..100/ccaa2222 1002/cpa. 22121
httpt::/://on httpooonnnnniaarrr wwwppw.wwww..................... com/ htte//doo/aaa///////ee meoi/abbs/ie
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.088774 2023.
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. In Forty-first International Conference on Machine Learning, 2024.
Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry Tworek, and Mark Chen. Efficient training of language models fill in the middle. arXiv preprint arXiv:2207.14255 2022.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, et al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint vrXiv:2303.12115 2023.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng. Jason Lee, Deming Chen, and Tri Dao. Medusa: Simple I1m inference acceleration framework with multiple decoding heads. arXiv preprint arXiv:2401.1077 2024.
Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 04805, 2018.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv preprint arXiv:2305.07755 2023.
Arvid Frydenlund. The mystery of the pathological path-star task for language models. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference Empirical Methods in Natural Language Processing, pp. 12493-12516, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10. 00111886611000 http://w11/222...mmmm//// 695. URL ott 1at.:/ooo..or/2222222mmmmmmmmm.. -mnnp-main.nn...m 695/
Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Synnaeve. Better & faster large language models via multi-token prediction. arXiv preprint 1977777404.19737 2024.
Alex Graves. Sequence transduction with recurrent neural networks. arXiv preprint 2012.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint arXiv:2312.00752, 2023.
Jiatao Gu, Qi Liu, and Kyunghyun Cho. Insertion-based decoding with automatically inferred generation order. Transactions of the Association for Computational Linguistics, 7:661-676, 2019.
Ramin Hasami, Mathias Lechner, Alexander Amini, Daniela Rus, and Radu Grosu. Liquid time- 04439
constant networks, 2020.
11

---

Published as a conference paper at ICLR 2025
Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. In The Twelfth Internationaa Conference on Learning Representations, 2024.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437. 2024.
Lemao Liu, Masao Utiyama, Andrew Finch, and Eiichiro Sumita. Agreement on target-bidirectionll neural machine translation. In Kevin Knight, Ani Nenkova, and Owen Rambow (eds.), Proceedings of the 2016 Conference of the North American Chapter of the Association for Computational
Linguistics: Human Language Technologies, pp. 411-416, San Diego, California, June 2016. Association for Computational Linguistics. doi: 10.18653/11/11- 1046. URL https://
acl ccllanhaooggyooo (aargggNN6612222 6-1046,
Sanidhya Mangal, Poorva Joshi, and Rahul Modak. Lstm vs. gru bidirectional rnn for script generation. arXiv preprint arXiv: 04332, 2019.
Ida Momennejad, Hosein Hasanbeig, Felipe Vieira Frujeri, Hiteshi Sharma, Nebojsa Jojic, Hamid Palangi, Robert Ness, and Jonathan Larson. Evaluating cognitive maps and planning large language models with cogeval. Advances in Neural Information Processing Systems, 36, 2024.
Lili Mou, Rui Yan, Ge Li, Lu Zhang, and Zhi Jin. Backbone language modeling for constrained natural language generation. CoRR, abs/1512.06612, 2015. URL /R//tttp///axxi.or////aaa
1512. 1512.06612
Anh Nguyen, Nikos Karampatziakis, and Weizhu Chen. Meet in the middle: new pre-training paradigm. Advances in Neural Information Processing Systems, 36, 2024.
Vassilis Papadopoulos, Jérémie Wenger, and Clément Hongler. Arrows of time for large language models. arXiv preprint arXiv:2401.15505, 2024.
ME Peters, M Neumann, M lyyer, M Gardner, Clark, K Lee, and L Zettlemoyer. Deep contextual- ized word representations. corr. arXiv preprint arXiv: (00:: 05366, 2018.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarith- mic depth. arXiv preprint arXi::240..09268, 2024.
Dmitriy Serdyuk, Nan Rosemary Ke, Alessandro Sordoni, Adam Trischler, Chris Pal, and Yoshua Bengio. Twin networks: Matching the future for sequence generation, 2018. URL https
06742
Adam S. Shai, Sarah E. Marzen, Lucas Teixeira, Alexander Gietelink Oldenziel, and Paul M. Riechers. Transformers represent belief state geometry in their residual stream, 2024. URL
15943
Shai Shalev-Shwartz, Ohad Shamir, and Shaked Shammah. Failures of gradient- deep learning. In International Conference on Machine Learning, pp. 3067-3075. PMLR, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Sean Welleck, Kianté Brantley, Hal Daumé lii, and Kyunghyun Cho. Non-monotonic sequential text generation. In International Conference on Machine Learning, pp. 6716-6726. PMLR, 2019.
Zhirui Zhang, Shuangzhi Wu, Shujie Liu, Mu Li, Ming Zhou, and Tong Xu. Regularizing neural machine translation by banget-iiireettinall agreement, 2018. URL https: //tpiv/.rr//v/////// abs/ 1 8088000000 04064
12

---

Published as a conference paper at ICLR 2025
A STATEMENT OF CONTRIBUTIONS
Edward Hu developed the BST, set up the code and infrastructure, proposed and ran the experiments. Kwangjun Ahn developed the BST, assisted with the Stargraph experiments, and contributed to theoretical analysis. Qinghua Liu contributed the parity proof and gave feedback to the project. Haoran Xu assisted in debugging, architecture code, ran the probe analysis, and gave feedback. Ada Langford contributed code for the Stargraph experiment. Dinesh Jayaraman advised EH and gave feedback. Alex Lamb oversaw the project, developed the BST, debugged models, proposed experiments, and ran the probe analysis. John Langford oversaw the project, developed the BST, contributed BST theory, proposed experiments. and assisted with writing.
baselines, and gave feedback. Manan Tomar contributed early
B PROOFS OF THEORETICAL RESULTS
B.1 PROOF OF THEOREM
In this section, we provide formal statement and proof of Theorem
As discussed in the main text, the starting point of our argument is the observation by (Bachmann) & Nagarajan 2024 Appendix F.2) that the model quickly learns the aforementioned "flawed cheat" solution, which prevents it from learning the true solution. Specifically, next-token predictors rapidly adopt this flawed shortcut, but make little progress in correctly predicting the first vertex the path. Below, we provide formal evidence that learning the flawed cheat indeed hinders the model from learning the true solution.
In essence, once the flawed cheat solution is perfected, the model receives supervision only from the prediction the first vertex after the starting node. This effectively reduces the star graph task the following simplified task:
Task 1. Given a stargraph and a goal node, predict the correct neighbor of the starting node.
The formal statement of Theorem I is that Task is at least as difficult as learning the full parity function (i.e.. predicting whether the sum of the elements in a binary string is even or odd).
Theorem 5. For every full parity problem, there exists stargraph such that solving Task provides a solution to the parity problem.
Notably, empirical evidence shows that gradient-based methods struggle to learn the full parity function in standard training setups, especially in high dimensions (e.g.. Shalev-Shwartz al. 2017] Abbe & Sandon 2023). It has been also conjectured that learning the full parity function requires number of samples and computations exponential in the input dimension (e.g..AAbe & Boix-Adseral 2022). Consequently. Task inherits this difficulty, as it encapsulates learning the full parity function as special case.
Parity
x[1]
x[2]
x[3 =1
x|4] =
x|5] =1
1
2
-3
4
5
6
Star Graph
-1
-2
3
4
-5
Figure 9: Reduction from learning parity learning star graph. Note that if parity(#) = 0 (even), then the shortest path from center () to vertex n passes 1. Otherwise, it passes -1. As a result, solving the star graph instance induced by æ implies correctly predicting the parity of ...
Proof of Theorem 3, Consider a star graph with 2n + vertices. For simplicity, we index the vertices by .., -1, 0, 1, .,..., n. Given binary string x E (0, 1)", we generate star graph instance as
13

---

Published as a conference paper at ICLR 2025
follows: First make vertex 0 the center and connect it to vertex #1. For each i fl...., n 1}. ar((] = 0 we connect vertex to + 1 and vertex -1 to -i - 1. Otherwise, we connect to -1 1 and -i 1. We illustrate this reduction in Figure 9
The parity of z can now be determined by the first step towards target vertex n. Specifically, parity x (even), then the first vertex is 1. Otherwise, is -. Consequently, if algorithm solve lask for star graphs generated in this way, then it equivalently solves learning the full parity problem.
B.2 PROOF OF THEOREM 3
We prove this theorem by constructing counterexample. Consider the distribution D that is uniform over the sequences (ACA, BCB). Our goal is to show that there exists an ideal next-token predictor, in the sense of Equation (5) that does not learn a belief state.
First, note that an ideal next-token predictor achieves a log loss of 1 bit on the first token and 0 log loss on all subsequent tokens. Define the forward encoder F( output a two-dimensional vector for any partial sequence 21:1- Specifically, we define the encoding as follows:
F(0) (-1, 1),
F(A) = F(B) = (-1, 1);
F(AC) (1, -1);
F(BC) = (1.
Next, we define the output head T for the next-token predictions:
T(-1,-1) uniform (A, B),
T(-1, = C,
T(1, -1) A,
T(1, = B.
It can be verified that this setup achieves the optimal performance: 1-bit log loss on the first token (since and are equally likely), and log loss on subsequent tokens (since the continuation of the sequence is fully deterministic). However, the key observation is that F(A) F(B) (-1, is not a belief state for the remainder of the sequence. This is because any function g(-1, 1) applied to this encoding outputs distribution that is independent of whether the initial token was A or B. Thus, the encoding fails to capture the information necessary to distinguish between the two possible continuations of the sequence, violating the definition of belief state. Therefore, this next-token predictor does not output a belief state, completing the proof.
B.3 PROOF THEOREM 4
We again use a counterexample construction to prove this theorem. Consider the distribution D. which is uniform the four sequences (DAA, DBB, SAB,AAB, SAB, SBA). Here, the initial tokens D and S determine whether the next symbol is doubled or not. Our goal is to demonstrate that a multi-token model, in the sense of Equation (6) does not learn belief state.
First, we construct an ideal multi-token model that suffers a log loss of 2 bits on each token except for the last token. Note that this is the minimal log loss achievable. Define the forward encoder F( to output a -wo-dimensional vector for any partial sequence T1:e. Specifically, we define the encoding as follows:
F(0) = (-1, -1);
F(D) F(S) (-1,1),
F(DA) = F(SB) = -1,-1))
F(DB) = F(SAA) (1,1).
14

---

Published as a conference paper at ICLR 2025
Now, define the output head for predicting the next tokens at different time steps as follows:
Ti((-1,-1) uniform S, D),
Tz((-1,-1)) uniform A, B);
Ta((-1,-1)) = uniform A, B),
Tı((-1,1)) uniform A, B),
T2((-1,1)) uniform A, B),
T1((1,-1) A,
Ti((1,1)) B.
This multi-token model clearly achieves the minimal log loss. However, the key observation that F(S) = =F(D) (-1, 1) is not belief state for the remainder of the sequence. The reason is that any function 9(11,1) 1) applied to this encoding outputs a distribution that is independent of whether the initial token was 5 or D. Thus, the encoding fails to capture the necessary information for predicting the future tokens. Therefore, this multi-token model does not output a belief state, completing the proof.
C STAR GRAPH DETAILS
In this section, we provide detailed descriptions of the star graph experiments from Section 3.4
Data Generation. We use the data generation code from the official codebase. To generate the star graph structure, each node n, is sampled uniformly from the set [I....,N} For all experiments, set N 50, and generate 8M examples for each data set.
Data Tokenization, We follow the same tokenization settings as (Bachmann Nagarajan
2024 Section G.1).
Architectures. Both the forward and backward encoders consist of Mayers = layers with an embedding dimension of 768, 7hhaad attention heads, and an MLP expansion factor of 1. The baselines use the same configuration.
Training Details. In all cases, we use the AdamW optimizer with weight decay strength of 0.1. For G(2,3 5). the learning rate is set to 1) 3 10 while for G(5, 5) and G(2, 20). smaller learning rate of r) = 10-4 used. We run all experiments for 100 epochs to ensure convergence. All models run quickly, finishing 100 epochs in less than hours. The belief state transformer reaches 100% accuracy in a few minutes on the casier graphs. Each model is trained single A100 / H100 GPU with 80GB memory. Evaluation. We evaluate the models on 10,000 unseen graphs. To succeed, the model must output the entire path correctly, and we report the success rate.
C.1 BASELINE DETAILS
We the official codebase from Bachmann & Nagarajan (2024 instantiate and train the baselines.
Forward-only. This baseline trains a transformer to do next-token prediction over the sequence data. Teacher foreing is used during training, the model gets ground truth intermediate sequences as input.
Data Augmentation. Here, data augmentation process entails simply modifying the data using some domain knowledge, before feeding it into the standard forward-only transformer training process.
Given the original training sequences [g R. ni. R22 R3.... ny), we propose to replace the an intermediate node in the path. The resulting training sequence would then look like: [g n1.n' nı, f2, ...... ..). The idea here is that replacing the long horizon goal with a subgoal closer to the start would make learning the goal conditioning easier.
task specification n1, ny with ni, n' where n' {n2 ... .1-...
Fillin-the-MMddlle This similar to the domain augmentation baseline, except we replace intermediate nodes with intermediate suffixes n...... ny. The input the transformer is then: 71, n2, nz, nt). During training time, we uniformly sample
[E n, Ry ny. ni suffixes of different lengths.
.. ny
15

---

Published as a conference paper at ICLR 2025
Multi-token Prediction. This baseline changes the objective from next-token prediction, to multiple token prediction. Given the graph description, start and goal. [[ 1i,, ny] the model needs to predict the path [n,, n2, ...... ... Ry] in single forward pass.
D TINYSTORIES EXPERIMENTS
We train all models on a single A100 / H100 GPU with 80GB memory.
D.I TRAINING
The Belief State Transformer's encoders have the following settings: Mlayers 8, blocks with embedding dimension édim 768, and Theads = 8. The textheads T.. T are implemented as single 3-layer MIPP with dimensionality 512 and ReL.U activations that outputs two predictions. The total model size is 80 million parameters.
The FIM baseline trains forward only transformer with the following settings: Mayyrr = 12, blocks with embedding dimension Edim 768, and Yiheads = 8. The total model size is 85 million parameters.
The FIM baseline's training is simple-for a given prefix 2iss we prepend a random suffix #kr for l ~ U(L+ + 2, 7') and train the transformer to predict 1 We create as many prefix-suffix inputs as the GPU memory allows. Even then, for longer sequences, it becomes increasingly intractable for the FIM baseline train on all possible prefix-suffix inputs.
The Belief State Transformer on the other hand, is able to train on all O(T²) possible prefix-suflix combinations for sequence of length T, by carefully accumulating text head gradients over all pairs before computing the forward and backward encoder gradients. Without such scheme, training goes from few hours to multiple days. As result, the Belief State Transformer is only around 3 5x slower than the FIM baseline's suffix-samling update, yet able to train on all O(T2) possible prefix-suffix inputs.
We trained the Belief State transformer on epoch of the TinyStories dataset with 2.7M stories, with batch size of 256. The training takes around hours for epoch. Because the FIM baseline goes through an epoch more quickly than the Belief State transformer, we allow it to run for multiple epochs with a max training time of 5 hours, and use early stopping on the validation loss select the best checkpoint.
Algorithm 2 Beam Search
Require: text id., goal TT+k.7, number of steps K, beams n
1: Priority Queue Qo < (1,0)
2: for k = 1 to K do
3: for j == to 7 do
4: if 12=11
then
5:
(priority r, sequence u) pop((kkk) 6: for Possible & do
7:
Qx + (T (ā F(Z + T1a + u)); +
8: Output: top(n, Qk)
The FIM model uses beam search (Graves 2012) to find high probability sequences at inference time. At a high level, beam search only examines and expands upon fixed amount of paths during search to keep search tractable. To keep the inference time resources comparable, We use the same number of paths trajectories we generate in Algoorithm I which is 120.
D.2 EVALUATION
We evaluate the models on 100 held out stories from the TinyStories dataset. For each story, we use the first 50 tokens for the prompt, and last 100 tokens for the suffix. The choice of 100 for the suffix length due to the varying length of the stories shorter stories would often just receive padding tokens as the suffix if the suffix window wasn't long enough.
16

---

Published as a conference paper at ICLR 2025
We use GPT4 to judge the quality of generated text. Given two generated stories, one from each model, we ask GPT4 to analyze the stories and output a recommendation on which story is superior. The stories anonymized and shuffled to prevent information leakage about its author before feed it to GPT4, For every evaluation, we shuffle the story ordering once, and repeat the the evaluation 3 times for total of 6 GPT4 evaluations for every story. Since we evaluate across 100 held out stories, this amounts to 600 total queries to GPT4, We use the OpenAl API to automate the process.
Here is an example input we feed into GPT4
System Prompt: You are evoluating two story generation models. First, andlyte the outputs and note their properiee. Then, make an argament for why one is better than another, or say that both are roughty equal.
Suory I:
Once upon a time, in One day, Tom lost his red ball. He was very sad. Tom asked his friend, Sam, "Did you see my red ball?" Sam said, fiod the red ball. Tom was very sad, Than, big bird came and dropped the red ball. The bird had the red ball in its beak. Tom and Sam were so happy! They said, "Thank you, bird!" bird flew away, Tom played with his red ball all day.
and sunny place, there
big pit. A little boy named Tom liked play the pit.
"No, but let's look together." They looked and looked. but they could
Story 2:
Once upon a time, in a warm and sunny place, there was big pit. A little boy named Tom liked play near the pit. One day, Tom lost his red ball. He was very sad. Tom asked his friend. Sam, for help. "Sam, can you help me lind my red hall?" Sam said, "Yes, will help you." They looked and looked, but they could not find the red bull. Just when they were about to give up, a big hird flew down from the sky. The hird had the red hall in its beak! The bird dropped the ball into the pir. Tom and Sam
happy. They thanked the bird and played with the red ball
Then, GPT4 responds with:
Analysis of Outputs:
Story I:
*rrammar and Style:* Simple and clear sentences, with correct grammar throughoot. - Repetitive usage of "Tom was very sad'. - "Tom and Sam were so happy!" creates clear and optimistic conclusion.
Naarativve Flow: Introduction: Sets the scene with Tom and the pit. Confiict: Tom loses his red ball. - Resolution:
big bird returns the ball.- Conclusion: Ends with Tom playing all day, providing a satisfying elosure. Story 2:
*Grammar and Style;* - Generally grammatically correct. Uses a more varied sentence structure. Dialogue includes direct speach ("Sam, can you help me find my red hall?"), adding a personal touch.
Narrattive Flow; - Introduction: Similar sctup with Tom and the pit. Conflict: Tom loses his red ball, seeks Sam's 's help.- Resolution: Big bird returns the ball, but drops into the pit. Conclusion: The story seems to cnd abruptly without further explanation about retrieving the ball from the pit.
Argument for Which Story Better:
*Conclusion Analysis:** w*Story 1tt concludes elearly with Tom playing all day, which establishes definite and satisfying cading for the reader. **Story 2 leaves a gap in logic-thoug the bird returss the ball. it drops it into the pit, and we don't I see how Tom and Sam retrieve it. The absence of this information feels like an incomplete ending.
Grammar and Surrative Both stories are grammatically correct and maintain simple language, making them repetition, which could be beneficial for reintorcement but polentially tedious for some readers. Story introduces more varied dialogue which might be more engaging.
suitable for young readers. Story employs
*Overrlll Conclusion:* Given that Story I has logical setup, development, and conclusion, **Story 1 is the hetter option. provides clear, satisfying end where Tom happy and continues to play. Story while containing more engaging dialogue, fails to wrap up the narative logically, leaving the reader with lingering questions about the outcome.
As seen in the example above, GPT4 is able to consider multiple factors about the story, from low level grammar and vocabulary choices, to higher level concepts like cohesiveness and logical flow, and uses these factors to make its final recommendation.
D.3 MODEL OUTPUTS
Here are additional stories generated by the Belief State transformer in the goal-conditioned setting.
Story I: You don't 't have to be scared of the loud dog. I'll protect you". The mole felt sn safe with the little girl. She was very kind and the mole soon came to trust her He was so gentle gentlle and kind. The little girl was so happy to have a friend. She gave the mole a big hug and they played logether all day. They ran around the gurden and laughed and had lots of fun. When it was time to go home, the little girl hugged the mole and said goodbye. The mole was so happy to have friend like her and they promised to always friends.
Story 2: Otse upon a time, in a warm and sunny place. there was big pit. little boy named Tom liked to play near the Pil One day, Tom lost his red ball. He was very sad. Tom asked his friend, Sum, "Did you my red ball!" San said, "No, but let's look together." They looked near the pit and found the red hall. Tom was very happy. He said, "Thank you. Sam!" But then. a big dog came and took the red ball. Tom and back. They ran after the dog. The dog zan into the pit and disappeared. Tom and Sam were sud. They went into the pit to search. It was dark and scary. They could not the They tried to get out, but the pit was too deep. Tom and
were scared. They wanted
ball
17

---

Published as a conference paper at ICLR 2025
Sum were stuck the pit. They called for help, but no one could hear them. They were and scared, they never out of the pit.
Story 3: Onee upon a time there was little girl named Lucy. She leved to go to the store to buy sweets with her mom and dad. On this special day, Lucy entered the store with ber mom. Lacy
big. red ball. She asked her mom, "Can I have the ball, please" Her mom said, "No, Lucy. We don't I have enough money today. We have to go home and get a new one." Lucy was sad, but understood. She asked smiled and said, "Yes, we can a new one, But first, we have to be peaceful and kind to her. Have patience and understunding. Together, you both happy!" So, Lacy smiled little girl softened and smiled back. She agreed share the toy and even let Lucy have tarn first. Lucy and the little girl played together happily.
mee. 'Can we get new ball"" Her
girl said, "Can we play together?" The
And some outputs for the FIM baseline, over the same prefix / suffix pairs from above.
Story I: You don't 'tt have to be scared of the loud dog. I'll protect you". The mole felt safe with the little girl. She was very kind and the mole soon carne to trust her. He said, "Don't be scared, I'm here to protect you". The little girl so happy and thanked the mole for keeping her safe. From then on, the little girl and the mole were the best of friends.
Story 2:
Once upon a time, in a warm and sunny place, there was big pit. little boy named Tom liked play near the pit. One day, Tom lost his red ball. He was very sad. Tom asked his friend, Sam, for help. "Sam, can you help me find my red ball!" Tom said. Sam said. "Yes, I will help you." They looked and looked. but they could not find the red bull. They
very sad. Then, big bird flew down from the sky. The bird had the red ball in beak! The bird dropped the hall into the pit. Tom and Sam were so happy. They said, "Thank you, They went into the pit to search. was dark scary. They could not lind ball. They to get out, but the pit was too deep. Tomm and Sam were stuck in the pit. They called for help, but no one could hear them. They were sad and scared, and they never got out of the pit.
Story 3: Once upon a time there was a little girl named Lacy. She loved go the store buy sweets with her mom and dad. On this special day, Lacy entered the store with her mom and dad. When they arrived at the store, Lucy was so excited see all the different kinds sweets. She couldn't believe her eyes when saw sweets. She couldn't believe her eyes when she saw all the different kinds of sweets. Lucy" 's mom smiled at her and patience and understanding Together, you can both happy!" So, Lucy
different kinds of
said," to be peaceful and kind her. smiled the girl and said, "Can we play together?" The little girl softened and smiled back. She agreed to share the toy and even let Lucy have a turn first. Lucy and the little girl played together happily.
E TRAINING THE BELIEF STATE TRANSFORMER
Prefix Transformer
Suffix Transformer
belief state
next token
prev token
Figure 10: BST input and outputs.
The belief state transformer, at a high level, takes in a prefix Ti:s, a suffix I.- and predicts the token right after the prefix, and the token right before the suffix. See Figure 10 for an example.
To create training data from sequence Ti.r, all valid prefix-suffis pairs are created from the sequence. We define a valid prefix-suffix pair as follows.
Let the prefix be Z1:1- Then a suffix Z++kr is valid for all k 2 2. This condition ensures that there is at least one token that separates the prefix and suffix, which makes it feasible to predict tokens in between the prefix and suffix.
For sequence n tokens long, there are O(n²) possible valid puefix-suffix pairs. See Figure 11 for example training inputs and their targets.
We present pytorch pseudocode illustrating a simplified implementation of the belief state trans- former (Figure 12). Additionally, we present a slightly more complex version which computes and accumu lates the gradients for the text head MLP before backpropagating gradients into the encoders (Figure 13). This implementation is more computationally efficient as it avoids backpropagating through the large transformer encoders multiple times. Gloeckle et al. (2024) adopts a similar strategy to efficiently compute gradients of an encoder with multiple heads.
18

---

Published as a conference paper at ICLR 2025
k=2
k=3
k=4
k=
input
next
prev
Figure 11: Forming training examples for the BST.
We define the training overheads. Let the sequence length be n. To scale with batch size, one would multiply the batch size with the following equations. Let h be the cost related to the output head size, and be the cost of soft attention.
GPT: O(hm++en)) first term is cost of doing text (fe] for fi;,, second term is cost of encoding the tokens Z1:4 with attention to produce f1:1.
BST: Ohhr² + en²), first term is cost of doing text head by) for all valid pairs in the O(n²) prefix-suffix pairs, second term cost of encoding the prefix and suffix tokens with attention. Note that at inference, BST is O + en²).
19

---

Published as a conference paper at ICLR 2025
import torch import torch. nn as nn
def belief_ belief_state state_objectivee (enc__F, (enc_F, enc_B, Lext_head x)
bs, shape
forward_state enc_F(x)
backward_state = (xc_BB(x.lii(()).
ft = torch. arange (I, stype=torch.int33
= torche (I, (type=torch.int322
combinatiens torch.catteiian.pooo.E.,, bt)
combinations combinations (combinations [+, 11-combinations 01>-2)]
fb_pairs combinations.clone
fb_pairs [o_pairs_comiiaaiion [1 <
f icxs ib pairs 01
b icxs fb_pairss:; 11
nt ixxixxs = (combinationsl;; 0] + 1)
f forward_statel;; f_idxs]
b backward_stateli, b_idxs]
singlc_labcls_f= x(;, nt_idxs] unsquecze (2)
single_labels_3 x[:, __iixss_ (2)
single_labels (or (siigle__aall single_labels_bb dim-2)
- ttw.__heaaooorcc.cctttoooooo
dim
fb fb_numpairs fb
logits shperss.
(oss, ib ib_numpairs, -1))
-1))
logits =
( (bs+tb__s__ cttp:ib___maarrs_2. -1)]
single_labels - single_labels. reshape (bs (oss-foonumaarr-222
loss nn. Crosstntropplooss () (legits, single_labels)
return
if name settch_size = 8
main
12
m num_tokens = 100
#Use dummy function in place of actual autoregressive transformer
enc_F nn Sm.equentiall (nn. R.m.mddedddg (nneedddn
m) on.Linear (m, m)
enc_8 Sn.eunttt (nn. Rmbeddi ng (num_tokens, m) Linerr (m, m)
text head nn Sequential in.Linear (m*2, m) In. nn In.ein
m, num num_tokens=2)1
forch.randien (0, num tokens, size= (batch_size,
loss belief_staeee belief__taae__ojjettteeeeee (enc_F, enc_B, text_head, x)
print (loss)
Figure 12: A simple implementation of the belief state transformer objective
20

---

Published as a conference paper at ICLR 2025
import torch import torch.nn as
def belief_state srrate_ojjectivee (all_f, all_b, b, text text_head, x)
bs, T = x. shapp
forward, state all
backward_state all_b.flip(1)
ft torch. (T, dtype-torch.ints int32)
bt - torch. dtype=torch.int32)
combinations combinations combinations (combinations [:, 1]-combinations(:: 0]>-2)]
corch.cartesian_pcce
bt)
fb fb_pairs- combinations.clone//
fb pa_rs fb_ pt_.parrs[comiiaaiioo [;,1] < T]
f_icxs = fb_pairsl;, 0]
b icxs ib_pairs 11
nt iaxs (combinationsl;; 0] + 1)
f forward_statel;, f_idxs]
b backward_statel:: b_idxs]
single_labels__ x[;, unsqueeeess (2)
single_labelse xl:, b b_idxs]
single_ cat ((wingll_aaaeeeeeeeeeeeeeeeeeeee labels f, single_labels b), dim=2)
logits text wtw.__eaadtoocc..aattt//oo
b], dim=2))
fb_numpairs - wwwppais..ssssiee..eeeerrrreesss shape
logits = legits. (bs, fb_numpairs, -1))
logits logits. (bms-ff_numparrs-2,, -1)]
single_labels,
reshape fbbs=fo_uupairss
loss nn. rrosssstttrrooyoosssssss...sssss () (legits, single_labels)
return lcss
name batch size =
main
T 12
512
num_tokens 100
#Use dummy function place of actual autoregressive transformer
enc_F F nn Sequential Embedding com__okeens, m), an. (m, m)
enc_ B nn. Sequentia (nn. Embedding inun_tokens m) nn. Lineir (m, m)
text head nn Se.uennntial (nn Linear (m*2, LeakyReLU) Linear
m, num_tokens=2)) x torch. randint (0, num_tokens, size= (baccbaatchssie,,11
f enc_F (x)
b exc_B(x)
$ just get the text head computation graph f = f.detach()
b = b.detach()
_f. requireeaaaad True
_b. beeuuuires_graa True
loss
_, text_ x)
$ compute text head gradients over all prefix/suffix pairs.
loss .as.kwawcrrd)))
# Update encoders with backward pass.
babckward (_..grad)
b. backward b. grad)
Figure 13: Efficient computation of all prefix-suffix losses.
21


Reviews:
[{"review_id": "S65DUV0c5q", "reviewer": "Reviewer_QZxB", "summary": "This paper trains a belief state transformer, that learns to predict the next token given the suffix and previous token given the prefix. The training is done jointly on both these objectives. The paper proposes that the training in this fashion can lead to better look-ahead when the goal state is present in the input representation. This is demonstrated with a simple planning task of star graphs where the model needs to choose a correct path to reach the goal state. Results show that this training results in a 100% success, whereas next token prediction results in chance performance. Further evaluation on tinyStories (with prefix) and unconditional generation is also performed showing superior performance.", "questions": "- For the StarGraph experiments, how was the backward latent $B(\\emptyset)$ computed? How would the performance change if the training objective had $B$ but during inference $B(\\emptyset)$ was not used?\n- In the ablation section, what is the difference between `Belief w/o Backward` and single transformer next token prediction setup?\n- Was this model trained on any other planning tasks that required look-ahead (for example stacking blocks to reach a goal state) and how did it perform in those cases?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a new form of training that results in better belief state representation learning for LMs. This is shown with the performance on the star graph task where the belief state transformer achieves 100% performance. Similarly, the results on TinyStories also shows that this training regime achieves better LLM-as-Judge scores for completion of stories with prefix and suffix provided.", "weaknesses": "The task StarGraph seems to be very specific to the type of training objective proposed in the paper. The training involves learning the tokens left-to-right and right-to-left which seems aligned to the problem. Specifically, such an solution was mentioned in Bachmann & Nagarajan (2024), where the LM can learn to predict from \"right-to-left\" starting at the goal and ending in the start state. In order to claim that the proposed training objective results in a better belief state representation in general, would require more evaluations on planning problems like BlocksWorld etc. Specifically, showing the effectiveness on planning problems from benchmarks like PlanBench [1] and ACPBench [2] would make the claim of the paper stronger. \n\n[1] Valmeekam, K., Marquez, M., Olmo, A., Sreedharan, S. and Kambhampati, S., 2024. Planbench: An extensible benchmark for evaluating large language models on planning and reasoning about change. Advances in Neural Information Processing Systems, 36.\n\n[2] Kokel, H., Katz, M., Srinivas, K. and Sohrabi, S., 2024. ACPBench: Reasoning about Action, Change, and Planning. arXiv preprint arXiv:2410.05669.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1cBYFYvFuL", "reviewer": "Reviewer_Yi5e", "summary": "This paper introduces the Belief States Transformer, a model that predicts the next token for the left context (prefix) and the previous token for the right context (suffix) using two distinct encoders. The authors evaluate the proposed approach on two tasks: finding paths on star-shaped graphs and generating stories using the TinyStories dataset. The results show that the proposed method outperforms models trained only in the left-to-right direction and the Fill-in-the-Middle (FIM) approach, which predicts the next token conditioned on both left and right contexts.", "questions": "- It is unclear how the model handles unconditional text generation. Specifically, does the generation proceed left-to-right with only the scoring head for the reverse direction, or some other way?\n- The term \"teacherless\" (Fig. 2) might be confusing. Would it be more appropriate to refer to this approach as \"multi-token prediction\"? Clarifying the terminology could improve the paper’s readability.\n- Figure 5, left: It is not clear whether the probe tests the model's ability to predict nodes on the correct path several steps ahead, or nodes that the model would generate an inference (including possible errors). Forward-only model has low scores on star graph tasks, but could it possibly predict its own outputs based on internal state?\n- Figure 6: What is meant by “graph description”? What is on the x-axis? It is challenging to interpret the results presented in Figure 6.\n- On TinyStories authors use goal-conditioned planning algorithm for inference (Alg.1). Can this algorithm be adapted to FIM?\n- The paper does not specify whether FIM was trained on full sequences or if samples were truncated during training. Was FIM trained on full sequences or truncated samples? Could training on incomplete sequences or limitations in generation length impact its ability to generate complete outputs? Could FIM and BST results from Sec 5.2 be improved by setting suffix to “.”, so both models would know that they should end generation with end of sentence.\n- There are several aspects that should be made clear for each experiment: is generation done from left-to-right or right-to-left? Are generations scored based on the next-token head, previous-token head, or both?\n\nRelated work could mention BERT as a bidirectional model trained with MLM objective and ELMo as one of the first popular language models that incorporated information from two directions.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The paper addresses the important challenge that modern architectures lack planning ahead mechanisms to rely on before generating solutions.\n- By predicting both the next token for the left context and the previous token for the right context, the model uses a more complex training scenario with more training signal (O(n^2)) compared to FIM and regular next-token language modeling.\n- Despite the more complex training, the model retains the same autoregressive inference mechanism as regular language models, ensuring compatibility with existing model inference scenarios.\n- The authors select a set of tasks (path finding on star-shape graphs and story generation) that test models' abilities to make generations that follow certain goals. They show that proposed Belief State Transformers outperform regular next-token prediction models and FIM, indicating the effectiveness of their approach.", "weaknesses": "- Comparison with FIM on Star-Shaped Graphs. The paper does not provide results for the Fill-in-the-Middle (FIM) approach on the star-shaped graph task. Including FIM as a baseline in this task would offer a more comprehensive comparison and help determine whether the improvements are due to the model architecture or the specific training objectives.\n- The paper does not discuss whether a backward language model could solve the star-shaped graph task more effectively than a forward LM. Exploring this could provide insights into the benefits of using bidirectional information in the model.\n- Models like BERT, trained with a Masked Language Modeling (MLM) objective, consider information from both directions and can be adapted for generation with additional techniques (e.g., refer to “BERT has a Mouth, and It Must Speak: BERT as a Markov Random Field Language Model”) or generate multiple tokens at once. The paper does not compare Belief State Transformer with MLM models, which could highlight the advantages or limitations of the proposed approach.\n- Since the star-shaped graph task is stated to be equivalent to parity problems, the authors do not demonstrate the model’s ability to solve parity tasks. Including such experiments would strengthen the claim about the model’s capabilities in handling these types of problems and prove that the proposed approach is superior to regular next-token transformers.\n- The paper lacks clarity in description of some results and experiments (refer to questions section).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DnhvqRih7K", "reviewer": "Reviewer_Aowt", "summary": "This work presents a new Transformer model architecture capable of both next and previous token prediction. The model leverages two separate auto-regressive Transformer decoders, one processing input sequences forward and the other backward, the outputs of which are concatenated together and fed into the model’s output head to predict each token conditional on its prefix and suffix.  Apart from efficient training, this design enables the model, called Belief State Transformers (BSTs), to effectively learn good representations for token prediction, supports more flexible inference (such as goal-conditioned decoding or backward evaluation),  and demonstrates success on complex tasks, such as star graphs, where conventional next-token-prediction Transformers are known to fall short.  In addition, BSTs also achieve better performance in story-infilling tasks than prior baselines.", "questions": "1. In Line 269, the footnote mentions “the fact that the forward-only approach does not produce a compact belief state is well known”. Could you clarify what defines a belief state as “compact”?\n2. How does FIM perform specifically on star graph tasks?\n3. What is the training overhead of BSTs compared to conventional auto-regressive Transformer decoders, particularly in terms of sequence length and batch sizes?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- This paper addresses an important problem in improving the non-myopic prediction capabilities of current auto-regressive language models. Although this study is limited to small-scale settings, its design contributes meaningfully to expanding the architectural space of flexible language models.\n- The proofs in Section 4 further strengthen the findings and nicely highlight the biases present in modeling non-myopic dependencies across various objectives (namely BSTs, next token prediction, and teacher-less training).\n- The probing experiments, as demonstrated in Figure 5, Section 5.4, are particularly interesting, as they effectively showcase the improved quality of learned bidirectional representations.", "weaknesses": "The study lacks experiments on more realistic, large-scale settings or with larger model sizes. For instance, evaluating bidirectional representations in tasks like code infilling could provide more robust insights.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WFeHAYKhsO", "reviewer": "Reviewer_1PYd", "summary": "The paper introduces the Belief State Transformer, a next-token predictor that leverages both prefix and suffix inputs to enhance prediction capabilities. By learning a compact belief state, the model overcomes the limitations of traditional forward-only transformers, excelling in goal-conditioned tasks such as story writing.", "questions": "- Based on the model parameter details in section C.1, are the forward and backward encoders shared? If so, does this mean that by simultaneously learning the backward objective, the forward encoder obtains some future information, leading to a better representation of $f_t$? Furthermore, would not using a shared encoder cause the model to fail?\n- I'm a bit confused about the proof about Theorem 3. Given the proof of Theorem 2, it seems we can also derive a similar proof to show AR also encodes a belief state through forward decomposition. The authors provide a counterexample for AR, but stating that \"F(A)=F(B)=(−1,1)” seems unreasonable, as a normal encoder won't produce exactly the same encodings for different inputs.\n- Could the authors elaborate more on lines 725-726 and show how the belief state transformer constructs a belief state that corresponds to the {ACA, BCB} distribution?\n- Why the belief state transformer performs better than FIM training? Could the authors give some intuitions about this?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The idea of learning a compact belief state that captures all relevant information for future predictions is pretty interesting.\n- The paper is mostly well written, with controlled experiments and analysis of the proposed method.", "weaknesses": "- The methods and implementations are somewhat confusing, as detailed in the question section.\n- There is no analysis regarding scaling AR and the Belief State Transformer. Considering cost issues, it might be possible to examine this on relatively small series models. The concern is whether the forward AR inherently encodes a \"belief state\" as AR model size increases.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)