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

TL;DR

This work investigates how embedding dimension shapes the emergence of an internal world-model in a tiny transformer trained via reinforcement learning to perform adjacent swaps toward sorting. By combining a minimal single-head transformer with PPO in a controlled bubble-sort task, the authors show that larger embedding dimensions yield more faithful and robust internal representations, manifested in two consistent mechanisms: the last attention-row encodes the global token order and the chosen swap aligns with the largest adjacent difference. Accuracy saturates at low embedding dimensions, but the fidelity of the internal circuit and its interpretability improve with dimension up to about 30, indicating that capacity primarily enhances representation quality rather than end performance. These results provide quantitative evidence that transformers build structured world-model-like representations in constrained RL settings and offer metrics and methods to probe similar algorithmic tasks.

Abstract

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.

Figures (8)

• Figure 1: Accuracy vs. embedding dimension for length 6 and length 8 agents. Almost all agents achieve 100% accuracy for length 6, but this is not the case for length 8, where agents either achieve very high accuracy or get stuck at much lower accuracy. The mean and 95% confidence intervals are shown.
• Figure 2: Proportion of non-inversions vs. embedding dimension for length 6 and length 8 sequences where the model has near-perfect accuracy. This reaches over 80% on average for both sequences when the embedding dimension is large enough. The mean and 95% confidence intervals are shown.
• Figure 3: Proportion of moves which are the top 1/2 prediction according to our observation vs. embedding dimension for length 6 and length 8 sequences where the model has near-perfect accuracy. Almost all moves are chosen according to this observation, with the top 2 proportion reaching above 90% for a large enough embedding dimension. The mean and 95% confidence intervals are shown.
• Figure 4: The proportion of top 1 predictions vs. proportion of non-inversions for all agents. These two metrics have an $r^2$ value of 0.56, so there is a fairly strong correlation between the two values. The line of best fit and 95% confidence interval for the parameters of this line are shown.
• Figure 5: Visualization of the attention weights for a fully ordered permutation (using the same agent as in Figure \ref{['best-violin']}). The weights in the last row are ordered just like the permutation, showing that for this agent and permutation there are no inversions.