Title: Higher Embedding Dimension Creates a Stronger World Model for a Simple Sorting Task
Primary area: interpretability and explainable AI
Keywords: embedding dimension, world model, transformer, reinforcement learning, sorting, PPO, interpretability, interpretation of learned representations

Abstract: We study 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. While even very small embedding dimensions are sufficient for models to achieve high accuracy, larger dimensions yield representations that are more faithful, consistent, and robust. In particular, higher embedding dimensions strengthen the formation of structured internal representation and leads 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 metrics and analyses that can be reused to probe similar tasks.

TL;DR: We study 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.