How language models extrapolate outside the training data: A case study in Textualized Gridworld (2406.15275v4)

Abstract: LLMs' ability to extrapolate learned behaviors to novel, more complex environments beyond their training scope is highly unknown. This study introduces a path planning task in a textualized Gridworld to probe LLMs' extrapolation capabilities. We show that conventional approaches, including next token prediction and Chain of Thought (CoT) finetuning, fail to extrapolate in larger, unseen environments. Inspired by human cognition and dual process theory, we propose cognitive maps for path planning, a novel CoT framework that simulates humanlike mental representations. Our experiments show that cognitive maps not only enhance extrapolation to unseen environments but also exhibit humanlike characteristics through structured mental simulation and rapid adaptation. Our finding that these cognitive maps require specialized training schemes and cannot be induced through simple prompting opens up important questions about developing general-purpose cognitive maps in LLMs. Our comparison with exploration-based methods further illuminates the complementary strengths of offline planning and online exploration.

Cognitive Map for LLMs: Optimal Planning via Verbally Representing the World Model

This paper addresses the challenge of enabling LLMs (LMs) to perform robust, multi-step planning tasks that require detailed simulations, drawing inspiration from human cognitive processes. Specifically, it introduces a novel approach where a LLM constructs a "cognitive map" of a given environment to enhance its planning capabilities. The cognitive map is a verbally represented world model that allows the LLM to simulate various environmental states and plan optimally. The authors demonstrate the efficacy of this method through experiments involving the Gridworld path planning task.

Key Contributions

1. Cognitive Map Construction: The paper proposes an approach where the LLM constructs a tree-structured verbal representation of the world model, termed the cognitive map. This map is generated in a sequential manner through three key processes:
   • Sampling: The model samples plausible actions for each state.
   • Propagation: The model simulates the outcome of each action to explore new states.
   • Backtracking: Once a goal state is reached, the model traces back the optimal path to refine its plan.
2. Human-Cognitive Characteristics: The cognitive map method has been shown to exhibit two characteristics similar to human cognition:
   • Generalization: The ability to apply learned knowledge to solve problems in larger, unseen environments.
   • Rapid Adaptation: The ability to effectively learn and perform with limited training data.
3. Empirical Validation: The authors conduct rigorous experiments on the Gridworld path planning task, showing that the cognitive maps significantly enhance both optimal and reachable planning abilities. The results indicate a substantial improvement, with up to 57.5% enhancement in optimal planning and 56.4% in reachable planning compared to baseline models.

Numerical Results and Analysis

• Optimal Planning: The best-performing configuration for the optimal planning task (bwd marking deadend) achieved a success rate of 76.5%, a significant improvement over the baseline implicit learning model (none), which had a success rate of only 19%.
• Reachable Planning: For reachable planning tasks, the cognitive map approach (bwd marking deadend) obtained an 88.5% success rate, highlighting its robustness in generating effective plans even when the goal is simply to reach a state rather than to find the optimal path.
• Rapid Adaptation: The cognitive map model demonstrated rapid learning convergence, achieving a significant success rate early in the training (79.13% at 500 steps), which suggests the ease with which the model learns the cognitive map construction.

Theoretical Implications and Future Directions

The cognitive map approach aligns with theories in cognitive science, particularly dual-process theories which distinguish between fast, automatic, and intuitive thinking (System 1), and slow, deliberate, and analytical thought (System 2). This method essentially models System 2-like reasoning within LLMs by enabling internal simulations before making decisions, thereby allowing for more sophisticated and long-horizon planning.

The theoretical implications are profound:

• Enhanced Generalization: The ability to generalize to larger and more complex environments mirrors human cognitive flexibility and sets a new benchmark for artificial planning systems.
• Human-Like Planning: By modeling human cognitive processes, this approach potentially paves the way for developing LLMs that not only understand language but can also apply this understanding to navigate and interact with the world in a more human-like manner.

Future research directions may include:

• Scalability: Investigating sampling strategies that enable cognitive map construction in significantly larger and more complex environments, such as those encountered in web or travel agent tasks.
• Broader Applications: Extending the cognitive map approach to various domains where planning and decision-making tasks are crucial, such as robotics, logistics, and automated personal assistants.
• Integration with Model-Free Planning: Exploring hybrid approaches that combine the strengths of model-free (System 1) and model-based (System 2) planning methods to enhance overall performance and adaptability.

Conclusion

The paper provides a compelling blueprint for leveraging cognitive maps in LLMs to achieve robust and optimal planning abilities. By systematically constructing and utilizing verbal representations of the environment, the LLM closely emulates human cognitive processes, demonstrating significant advancements in task performance and learning efficiency. This research opens new avenues for developing AI systems that not only process language but also interact with the world in a meaningful and intelligent manner.