How language models extrapolate outside the training data: A Case study in Textualized Gridworld

We greatly appreciate your thoughtful comment. We outline the key revisions made in response to your concerns. We also revised the paper upon your comments and concerns, so it would be much appreciated if you take a look.

1. Clarification on Simulation and Cognitive Maps:

• We agree that "simulative reasoning" needs clearer definition
• In cognitive science, simulation refers to mental construction and manipulation of future states [1, 2]
• Our implementation makes this concrete through:
  • State-space exploration before action selection
  • Explicit representation of possible future states
  • Integration of these representations into decision-making
• This differs from sequential CoT which:
  • Relies on step-by-step verbal reasoning
  • Doesn't construct global representations
  • Makes local decisions without explicit future simulation
• According to the o1 system card, o1 learned to refine their thinking process, try different strategies, and recognize their mistakes (https://openai.com/index/openai-o1-system-card/).
  • We initially expressed this as a “simulative reasoning” of o1 in Section 6.2., but revised as "refine its reasoning process, explore alternative strategies, and iteratively recognize and correct its mistakes"
  • It iteratively makes local decisions, and refine its own decision without explicit interaction
  • We fixed the overstatement about the o1’s capability to do tree search in Section 1.
• Our intention was to highlight that:
  • O1's performance may suggest some form of structured reasoning as a CoT
  • This aligns with our broader argument about the need for structured representations
• Our work's value stands independently of o1's specific implementation (which is elusive)

2. Relationship to Prior Works on Evaluating Cognitive Maps in LLMs and Result Differences:

• Our work can be viewed as a following work since we evaluate the challenges language models face in demonstrating spatial reasoning , particulary in extrapolated environments that are unseen during training.
• On top of them, our approach proposes a specific design of viable cognitive maps for path planning as a CoT for potential solutions to address these limitations.
• We added a list of papers and discussions of prior works in evaluating cognitive maps in language models in Appendix A.4, including your suggestions.

3. Scalability and Real-world Applications:

• Probing mental representation of the spatial layout is foundational in cognitive science for studying cognitive maps [1 - 4]

• Like seminal cognitive science studies, our work provides valuable insights about a specific cognitive capability in a controlled environment

• The scientific method often progresses from controlled experiments to broader applications - many breakthrough cognitive science papers focused solely on Gridworld experiments

• The value lies in definitively proving that current LLMs lack a crucial cognitive capability, which would be harder to demonstrate in more complex environments

• Our paper suggests that we should first establish what cognitive capabilities are missing (through controlled experiments), then develop scalable architectures to enable them

• We agree that both types of generalization are important, but our work specifically addresses a fundamental limitation in current language models' cognitive abilities

• We revised the overall storyline of the paper to better deliver our main points, throughout Abstract, Section 1, and Section 7

4. Future Work and Scalability:

• Our findings highlight a clear path forward: language models need architectural innovations to support cognitive map-like structures
• This insight opens several promising research directions:
  • Developing more general representations of cognitive maps beyond spatial reasoning
  • Creating architectures that naturally support tree-structured thinking
  • Exploring how cognitive maps could enhance other types of complex reasoning
• We revised our detailed discussion of future work throughout Section 6, and summarize them in Section 7
• For your information, we are currently pursuing one such direction through modifying current sequential language modeling to enable native generation of decision trees:
  • Training language models to generate sequences of actions requiring expansion, rather than single actions in traditional sequential modeling
  • Implementing separated generation during inference so that different branches evolve independently (similar to beam search, but in a sentence level)
  • While detailed architecture and performance analysis will be presented in future work, initial tests on challenging reasoning/planning domains (Gridworld, Game of 24, GSM8K) show promising results
• These early results suggest the cognitive map insights from our controlled study can indeed generalize to broader reasoning tasks, opening a new avenue toward cognitive language models

Thanks for your responses and updates. The presentation of the paper has improved with the updated figures, which also helps with my understanding the basics of the methods. The points around o1 have been sufficiently clarified and reflected in the paper. I also appreciated the updates which ground this in related prior work. I appreciate the authors' clarification on "simulative reasoning", although I would have appreciated a pointer to where in the paper this definition was updated (if it has been updated).

Scalability and real-world applications: The utility of a toy domain in AI and cognitive science depends on how well it reflects real world domains (ecological validity). This is important not only for applications, but also for understanding what exactly are the "cognitive capabilities" that are being studied.

I think this paper would benefit from brief descriptions of how it could be applied to other domains. The improved methods description helps with my ability to make guesses at this, but it would help if the authors filled in some of these gaps. E.g. for "Exploring how cognitive maps could enhance other types of complex reasoning", I'm not sure how cognitive maps could be designed and used for other types of problems that aren't spatial reasoning. There is also this line "Cognitive science literature refers to this aspect of human cognition [System 2] as cognitive maps — mental representations of environments that enable flexible planning." which seems to imply that all cognition that's considered "System 2" relies on cognitive maps (minor: the System 1 vs. 2 distinction is also contentious among cognitive scientists [1]).

I've updated my score from 3 to 5.

[1] https://www.psychologytoday.com/us/blog/a-hovercraft-full-of-eels/202103/the-false-dilemma-system-1-vs-system-2

We appreciate your comments about task generallization and related works. We'd like to answer your questions:

1. Generalizability of Our Cognitive Map:

• Our findings opens several promising research directions:
  • Developing more general representations of cognitive maps beyond spatial reasoning
  • Creating architectures that naturally support tree-structured thinking
  • Exploring how cognitive maps could enhance other types of complex reasoning
• For your information, we are currently pursuing one such direction through modifying current sequential language modeling to enable native generation of decision trees:
  • Training language models to generate sequences of actions requiring expansion, rather than single actions in traditional sequential modeling
  • Implementing separated generation during inference so that different branches evolve independently (similar to beam search, but in a sentence level)
  • While detailed architecture and performance analysis will be presented in future work, initial tests on challenging reasoning/planning domains (Gridworld, Game of 24, ProntoQA) show promising results
• We also appreciate your suggestion on Physics of Language Models
  • Our direct application could be Part 1: Learning Hierarchical Language Structures [1]
  • We would try pretraining a GPT-2 from scratch to learn such hierarchical language structures as a future work

2. Comparsion with previous works:

• While the main goal of previous works such as Act-Re [2] or StateAct [3] is to interact with the environment to track and self-refine its own plan to the goal (online planning), our goal is to simulate the whole plan beforehand (offline planning)
• Also, our main analysis is to probe the capability of the language model to “extrapolate” in complex environments in controlled setting, while previous works focus on the performance of the agent in a more practical setting

References:

[1] Zeyuan Allen-Zhu and Yuanzhi Li. “Physics of Language Models: Part 1, Learning Hierarchical Language Structures.” arXiv, 2024

[2] Zonghan Yang and Peng Li and Ming Yan and Ji Zhang and Fei Huang and Yang Liu. “ReAct Meets ActRe: When Language Agents Enjoy Training Data Autonomy.” arXiv, 2024

[3] Nikolai Rozanov and Marek Rei. “StateAct: State Tracking and Reasoning for Acting and Planning with Large Language Models.” arXiv, 2024

We greatly appreciate your thoughtful feedback. Below, we outline the key revisions made in response to your concerns. We also revised the paper based on your concerns and requests, so take a look!

1. Methodology: Choice of Extrapolation Metric:

• Grid size provides an objective complexity metric, not just a dimensional increase
• We quantitatively observed increased complexity:
  • Larger grids require longer planning horizons and handling more potential paths, hence increased complexity
  • We additionally analyzed that success probability decreases exponentially upon grid size
  • You can check Section 2.2 and corresponding Appendix B.3 for the details
• This controlled complexity measure allows us to definitively show extrapolation capabilities
• The task design ensures exactly one valid path, making success a clear indicator of true reasoning rather than chance

2. Methodology: Comparison with Exploration-based planning methods:

• Our baseline choices reflect the paper's focus on decision-making rather than exploration

• ToT, RAP, and similar methods are exploration techniques that:

  • Separate tree generation (sampling) from tree searching (external search algorithms)
  • Are bounded both by sampling coverage and heuristic search performance

• Our cognitive map approach differs fundamentally by:

  • Integrating the tree generation + searching process into the model generation
  • Enabling direct decision-making rather than exploration
  • Showing true extrapolation beyond training environments

• Since exploration-based methods need “actual” interaction to reach the goal state, we could not experiment on optimal planning analysis - only reachable planning analysis is available

• Even for reachable analysis, we could not compare to ToT and RAP with our main experiments

  • They both struggled with Gridworld navigation - the pretrained models’ sampling tends to be overconfident in certain directions, preventing effective tree exploration

• So as an alternative, we conducted additional experiments by enabling the language model to explicitly follow the DFS search in the gridworld

  • DFS slightly overperformed in reachable analysis than our method
  • However, when calculating the “optimality” of the generated plan, DFS required $O(n^2)$ steps, while ours only require $O(n)$ (n: optimal plan length)
  • We interpret this result as a complementary feature of planning; While Exploration-based methods are better at adapting to revise its plan to reach to the goal (online planning), our methods are better at making an optimal plan at the first place (offline planning)

• We provided the whole experiment description and analysis in Section 6.3 and corresponding Appendix E.3.

3. Justification: Scientific Value of Domain-Specific Studies:

• Probing mental representation of the spatial layout is foundational in cognitive science for studying cognitive maps [1 - 4]
• Like seminal cognitive science studies, our work provides valuable insights about a specific cognitive capability in a controlled environment
• The scientific method often progresses from controlled experiments to broader applications - many breakthrough cognitive science papers focused solely on Gridworld experiments
• The value lies in definitively proving that current LLMs lack a crucial cognitive capability, which would be harder to demonstrate in more complex environments
• We added a corresponding discussion about the justification in Section 2.2 and Appendix A.4.

4. Justification: Intended Scope:

• This work is deliberately focused on a controlled environment where we can make definitive claims about extrapolation
• The choice of Gridworld allows us to:
  • Precisely measure complexity through grid size
  • Control for confounding variables that exist in more complex domains
  • Definitively prove the failure of conventional approaches
• Attempting to simultaneously address scalability would have diluted these crucial findings and made it harder to draw clear conclusions
• This focused approach follows successful precedents in both cognitive science and AI research, where fundamental capabilities are first established in controlled settings before being scaled up
• We revised the overall storyline of the paper to better deliver our main points, throughout Abstract, Section 1, and Section 7

Thanks for your thoughtful comments about generalization and scalability. We'd like to clarify several key points:

1. Validation of Domain-Specific Studies Outside AI:

• Probing mental representation through spatial tasks like Gridworld is foundational in cognitive science [1-4]
• Like seminal cognitive science studies, our work provides valuable insights about specific cognitive capabilities in controlled environments
• This approach can extend beyond navigation to any structured reasoning tasks such as games
• The scientific method often progresses from controlled experiments to broader applications

2. Challenge to Current AI Development:

• Despite vast training data, current LLMs fundamentally lack extrapolation capabilities
• This suggests we need to first focus on designing models with basic cognitive capabilities before pursuing broad generalization
• While both complexity and task generalization matter, establishing fundamental capabilities in controlled settings is our priority

3. Controversy of the Role of the Cognitive Map and its Scalability:

• We agree that the role of the current version of the cognitive map is just giving extra supervision to the model
• Ironically, this “mere extra supervision” suddenly enables extrapolability beyond the training boundary - suggesting that the extra supervision fundamentally changes how the model represents and reasons about the environment
• On the other hand, conventional intermediate reasoning structures alone are insufficient for extrapolation - even state-of-the-art LLMs with conventional CoT fail to extrapolate in our controlled environment
• This implicitly tells you that current language models don’t have cognitive maps, or at least there is no conventional way to give model extra supervision that enables extrapolability, even with a vast pre-training corpus, large parameters, and enormous training compute
• Our work is a proof-of-concept of how we can inject such extra supervision into the model

4. Intended Scope:

• Our controlled environment enables definitive claims about extrapolation
• Gridworld allows us to:
  • Precisely measure complexity through grid size
  • Control for confounding variables
  • Definitively prove conventional approaches' failure
• Attempting to simultaneously address generalizability would dilute these crucial findings
• This follows successful precedents in cognitive science, where fundamental capabilities are first established in controlled settings before being scaled up

5. Future Work and Generalizability:

• Our findings highlight a clear path forward: language models need architectural innovations to support cognitive map-like structures
• This insight opens several promising research directions:
  • Developing more general representations of cognitive maps beyond spatial reasoning
  • Creating architectures that naturally support tree-structured thinking
• For your information, we are currently pursuing one such direction through modifying current sequential language modeling to enable native generation of decision trees:
  • Training language models to generate sequences of actions requiring expansion, rather than single actions in traditional sequential modeling
  • Implementing separated generation during inference so that different branches evolve independently (similar to beam search, but in a sentence level)
  • While detailed architecture and performance analysis will be presented in future work, initial tests on challenging reasoning/planning domains (Gridworld, Game of 24, ProntoQA) show promising results
• These early results suggest the cognitive map insights from our controlled study can indeed generalize to broader reasoning tasks, opening a new avenue toward cognitive language models

References:

[1] Epstein RA, Patai EZ, Julian JB, Spiers HJ. “The cognitive map in humans: spatial navigation and beyond.” Nat Neurosci. 2017 Oct 26;20(11):1504-1513.

[2] John O'Keefe & Lynn Nadel (1978) The Hippocampus as a Cognitive Map, Oxford University Press.

[3] Kessler, F., Frankenstein, J. & Rothkopf, C.A. “Human navigation strategies and their errors result from dynamic interactions of spatial uncertainties.” Nat Commun 15, 5677 (2024).

[4] Kadner, Florian, et al. "Finding your Way Out: Planning Strategies in Human Maze-Solving Behavior." Proceedings of the Annual Meeting of the Cognitive Science Society. Vol. 45. No. 45. 2023.

Hello reviewer 3p1W, it has been a while since we last responded to your review. We further revised the paper based on your concerns and requests, and it would be much appreciated if you could take a look and respond if there are any other concerns left.

We are writing this comment to emphasize once more that we respectfully disagree with your assessment and would like to clarify several key points (in a more concise form), and hope you read the comment and reply:

1. Distinction between Extrapolation and Generalization

• Our work specifically focuses on extrapolation (capability to solve problems of higher complexity than seen during training) rather than general domain generalization
• As clearly defined in Section 1.2, extrapolation requires: (1) learning from simple demonstrations and (2) applying this knowledge to more complex environments
• This is fundamentally different from the type of generalization achieved through large-scale pretraining, which typically enables interpolation within similar complexity levels

2. Experimental Design and Results

• Our experimental setup deliberately controls for environment complexity to isolate and test extrapolation capabilities
• The significant performance gap between cognitive maps and baselines (both implicit and conventional CoT) under identical training conditions demonstrates that the improvement cannot be attributed to mere additional supervision, but the change of the mental representation of how language models plan in Gridworld environments
• If this were simply an effect of overfitting or additional supervision, we would expect similar performance improvements in the baseline approaches, which was not observed

3. Scope and Contribution

• Our work presents a novel representation learning approach for spatial reasoning in language models, aligning with ICLR's focus on representation learning
• Our work is the first work to propose specific CoT configurations to make conventional language models "extrapolate" in a certain domian, not just in theory
• While we agree that generalization through large-scale pretraining is important, it represents a different research direction from our focus on structured representations for extrapolation
• The limitations and scaling considerations are thoroughly discussed in Section 7 of our revised version

Dear Reviewer 3p1W, the discussion period ends in 8 hours. Could you reply to our comments and revisions? We further revised the paper based on your concerns and requests, and it would be much appreciated if you could take a look and respond.

We appreciate your thoughtful feedback about our paper's relationship to human cognition. We address each point:

1. Connection to Human Cognition:

• While the direct connection between explicit decision tree and the mental representation of the cognitive map stays elusive, research in cognitive science shows that maze-solving tasks are particularly valuable for studying planning behavior in controlled environments [1 - 3]
• Especially, eye-movement studies of human maze-solving reveal three key aspects of planning behavior [4]:
  • Mental simulation through gaze patterns that reflect the maze's structure
  • Balance between depth and breadth searches based on environmental complexity
  • Adaptive planning strategies that adjust based on the number of alternatives
• Our model's cognitive map implementation was designed by similar patterns:
  • Constructs mental representations before taking actions
  • Adapts its planning depth and breath based on environmental complexity
• We will strengthen these connections in the paper by adding more detailed analysis of how our design aligns with human cognitive patterns

2. Ecological Validity:

• While true representation of human cognitive map being much more complex, the key similarity is that both humans and our model construct mental representations before taking actions, rather than using purely reactive strategies
• Our goal is not to replicate the exact human learning processes, but to demonstrate that enabling cognitive map construction leads to better extrapolation abilities
• Future work could explore more naturalistic ways of developing these capabilities, such as using self-training RL

3. Methodological Clarity:

• We agree with your suggestion about reorganizing Section 3.3
• We will move implementation details to the appendix while keeping the core experimental results in the main paper

4. Novelty and Contribution:

• Our work fundamentally differs from existing intermediate structure approaches:
  • We demonstrate that conventional intermediate reasoning structures alone are insufficient for extrapolation - even state-of-the-art LLMs with conventional CoT fail to extrapolate in our controlled environment
  • Our cognitive map implementation is not merely a different structure, but enables a crucial cognitive capability (extrapolation) that conventional approaches fundamentally lack
  • Furthermore, while conventional CoT can be learned through few-shot demonstrations, our cognitive map structure cannot - suggesting it represents a fundamentally different type of reasoning
• Our implication challenges the current trajectory of AI development:
  • Despite vast amounts of training data and sophisticated architectures, current LLMs fundamentally lack extrapolation capabilities even in simple controlled environments
  • This suggests we need to first focus on designing language models with basic cognitive capabilities demonstrated through controlled experiments before pursuing broad generalization
  • We hope this validates our experiment on controlled tasks rather than considering generalizability

Final remark:

We hope we've addressed your concerns about the paper's positioning and its relationship to human cognition. We appreciate your constructive feedback, which has helped us better revise our contributions and limitations. We look forward to incorporating these improvements in the revision. If you have any other questions, we're happy to address them. Thank you in advance!

References:

[1] Epstein RA, Patai EZ, Julian JB, Spiers HJ. “The cognitive map in humans: spatial navigation and beyond.” Nat Neurosci. 2017 Oct 26;20(11):1504-1513.

[2] John O'Keefe & Lynn Nadel (1978) The Hippocampus as a Cognitive Map, Oxford University Press.

[3] Kessler, F., Frankenstein, J. & Rothkopf, C.A. “Human navigation strategies and their errors result from dynamic interactions of spatial uncertainties.” Nat Commun 15, 5677 (2024).

[4] Kadner, Florian, et al. "Finding your Way Out: Planning Strategies in Human Maze-Solving Behavior." Proceedings of the Annual Meeting of the Cognitive Science Society. Vol. 45. No. 45. 2023.

Dear Reviewer 5Pyx, the discussion period ends in 8 hours. Could you reply to our comments and revisions? We further revised the paper based on your concerns and requests, and it would be much appreciated if you could take a look and respond.

Our work advances the field of representation learning by introducing a novel approach to spatial reasoning in language models through targeted Chain-of-Thought configurations. Notably, we demonstrate that through careful fine-tuning alone, language models can achieve extrapolation capabilities in spatial reasoning tasks - a significant finding that bridges theoretical possibilities with practical implementation. This result has important implications for representation learning and opens new research directions in the field. Given ICLR's focus on representation learning advances, we believe our contribution aligns well with the conference's scope and could spark valuable discussions in the community. We would deeply appreciate the reviewers' careful consideration of our revisions and responses to their concerns.

Overall, we believe this work meaningful, as is reflected in our score (8). We believe some more comparison with length generalisation work would be good to include in this paper.