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

TL;DR

This paper investigates how language models extrapolate in complex environments using a textualized Gridworld, revealing limitations of conventional methods and proposing cognitive maps inspired by human cognition to improve extrapolation capabilities.

Contribution

Introduces a novel cognitive map framework for path planning in language models, inspired by human mental representations, to enhance extrapolation in unseen environments.

Findings

Cognitive maps improve extrapolation to larger, unseen environments.

Conventional methods like Chain of Thought fail in larger, novel environments.

Cognitive maps exhibit humanlike mental simulation and rapid adaptation.

Abstract

Language models' 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 language models' 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 language models. Our comparison with exploration-based methods further illuminates the complementary strengths of offline planning and online exploration.