The Belief State Transformer

Thanks for the review!

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.

This is a real concern, but it appears to require significant scale in both data and parameters to test this concern since small-scale solutions have been ruled out in the star graph paper [1].

[1] Bachmann, Gregor, and Vaishnavh Nagarajan. "The pitfalls of next-token prediction." ICML (2024).

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?

No, they are not shared. Note that the BST uses 2 encoders with 8 layers, and the combined parameter count of these 2 shallower encoders is roughly the same as the FIM’s 12-layer encoder.

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.

We agree that empirically F(A) and F(B) might sometimes have different representations. However, that is not sufficient to establish that a compact belief state is learned since the proof in theorem 2 applies always, not merely "sometimes".
The normal autoregressive model (GPT) does not learn a compact belief state through forward decomposition, because there is a dependence on all previous tokens via attention across time, with no explicit dependence on a belief state.

If you want another intuitive counterexample for GPT, there’s an old toy problem called copying where a sequence is modeled [x1 x2 x3 x4 0 0 0 0 0 x1 x2 x3 x4], i.e. the model is tasked with copying a sequence. This is fairly easy for transformers, and the true belief state at the midpoint should have enough information to reconstruct [x1 x2 x3 x4]. However, a GPT-style transformer does not need to express any information about the sequence on the midpoint step’s final layer representation, because it only needs to predict the padding token on that step.

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?

The key observation is that F(A)=F(B) does not optimize the Belief State Transformer objective since T_p(F(A), B(∅)) is optimized to predict A and T_p(F(B),B(∅)) is optimized to predict B.

Why the belief state transformer performs better than FIM training? Could the authors give some intuitions about this?

There are two important effects here: (a) The Belief State Transformer is trained on O(n^2) gradients whereas FIM is trained on only O(n) gradients per update. Additionally the gradients include both the next and prev predictions, which lead to learning a belief state and could also be a richer training signal. (b) The BST next head can provide an implicit valuation of generation by calculating the probability of goal tokens given the generation. In the FIM paradigm, this is not possible because the “middle” comes after the suffix.

Dear reviewer 1PYd,

As the discussion period draws to a close, we’d like to check back to see if you have any remaining concerns. We believe we sufficiently answered your questions, but would be happy to clarify further. We would also be grateful for any additional comments or feedback.

Thanks for the review!

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.

We look forward to scaling and investigating harder domains in future work. Code infilling is a natural task that we are investigating.

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”?

We refer to the embedding input to the text head, as the compact belief state. It is desirable to capture the belief state into a low-dimensional representation like the text head embedding input, as this makes it ergonomic for downstream tasks.

How does FIM perform specifically on star graph tasks?

Multiple reviewers have asked about FIM on the Stargraph setting; we will look into running FIM on the Stargraph and aim to report results before the end of the discussion period.

What is the training overhead of BSTs compared to conventional auto-regressive Transformer decoders, particularly in terms of sequence length and batch sizes?

Here are the training overheads in terms of the sequence length $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 $e$ be the cost of soft attention.

GPT: $O(h n+e n^2)$, first term is cost of doing text_head$(f_i)$ for $f_{1:n}$, second term is cost of encoding the tokens $x_{1:n}$ with attention to produce $f_{1:n}$.

BST: $O(h n^2 + e n^2)$, first term is cost of doing text_head$(f_i,b_j)$ for all valid pairs in the $O(n^2)$ prefix-suffix pairs, second term is cost of encoding the prefix and suffix tokens with attention. Note that at inference, BST is $O(h n + e n^2)$.

Dear reviewer Aowt, Thank you for the response. We have some updates to report:

Per your suggestion, we ran the FIM baseline on the Stargraph, and found that it performs poorly, picking random paths. We believe this is due to the next token prediction objective of FIM, which is known to fail on Stargraph. We added these findings to the Stargraph experimental section.

We also incorporated the discussion on training overhead to the appendix, which will be helpful for future readers.

We would be grateful for any additional comments or feedback before the discussion period closes.

Thank you for your review!

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?

As mentioned in Section 2.2 Belief State Inference, we always do left-to-right generation. When doing conditional planning, we score based on the next token head, and when doing unconditional, we use the previous token head.

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.

It seems like the reviewer would like to know where the performance improvements of the BST are coming from for the Stargraph task, and the reviewer thinks a FIM baseline would give more information on the performance gains of the BST.

However, this is already answered by the experiments. The Belief State Transformer uses standard left-to-right autoregressive generation during inference time, identical to the standard next token predictor models, and outperforms the standard next token predictor model. Hence, all the change in performance is due to the change in the objective, which is detailed in the ablation section.

On a related note, Reviewer Aowt also questioned how well FIM would do in the Stargraph setting. Because multiple reviewers have asked about FIM on the Stargraph setting, we will look into running FIM on the Stargraphs and aim to report results before the end of the discussion period.

Does not discuss whether a backward language model could solve the star-shaped graph task more effectively than a forward LM.

A language model that processes sequences from right to left could solve the Stargraph, and was in fact proposed by the Stargraph paper as a solution. However, such a solution is very task-specific and requires modification of both input data and model outputs (specifically, the user must reverse the training sequences and model outputs).

As a result, we limited our comparisons to more general baselines like teacherless (multi-token prediction), and data augmentation. The reviewer may be interested in [1] which shows that LLMs trained to predict right-to-left are generally worse than LLMs trained left-to-right in natural language tasks, which provides further evidence on the limitation of backward models. We will update the text to make it clear that the Stargraph can be easily solved using domain knowledge and that our motivation is to develop a general model that can solve hard tasks like Stargraph without domain knowledge.

