Larger Embedding Dimensions Strengthen Transformer "World Models" in Sorting Tasks

Researchers investigating how transformers build internal understanding of algorithmic tasks have found that embedding dimension—a core architectural parameter—plays a crucial role in the quality of learned representations, even when task performance plateaus.

How the Study Works#

The work examined transformer models trained via reinforcement learning to perform bubble-sort operations: arranging elements by swapping adjacent items. Although models achieved high accuracy across a range of embedding dimensions, the researchers discovered a striking pattern in how these internal structures differ.

Key Findings#

Two consistent mechanisms emerged across hundreds of experiments:

• The last row of the attention weight matrix monotonically encodes the global ordering of tokens, creating a structured representation of where each element belongs in the sequence.
• Selected transpositions align with the largest adjacent difference in these encoded values, allowing the model to identify which pairs need swapping.

Larger embedding dimensions strengthened both the formation and consistency of these structured representations. The researchers noted that higher dimensions lead to more interpretable internal models and improved representation quality beyond what end-task accuracy alone would suggest.

Why This Matters#

The findings provide quantitative evidence that transformers actively construct "world models"—coherent, internal understandings of problems—rather than simply learning surface-level statistical patterns. This suggests that model capacity improvements benefit not just raw performance but the reliability and robustness of the learned solution. The released metrics and analysis tools can probe similar algorithmic tasks in future work.

View the full paper on arXiv