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

arXiv: https://arxiv.org/abs/2510.18315 · HTML: https://mlsanity.com/paper/2510.18315/

## 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.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 18 of 18)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [Reasoning with Language Model is Planning with World Model](https://mlsanity.com/paper/2305.14992/) (2023)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Emergent Linear Representations in World Models of Self-Supervised Sequence Models](https://mlsanity.com/paper/2309.00941/) (2023)
- [Towards Understanding Grokking: An Effective Theory of Representation Learning](https://mlsanity.com/paper/2205.10343/) (2022)
- [A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operations](https://mlsanity.com/paper/2302.03025/) (2023)
- [STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning](https://mlsanity.com/paper/2310.09615/) (2023)
- [On the Dimensionality of Word Embedding](https://mlsanity.com/paper/1812.04224/) (2018)
- [Can Transformers Learn to Solve Problems Recursively?](https://mlsanity.com/paper/2305.14699/) (2023)
- [The Role of Context Types and Dimensionality in Learning Word Embeddings](https://mlsanity.com/paper/1601.00893/) (2016)
- [SortBench: Benchmarking LLMs based on their ability to sort lists](https://mlsanity.com/paper/2504.08312/) (2025)
- [ConvD: Attention Enhanced Dynamic Convolutional Embeddings for Knowledge Graph Completion](https://mlsanity.com/paper/2312.07589/) (2023)

## 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.