[1] Papadopoulos, Vassilis, Jérémie Wenger, and Clément Hongler. "Arrows of Time for Large Language Models." arXiv preprint arXiv:2401.17505 (2024).

Models like BERT, trained with a Masked Language Modeling (MLM) objective, consider information from both direction… The paper does not compare Belief State Transformer with MLM models …Related 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.

Thanks for the suggestion, and we will update related work. Non-causal models like BERT could ostensibly handle the Stargraph dataset better than GPT-style transformers. However, as motivated in our intro, we are interested in examining the flaws of the dominant sequence modeling paradigm (GPT + next token prediction), and addressing them.

Our contribution, the Belief State Transformer, is a generalization of the GPT architecture, and inherits key desirable properties (i.e. causal attention, left to right decoding) while overcoming existing weaknesses in GPT. Thus, our main point of comparison is against prior GPT-style approaches.

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.

Whether or not BSTs can solve parity isn’t an interesting question. Our primary objective with the Stargraph experiment, is to present a simple and well-known problem in which GPT models empirically fail, and show that BST outperforms the prior work. The reduction between the Stargraph task with the parity problem is to give intuition behind why the Stargraph is hard to solve for next token predictors.

Dear reviewer Yi5e, Thank you again for your detailed and constructive comments. We’ve updated the draft and ran some more experiments, and would like to report in.

Per your suggestion, we ran the FIM baseline on the Stargraph, and found that it performs poorly, picking random paths. We believe this is due to the next token prediction objective of FIM, which is known to fail on Stargraph. We added these findings to figure 2 and section 3.4.

We’ve updated the text to emphasize that we use left-to-right ordering in all experiments, and that we use the next head for the conditional experiments, and previous head for unconditional experiments. See lines 116, 354, and 416.

Following your suggestion, we’ve updated the related work to mention non-causal approaches like ELMo and BERT, see Section 6.

We believe that we’ve sufficiently answered all of your questions, but would be happy to engage further, and would be grateful for any other feedback.

Thanks for the review!

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.

We respectfully disagree with the point that BST seems designed to solve Stargraph here. BST is a general objective for sequence modeling motivated by belief state discovery, and this learning of belief states allows it to solve tasks that require lookahead, like the Stargraph.

Conceptually, the BST objective is radically different when compared to the right-to-left solution. The right-to-left solution only uses the suffix to predict the previous token, which does not guarantee the discovery of belief states. For BST, both prefix and suffix embeddings are used to predict the previous token (as well as the next token), and this distinction is crucial for a belief state to arise (see Theorem 2).

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

The right-to-left solution is a task specific solution engineered for StarGraph, using domain knowledge for both training and inference. First, the solution exploits the structure of the stargraph by reversing the path data. Next, after training, the user must manually reverse the outputs to get the correct solution.

In contrast, BST’s training is task-agnostic - we can train the BST without altering the data, and perform the standard left-to-right decoding to get the answer. So while the BST and the right-to-left solution both solve the Stargraph and share a similar element of a previous token prediction term, the BST solution is more general and provably discovers belief states.

In order to claim that the proposed training objective results in a better belief state representation in general, would require more evaluation

We provided formal proof that the Belief State Transformer training process converges to a compact belief state, and that methods like next-token and multi-token prediction do not. Could the reviewer point out a specific point of concern in the proof? Additionally, we empirically verified on stargraph that the belief state transformer perfectly learns the correct belief state while the next-token transformer does not.

Showing the effectiveness on planning problems from benchmarks like PlanBench [1] and ACPBench [2] would make the claim of the paper stronger.

Showing BST’s effectiveness on more domains is valuable, as we have written in our limitations. However, showing a full-scale experiment before the end of the rebuttal period is difficult, and we kindly ask the reviewer to look at our current contributions:

We proposed a new form of training for language models motivated by belief state discovery, showed that it overcomes hard planning problems known to confound standard approaches, provided theoretical proofs, and demonstrated good results in text generation.

Our current contributions serve as proof-of-concept, and will hopefully inspire others to investigate the work further. The reviewer may be interested in our preliminary results in planning in a maze navigation setting, see below.

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?

We began experiments (after submission) of the BST on a 4-room maze navigation task, and found some promising results where BST greatly outperforms FIM in navigating to different goal locations. See below for some more details.

We use the four-room maze task. In the four-room maze, the agent must navigate in a maze composed of four rooms interconnected by 4 gaps in the walls. We collect an offline dataset of 100K transitions using a uniform random data collecting policy, and train all models for 50 epochs.

We evaluate the models on reaching goals of increasing difficulty. Type-0 goals are cells in the same room, Type-1 goals are cells in adjacent rooms, and Type-2 goals are cells in the most distant room. We compare BST with the FIM baseline on the mean success rate, results are listed as follows.

For the StarGraph experiments, how was the backward latent B(∅) computed?

We compute the backward latent by passing in the End-Of-Sentence (EOS) token, which is appended to every sequence in the stargraph training data.

How would the performance change if the training objective had B but during inference B(∅) was not used?

The text head is conditioned on both forward and backward latents, so it is not possible to do inference without B(∅).

In the ablation section, what is the difference between Belief w/o Backward and single transformer next token prediction setup?

The Belief w/o Backward still predicts next and previous tokens.

Dear reviewer QZxB, We thank you again for your feedback. As the open discussion period draws to a close, we’d like to check back to see if you have any remaining concerns.

We believe that we’ve sufficiently addressed your concerns and questions, but would be happy to clarify further. We would also be grateful for any other feedback.