Pith. sign in

REVIEW 3 major objections 9 minor 25 references

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

T0 review · 3 major / 9 minor · reviewed 2026-08-04 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read Larger embedding dimensions make a sorting transformer's internal ordering representation more faithful, even after task accuracy saturates.

desk verdict The empirical trend is real and worth discussing, but the mechanism claim is correlational — referee for the observation, ask for ablations. read the letter →

arxiv 2510.18315 v2 pith:PYUDPFWA submitted 2025-10-21 cs.LG cs.AI

Brady Bhalla , Honglu Fan , Nancy Chen , Tony Yue YU This is my paper
classification cs.LGcs.AI
keywords mechanisticinterpretabilityworldmodelsembeddingdimensionsortingreinforcementlearningattentionweightsrepresentationfidelitybubblesort
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper asks whether a minimal transformer trained with reinforcement learning to sort permutations by adjacent swaps builds an internal "world model," and whether that model sharpens as embedding dimension grows. It claims that even though task accuracy saturates at small embedding sizes, larger embeddings make two mechanisms increasingly consistent: the last row of the attention matrix comes to encode the global ordering of the tokens, and the agent's chosen swap matches the largest adjacent gap in those encoded values. The paper takes this as quantitative evidence that transformers form structured internal representations of environment state, and that capacity buys representation quality even when it no longer buys accuracy. A sympathetic reader would care because it suggests width alone can improve interpretability and robustness without hurting performance.

What carries the argument

The load-bearing object is the final row W_ℓ of the single-head causal attention matrix W = QK^T/√d_k. Because the model is causal, this last row can attend to every token, making it the natural readout of global sequence state. The paper defines two metrics on that object: the proportion of non-inversions between the token order and the order of the last-row attention weights, and the frequency with which the chosen swap ranks among the top gaps between consecutive attention-output values. The monotonic ordering of W_ℓ plus the largest-difference rule together form the claimed "simple circuit" for sorting.

What would settle it

Find or train a 100%-accurate agent whose last attention row is not monotonic (non-inversion proportion near 50%) and whose chosen swaps do not match the largest adjacent gaps (top-1 match below chance); or, alternatively, ablate the last row of the attention matrix and observe no change in swap accuracy. Either observation would show the claimed ordering circuit is not causally responsible.

Watch

Extended reading notes

Core claim

The paper's central claim is that a consistent, interpretable circuit underlies the sorting behavior. Across hundreds of agents, the final row of the attention weight matrix assigns each token a value, with lower values for smaller numbers, so reading that row recovers the global order of the permutation. The value matrix and output layer then compute consecutive differences, and the selected adjacent swap is the one with the largest difference. For accurate agents, 76-77% of moves match the top-1 largest gap and over 90% match top-2. Larger embedding dimensions do not change the circuit qualitatively but make it more reliable: both the ordering metric and the swap-rule match rise with embed

Load-bearing premise

The claim stands or falls on the assumption that the last-row attention weights and the largest-difference swap rule are the circuit actually driving decisions, rather than byproducts of training that correlate with accuracy.

Editorial extensions

If this is right

  • If the last-row monotonic encoding is genuine, a model's swap choices can be read off directly from its attention outputs: picking the largest adjacent difference in that row predicts behavior over 90% of the time at top-2.
  • Because representation metrics keep rising after accuracy plateaus, performance alone understates what larger capacity does; internal-structure metrics are needed to evaluate model quality.
  • The ordering is learned without ever being given token numerical values, so the circuit is an emergent inference about global order from the reward signal alone.
  • The strong cross-metric correlation (r² = 0.56) means either the ordering metric or the swap-rule metric can serve as a proxy for the other in similar studies.
  • Overprovisioning width may be a low-cost way to make learned policies more interpretable and less brittle, with no loss of task accuracy.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • Because the evidence is correlational, the causal claim that this circuit drives sorting remains open; an ablation of the last attention row or the value projection would close that gap.
  • If the capacity–representation link is general, the same probe could be applied to other discrete tasks — graph problems, dynamic programming — where a last-row ordering might expose similar structure.
  • The dimension threshold near 30 may depend on sequence length; testing length-10 or length-12 sequences would show whether the threshold shifts with state-space size.
  • The "local greedy trap" suggests a testable symptom: low-embedding failures should correspond to near-random last-row ordering, and initializing embeddings monotonically should improve accuracy.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 9 minor

