Research map / Architectures and efficiency / Attention and transformers
The Belief State Transformer
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes the Belief State Transformer, a jointly trained forward and backward encoder model that predicts the next prefix token and previous suffix token, and shows gains on star graph planning and TinyStories story generation.
Problem
Forward-only transformers struggle with planning-heavy, goal-conditioned tasks (e.g., star graph path finding) and with goal-conditioned story infilling because next-token training yields flawed shortcuts and non-compact belief states.
Approach
Train two encoders (forward over prefix, backward over suffix) with two output heads that jointly predict the next token after the prefix and the previous token before the suffix over all prefix-suffix pairs; theory shows the forward encoding becomes a belief state, and at inference the backward encoder is replaced by an empty-suffix embedding.
Key results
- On TinyStories goal-conditioned generation, GPT-4 prefers the Belief State Transformer over FIM 64% of the time (CI 54%-74%).
- On star graphs, the Belief State Transformer solves all graph settings while forward-only, data augmentation, FIM baselines stay near 1/d success; multi-token fails on G(5,5) and G(2,20).
- A shared-encoder improved BST with lambda=0.9 matches GPT next-token loss (0.926 vs 0.929) while training about 0.801 vs 0.205 iterations/s for the original BST.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use the improved BST configuration with shared encoders and subsampled training pairs | BST (original) · TinyStories | next token loss: 1.013 -> 0.926 | Helped |
| Use the improved BST configuration with shared encoders and subsampled training pairs | BST (original) · TinyStories | previous token loss: 1.110 -> 1.268 | Hurt |
| Use shared encoders and subsample 2% of prefix-suffix pairs | BST · TinyStories | training speed (iterations/second): $\sim0.205$ -> $\sim0.801$ | Helped |
| Use Belief State Transformer instead of Fill-in-the-Middle | TinyStories goal-conditioned story generation · TinyStories | GPT4 preference: 64\% preference; confidence interval $\left[54\%,74\%\right]$ | Helped |
| Use Belief State Transformer instead of Fill-in-the-Middle | TinyStories unconditional text generation · TinyStories | GPT4 preference | Helped |
| Remove the previous-token decoder and its objective (ablation) | Belief State Transformer · Star graph navigation | path accuracy | Hurt |
| Remove the backward encoder (ablation) | Belief State Transformer · Star graph navigation | path accuracy | Hurt |
| Use beam search instead of the priority queue scheme (ablation) | Belief State Transformer goal-conditioned generation · TinyStories | story quality | Hurt |
| Use the next head instead of the previous head to score the suffix (ablation) | Belief State Transformer unconditional generation · TinyStories | story quality | Hurt |
| Subsample prefix-suffix pairs during training | Belief State Transformer · TinyStories | final performance | No clear effect |
Surprising findings
- The forward-only and data augmentation baselines perform poorly on star graphs even with extra subgoal or suffix information, and the original multi-token baseline fails on harder graphs.
- Discarding the backward encoder at inference still yields goal-conditioned behavior, since the goal is present in the prefix.
Limitations
- Experiments are small scale (80M parameters, TinyStories with vocabulary 1000) and the authors state scaling is untested.
- Evaluation of story quality relies on GPT-4 as judge with winrates only, on 100 stories.
- Training on all O(T^2) prefix-suffix pairs is expensive; practical variant requires subsampling and changes the architecture.
Experimental setup
Datasets: TinyStories Star graph synthetic data (N=50)
Benchmarks: Star graph path accuracy G(2,5), G(5,5), G(2,20) TinyStories pairwise winrate
Metrics: path accuracy winrate binomial confidence interval probe accuracy log loss
Baselines: Forward-only next-token prediction Data augmentation Fill-in-the-middle Multi-token prediction Belief w/o Prev Belief w/o Backward
Builds on: GPT-2 Fill-in-the-Middle Next-Token-Failures star graph setup BERT
Closest papers by meaning
- Next-Latent Prediction Transformers Learn Compact World Models (NeurIPS 2026 · similarity 0.85)
- Distilling Sequential Computation in Transformer Language Models (NeurIPS 2026 · similarity 0.85)
- Transformers Represent Belief State Geometry in their Residual Stream (NeurIPS 2024 · similarity 0.85)
- Effective Context in Transformers: An Analysis of Fragmentation and Tokenization (NeurIPS 2026 · similarity 0.85)
- Hierarchical Autoregressive Transformers: Combining Byte- and Word-Level Processing for Robust, Adaptable Language Models (ICLR 2025 · similarity 0.85)
- Learning and Transferring Sparse Contextual Bigrams with Linear Transformers (NeurIPS 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.