更大嵌入维度让Transformer更精准构建排序任务的世界模型。
Higher Embedding Dimension Creates a Stronger World Model for a Simple Sorting Task
- 用强化学习训练Transformer完成相邻交换排序,研究嵌入维度影响。
- 高维嵌入使内部表示更结构化、一致且可解释,准确率超95%。
- 适合关注模型内部机制与可解释性的研究人员参考。
我们研究了嵌入维度对基于强化学习训练的Transformer在执行冒泡排序式相邻交换任务时,内部“世界模型”涌现的影响。即使嵌入维度极小,模型仍能实现高准确率,但更大的维度带来更忠实、一致且鲁棒的内部表征。特别是,高维嵌入增强了结构化内部表示的形成,提升了可解释性。经过数百次实验,我们观察到两个稳定机制:(1) 注意力权重矩阵的最后一行单调编码标记的全局顺序;(2) 选择的交换操作与这些编码值中最大的相邻差值对齐。结果为Transformer构建结构化内部世界模型提供了量化证据,并表明模型规模不仅提升最终性能,也改善表征质量。我们公开了相关度量与分析工具,可用于探测类似算法任务。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。