Summary. The paper studies how embedding dimension affects the internal representations of a minimal, single-head, decoder-only transformer trained with PPO to perform adjacent swaps that sort permutations of length 6 and 8. Across 475 runs with dimensions 2–128, the authors report that sorting accuracy saturates at small embedding dimensions, while two internal metrics keep improving up to a dimension of roughly 30: (1) the last row of the attention weight matrix increasingly encodes the global order of tokens, measured by a non-inversion proportion, and (2) the chosen swap increasingly matches the largest adjacent gap in the corresponding attention-output values, measured by top-1/top-2 agreement. The paper interprets these as evidence that larger embedding dimensions create a stronger, more faithful 'world model' of the sorting task, and that this difference is invisible in task accuracy alone.

Significance. If the causal interpretation held, the paper would provide a clean demonstration that model capacity affects internal representation quality beyond the performance plateau—an important result for mechanistic interpretability and for thinking about overparameterization. The study has clear strengths: a large number of trained agents, well-defined quantitative metrics, public release of metrics/analyses, and two specific, falsifiable observations about attention structure and decision rule. The main weakness is that the central 'circuit' and 'world model' claims are supported only by correlational alignment metrics; no intervention or ablation is performed. This is a load-bearing gap for the title's causal language. The aggregate empirical trends (accuracy saturation versus later representation improvement) are credible and useful even if the causal interpretation is weakened.

major comments (3)
  1. [Section 5, Observations 1–2 and §6.1] The paper labels the two observations as 'mechanisms' and 'a simple circuit underlying sorting behavior', but the evidence is entirely correlational. §4.3 and §5.2–5.3 measure alignment between attention weights/attention-output gaps and the agent's chosen action, but no experiment ablates, perturbs, or counterfactually modifies the last attention row, the value matrix, or the output linear layer. Thus a third factor—e.g., convergence quality or training dynamics—could drive both high task accuracy and ordered attention weights without the proposed circuit being causally responsible. To support the causal 'world model' claim, the authors should perform intervention experiments: e.g., replace or permute the last-row attention weights according to a different ordering and check that the chosen swap changes as predicted by the largest-gap rule, or clamp the attention row to uniform to see w
  2. [Section 4.2, §4.1, Table 1] Training timesteps vary between runs (1M, 2M, or 10M), and the paper does not report how these were allocated across embedding dimensions in the dimension sweeps that produce Figures 1–3. If larger-dimension runs received more timesteps, the observed trend of improving representation with dimension could partly reflect training budget rather than embedding dimension. This is load-bearing for the title claim that higher embedding dimension 'creates' a stronger world model. Please present the dimension sweeps disaggregated by timestep budget, or verify that the allocation was balanced/randomized across dimensions; if it was not balanced, the manuscript needs to state this and temper the causal language.
  3. [Section 5.2, length-8 analysis] For length 8, the cross-dimension comparison is performed after restricting to agents with 'near-perfect accuracy'. Since convergence to 100% accuracy is itself dimension-dependent (§5.1 reports only 37.4% of length-8 agents above dim 16 achieve 100%), this selection can induce a spurious relationship between embedding dimension and the non-inversion metric: one may be comparing the best high-dimension seeds with a non-representative subset of low-dimension seeds. The authors should show that the trend within the high-accuracy subset is not an artifact of selection, for example by reporting results for all agents, or by matching agents on achieved accuracy, or by applying a fixed accuracy threshold rather than the 'near-perfect' criterion. As written, the length-8 claim is not robustly supported.
