mlsanityData to 14 Oct 2026

Higher Embedding Dimension Creates a Stronger World Model for a Simple Sorting Task

arXiv 2510.18315 ↗First posted 21 Oct 2025cs.LG cs.AIcited by 0

On the map: World models → World models for language agents

also in 1 more
Abstract (arXiv)

We investigate how embedding dimension affects the emergence of an internal "world model" in a transformer trained with reinforcement learning to perform bubble-sort-style adjacent swaps. Models achieve high accuracy even with very small embedding dimensions, but larger dimensions yield more faithful, consistent, and robust internal representations. In particular, higher embedding dimensions strengthen the formation of structured internal representation and lead to better interpretability. After hundreds of experiments, we observe two consistent mechanisms: (1) the last row of the attention weight matrix monotonically encodes the global ordering of tokens; and (2) the selected transposition aligns with the largest adjacent difference of these encoded values. Our results provide quantitative evidence that transformers build structured internal world models and that model size improves representation quality in addition to end performance. We release our metrics and analyses, which can be used to probe similar algorithmic tasks.

Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in World models for language agents, one of 11 directions in World models. Builds on 18 of its 21 references that are in the map. First posted on arXiv 21 Oct 2025.

Builds on

18 of 21 references.

Closest ancestors · 2

Same direction, newest first

  1. 2023Emergent Linear Representations in World Models of Self-Supervised Sequence Models2 Sep
  2. 2023Reasoning with Language Model is Planning with World Model24 May

Other references · 16

  1. 2017Attention Is All You Needcited by 60,361
  2. 2017Proximal Policy Optimization Algorithmscited by 15,215
  3. 2022Chain-of-Thought Prompting Elicits Reasoning in Large Language Modelscited by 12,459
  4. 2021Training Verifiers to Solve Math Word Problemscited by 10,109
  5. 2013Efficient Estimation of Word Representations in Vector Spacecited by 4,375
  6. 2018World Modelscited by 1,513
  7. 2022In-context Learning and Induction Headscited by 964
  8. 2022Transformers are Sample-Efficient World Modelscited by 328
  9. 2022Towards Understanding Grokking: An Effective Theory of Representation Learningcited by 230
  10. 2023A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operationscited by 144
  11. 2023STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learningcited by 89
  12. 2018On the Dimensionality of Word Embeddingcited by 46
4 more
  1. 2023Can Transformers Learn to Solve Problems Recursively?cited by 17
  2. 2016The Role of Context Types and Dimensionality in Learning Word Embeddingscited by 4
  3. 2025SortBench: Benchmarking LLMs based on their ability to sort listscited by 4
  4. 2023ConvD: Attention Enhanced Dynamic Convolutional Embeddings for Knowledge Graph Completioncited by 2

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2024Transformers Use Causal World Models in Maze-Solving Tasks
  2. 2026Transformers Linearly Represent Highly Structured World Models
  3. 2023Emergence and Function of Abstract Representations in Self-Supervised Transformers
  4. 2026MetaOthello: A Controlled Study of Multiple World Models in Transformers
  5. 2026Local Intrinsic Dimension of Representations Predicts Alignment and Generalization in AI Models and Human Brain
  6. 2023Linear Latent World Models in Simple Transformers: A Case Study on Othello-GPT
  7. 2024Let's Think Dot by Dot: Hidden Computation in Transformer Language Models
  8. 2024Learning interpretable positional encodings in transformers depends on initialization