Saved in:
| Main Authors: | , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2510.18315 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866912662677356544 |
|---|---|
| author | Bhalla, Brady Fan, Honglu Chen, Nancy YU, Tony Yue |
| author_facet | Bhalla, Brady Fan, Honglu Chen, Nancy YU, Tony Yue |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2510_18315 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Higher Embedding Dimension Creates a Stronger World Model for a Simple Sorting Task Bhalla, Brady Fan, Honglu Chen, Nancy YU, Tony Yue Machine Learning Artificial Intelligence 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. |
| title | Higher Embedding Dimension Creates a Stronger World Model for a Simple Sorting Task |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2510.18315 |