minor comments (9)
  1. [§4.3, metric 2] The definition of 'non-inversions between (π(1),...,π(n)) and (Wℓ,1,...,Wℓ,n)' is informal. Please define the inversion count for a real-valued sequence against a permutation, and state how ties in Wℓ are handled.
  2. [§4.3, metric 3 and §5.3] The metric sorts 'differences between consecutive attention output values' and computes top-k agreement, but §5.3 says the agent selects the 'most positive or most negative difference'. Specify whether the ranking is by absolute difference or signed difference, and what 'top-1' means when both a large positive and a large negative gap exist.
  3. [§5.1] The phrase '99.2% of the agents reliably achieved 100% accuracy' is ambiguous: does it mean that 99.2% of trained agents reached perfect accuracy, or that the average accuracy was 99.2%? Clarify.
  4. [Table 1] The N/A entries and non-overlapping timestep choices between length 6 and length 8 are unexplained. State explicitly which (length, timestep) configurations were run.
  5. [§4.2] Reference [7] is cited for the PPO clipping mechanism, but [7] is a knowledge-graph embedding paper; the correct citation appears to be [20] (Schulman et al.). Please fix the citation.
  6. [§5.5] Use 'R²' or 'R^2' consistently rather than 'r2'; also, the correlation between the two metrics means Observation 2 is not fully independent evidence for Observation 1—this is acknowledged but should be discussed explicitly when claiming two separate mechanisms.
  7. [Appendix A.3] The claim that 'Figure 7 has a much wider spread of values for each token' is qualitative; consider reporting standard deviations or interquartile ranges for the violin plots.
  8. [Abstract/§5] The term 'transposition' usually denotes swapping any two elements, not necessarily adjacent; use 'adjacent swap' or 'adjacent transposition' throughout for precision.
  9. [§4.2 / Data availability] The paper says 'We release our metrics and analyses' but provides no link or repository identifier in the manuscript. Please add the URL or state that it will be provided in the camera-ready.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity: the observed metrics are direct measurements from model internals, not fitted inputs relabeled as predictions, and no load-bearing claim reduces to a self-citation.

full rationale

The paper's central empirical chain is: train PPO agents; measure sorting accuracy; measure whether the last attention row is monotonically aligned with token order (metric 2); and measure whether chosen swaps match largest attention-output gaps (metric 3); then relate these measures to embedding dimension. Each metric is defined directly from the model's attention weights and behavior (§4.3), with no free parameter fitted to the target outcome and then reported as a prediction. There is no equation in which a derived quantity equals an input by construction. The acknowledged relationship between Observations 1 and 2 (§5.5, r²=0.56) is a redundancy between two correlated measurements, and the paper explicitly notes they are 'very related'; this reduces independent-evidence strength but is not a circular derivation. The absence of ablations or interventions is a causal-evidence gap about whether the observed correlations are the operative circuit, i.e., a correctness or interpretability risk, not a circularity. No self-citations are load-bearing; the paper does not invoke a uniqueness theorem or import an ansatz from the authors' prior work. The §6.3 limitations concern scope and generalization, not a redefinition of inputs. Accordingly, no circular step meeting the quoted-evidence standard can be exhibited, and the honest finding is score 0.

Assumptions & free parameters 0 free parameters · 3 assumptions · 0 invented entities

No fitted constants appear in the metrics or conclusions; the embedding-dimension '~30' saturation point is an empirical observation from figures rather than a fitted parameter. The central reliance is on domain assumptions about the task design, metric validity, and the decision to analyze only successful agents.

assumptions (3)
  • domain assumption A single-head, single-layer, causal, MLP-free transformer can learn to execute adjacent-swap sorting from PPO rewards, and its final attention row is a faithful readout of token order.
    Invoked throughout §4 and §5; if the architecture or task cannot support this, the ordering circuit is an artifact rather than a learned world model. Location: §4.2, §5.2.
  • domain assumption The reward function (+1 when sorted, -0.001 otherwise) plus PPO causes successful agents to converge to a near-deterministic swap policy, so failed length-8 agents can be excluded without biasing the representation claim.
    Used when restricting analysis to near-perfect agents in §5.2 and Table 2. If exclusion selects on the outcome, the dimension trend is confounded with convergence ability.
  • domain assumption Monotonicity of the last attention row relative to input tokens is a valid operationalization of a 'world model' / internal state representation.
    The paper's central construct is defined by this metric (§5.2); the term 'world model' is not independently validated by behavioral prediction or counterfactual tests.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Higher Embedding Dimension Creates a Stronger World Model for a Simple Sorting Task." pith.science (2026). https://pith.science/paper/PYUDPFWA

@misc{pith2026251018315,
  author       = {Pith},
  title        = {Pith review of: Higher Embedding Dimension Creates a Stronger World Model for a Simple Sorting Task},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/PYUDPFWA}},
  note         = {Machine review of arXiv:2510.18315}
}
read the original 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

Figures reproduced from arXiv: 2510.18315 by the authors.

Figure 1
Figure 1. Accuracy vs. embedding dimension for length 6 and length 8 agents. Almost all agents [PITH_FULL_IMAGE:figures/full_fig_p006_1.png] view at source ↗
Figure 2
Figure 2. Proportion of non-inversions vs. embedding dimension for length 6 and length 8 sequences [PITH_FULL_IMAGE:figures/full_fig_p006_2.png] view at source ↗
Figure 3
Figure 3. Proportion of moves which are the top 1/2 prediction according to our observation vs. [PITH_FULL_IMAGE:figures/full_fig_p008_3.png] view at source ↗
Figures from the paper (5 more)
Figure 4
Figure 4. Figure 4: The proportion of top 1 predictions vs. proportion of non-inversions for all agents. These [PITH_FULL_IMAGE:figures/full_fig_p008_4.png]
Figure 5
Figure 5. Figure 5: Visualization of the attention weights for a fully ordered permutation (using the same [PITH_FULL_IMAGE:figures/full_fig_p013_5.png]
Figure 6
Figure 6. Figure 6: Attention weights by token for the agent with fewest average inversions. [PITH_FULL_IMAGE:figures/full_fig_p013_6.png]
Figure 7
Figure 7. Figure 7: Attention weights by token for a random agent with high accuracy. [PITH_FULL_IMAGE:figures/full_fig_p014_7.png]
Figure 8
Figure 8. Figure 8: Clipped value loss vs. global step in an example training of an agent on a length 6 sequence. [PITH_FULL_IMAGE:figures/full_fig_p014_8.png]

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

25 extracted references · 12 linked inside Pith

  1. [1]

    Quality over quantity in attention layers: When adding more heads hurts

    Noah Amsel, Gilad Yehudai, and Joan Bruna. Quality over quantity in attention layers: When adding more heads hurts. InThe Thirteenth International Conference on Learning Representa- tions, 2024

    2024

  2. [2]

    Openai gym.arXiv preprint arXiv:1606.01540, 2016

    Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym.arXiv preprint arXiv:1606.01540, 2016

    arXiv 2016

  3. [3]

    A toy model of universality: Reverse engineering how networks learn group operations

    Bilal Chughtai, Lawrence Chan, and Neel Nanda. A toy model of universality: Reverse engineering how networks learn group operations. InInternational Conference on Machine Learning, pages 6243–6267. PMLR, 2023

    2023

  4. [4]

    Training Verifiers to Solve Math Word Problems.CoRR, abs/2110.14168, 2021

    Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training Verifiers to Solve Math Word Problems.CoRR, abs/2110.14168, 2021

    arXiv 2021

  5. [5]

    MIT press, 2022

    Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein.Introduction to algorithms. MIT press, 2022

    2022

  6. [6]

    A mathematical framework for transformer circuits.Transformer Circuits Thread, 1(1):12, 2021

    Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits.Transformer Circuits Thread, 1(1):12, 2021

    2021

  7. [7]

    Convd: Attention enhanced dynamic convolutional embeddings for knowledge graph completion.IEEE Transactions on Knowledge and Data Engineering, 2025

    Wenbin Guo, Zhao Li, Xin Wang, Zirui Chen, Jun Zhao, Jianxin Li, and Ye Yuan. Convd: Attention enhanced dynamic convolutional embeddings for knowledge graph completion.IEEE Transactions on Knowledge and Data Engineering, 2025

    2025

  8. [8]

    World models.arXiv preprint arXiv:1803.10122, 2(3), 2018

    David Ha and Jürgen Schmidhuber. World models.arXiv preprint arXiv:1803.10122, 2(3), 2018

    arXiv 2018

Show all 25 references
  1. [9]

    Reasoning with language model is planning with world model.arXiv preprint arXiv:2305.14992, 2023

    Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, and Zhit- ing Hu. Reasoning with language model is planning with world model.arXiv preprint arXiv:2305.14992, 2023

    2023 arXiv

  2. [10]

    Sortbench: Benchmarking llms based on their ability to sort lists.arXiv preprint arXiv:2504.08312, 2025

    Steffen Herbold. Sortbench: Benchmarking llms based on their ability to sort lists.arXiv preprint arXiv:2504.08312, 2025

    2025 arXiv

  3. [11]

    Scale matters: Large language models with billions (rather than millions) of parameters better match neural representations of natural language.BioRxiv, 2024

    Zhuoqiao Hong, Haocheng Wang, Zaid Zada, Harshvardhan Gazula, David Turner, Bobbi Aubrey, Leonard Niekerken, Werner Doyle, Sasha Devore, Patricia Dugan, et al. Scale matters: Large language models with billions (rather than millions) of parameters better match neural represent...

    2024

  4. [12]

    Making large language models a better foundation for dense retrieval.CoRR, abs/2312.15503, 2023

    Chaofan Li, Zheng Liu, Shitao Xiao, and Yingxia Shao. Making large language models a better foundation for dense retrieval.CoRR, abs/2312.15503, 2023

    2023

  5. [13]

    Towards understanding grokking: An effective theory of representation learning.Advances in Neural Information Processing Systems, 35:34651–34663, 2022

    Ziming Liu, Ouail Kitouni, Niklas S Nolte, Eric Michaud, Max Tegmark, and Mike Williams. Towards understanding grokking: An effective theory of representation learning.Advances in Neural Information Processing Systems, 35:34651–34663, 2022

    2022

  6. [14]

    Daniel J. Mankowitz, Andrea Michi, Anton Zhernov, Marco Gelmi, Marco Selvi, Cosmin Paduraru, Edouard Leurent, Shariq Iqbal, Jean-Baptiste Lespiau, Alex Ahern, Thomas Köppe, Kevin Millikin, Stephen Gaffney, Sophie Elster, Jackson Broshear, Chris Gamble, Kieran Milan, Robert Tun...

    2023

  7. [15]

    The role of context types and dimensionality in learning word embeddings.arXiv preprint arXiv:1601.00893, 2016

    Oren Melamud, David McClosky, Siddharth Patwardhan, and Mohit Bansal. The role of context types and dimensionality in learning word embeddings.arXiv preprint arXiv:1601.00893, 2016

    2016 arXiv

  8. [16]

    Transformers are sample-efficient world models.arXiv preprint arXiv:2209.00588, 2022

    Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world models.arXiv preprint arXiv:2209.00588, 2022. 10

    2022 arXiv

  9. [17]

    Efficient estimation of word representations in vector space.arXiv preprint arXiv:1301.3781, 2013

    Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space.arXiv preprint arXiv:1301.3781, 2013

    2013 arXiv

  10. [18]

    Emergent linear representations in world models of self-supervised sequence models.arXiv preprint arXiv:2309.00941, 2023

    Neel Nanda, Andrew Lee, and Martin Wattenberg. Emergent linear representations in world models of self-supervised sequence models.arXiv preprint arXiv:2309.00941, 2023

    2023 arXiv

  11. [19]

    In-context learning and induction heads.arXiv preprint arXiv:2209.11895, 2022

    Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads.arXiv preprint arXiv:2209.11895, 2022

    2022 arXiv

  12. [20]

    Proximal policy optimization algorithms.arXiv preprint arXiv:1707.06347, 2017

    John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms.arXiv preprint arXiv:1707.06347, 2017

    2017 arXiv

  13. [21]

    Gomez, Lukasz Kaiser, and Illia Polosukhin

    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V . N. Vishwanathan, and Roman Garnett,...

    2017

  14. [22]

    Chain-of-thought prompting elicits reasoning in large language models

    Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837, 2022

    2022

  15. [23]

    On the dimensionality of word embedding.Advances in neural information processing systems, 31, 2018

    Zi Yin and Yuanyuan Shen. On the dimensionality of word embedding.Advances in neural information processing systems, 31, 2018

    2018

  16. [24]

    Can transformers learn to solve problems recursively?arXiv preprint arXiv:2305.14699, 2023

    Shizhuo Dylan Zhang, Curt Tigges, Stella Biderman, Maxim Raginsky, and Talia Ringer. Can transformers learn to solve problems recursively?arXiv preprint arXiv:2305.14699, 2023

    2023 arXiv

  17. [25]

    Storm: Efficient stochastic transformer based world models for reinforcement learning.Advances in Neural Information Processing Systems, 36:27147–27166, 2023

    Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. Storm: Efficient stochastic transformer based world models for reinforcement learning.Advances in Neural Information Processing Systems, 36:27147–27166, 2023. A Appendix A.1 Model The transformer architecture was fi...

    2023

Pith tools

Reviewed August 4, 2026 · model on record in the stance chip above.