GRAPH SCHEMAS AS ABSTRACTIONS FOR TRANSFER LEARNING, INFERENCE, AND PLANNING

Abstract

We propose schemas as a model for abstractions that can be used for rapid transfer learning, inference, and planning. Common structured representations of concepts and behaviors-schemas-have been proposed as a powerful way to encode abstractions. Latent graph learning is emerging as a new computational model of the hippocampus to explain map learning and transitive inference. We build on this work to show that learned latent graphs in these models have a slot structureschemas-that allow for quick knowledge transfer across environments. In a new environment, an agent can rapidly learn new bindings between the sensory stream to multiple latent schemas and select the best fitting one to guide behavior. To evaluate these graph schemas, we use two previously published challenging tasks: the memory & planning game and one-shot StreetLearn, that are designed to test rapid task solving in novel environments. Graph schemas can be learned in far fewer episodes than previous baselines, and can model and plan in a few steps in novel variations of these tasks. We further demonstrate learning, matching, and reusing graph schemas in navigation tasks in more challenging environments with aliased observations and size variations, and show how different schemas can be composed to model larger environments.

1. INTRODUCTION

Discovering and using the right abstractions in new situations affords efficient transfer learning as well as quick inference and planning. Common reusable structured representations of concepts or behaviors-schemas-have been proposed as a powerful way to encode abstractions (Tenenbaum et al., 2011; Mitchell, 2021) (Fig. 1A ). Humans excel at this ability, argued to be a key factor behind intelligence and a fundamental limitation in current AI systems (Shanahan & Mitchell, 2022) . A computational model with the ability to discover and reuse previously-learned structured knowledge to behave and plan in novel situations can advance AI systems. Experimental evidence suggests that some animals have this ability. Rats and mice tend to learn new environments faster if they can reuse past schemas (Zhou et al., 2021; Tse et al., 2007) , and there are cells in the macaque hippocampus encoding schemas for spatial abstractions (Baraduc et al., 2019) . Neural circuits in hippocampus (HPC) and prefrontal cortex (PFC) are both implicated in different aspects of schema learning, recognition, update, and maintenance. Recent work in mice suggested that PFC encodes the schema while HPC binds this schema to sensorimotor specifics of a task (Samborska et al., 2022) . New experiences are also shown to be learned within a single trial if it fits an existing schema, and such knowledge rapidly becomes HPC-independent much faster than typically expected for memory consolidation. Kumaran et al. (2016) proposed an updated complementary learning systems (CLS) theory based on this evidence, but there is so far no explicit demonstration of such rapid learning with schema reuse as far as we know. Schema-based mechanisms also aid in memory consolidation over longer time-scales into reusable knowledge structures (Gilboa & Marlatte, 2017) . Even though spatial navigation and memory dominated research in HPC,it is also known to be important for non-spatial declarative memory and relational reasoning tasks. Structured relational representations have been proposed as a common mechanism that generalizes to spatial and non-spatial tasks and memory (Eichenbaum & Cohen, 2014; Stachenfeld et al., 2017) . Recent work on cognitive maps in HPC model these representations using higher order latent graph structures and show generalization to disparate HPC functions (Whittington et al., 2020; George et al., 2021; Sharma et al., 2021; Whittington et al., 2021) . In this work, we take one such model and provide a concrete computational model of abstractions using graph schemas. We describe how graph schemas can be learned and then reused for transfer learning, quick inference, and planning for behavior in new situations by rapidly learning observation bindings and discovering the best schema online. We build these schemas using clone-structured cognitive graphs (CSCG), a computation model of cognitive maps in HPC (George et al., 2021) . In particular, CSCG work showed that latent high order graphs can be learned in highly aliased settings using a smooth, probabilistic, parameterization of the graph learning problem using gradient-based optimization. We use navigation as our setting, which requires handling perceptual aliasing, a difficult problem in which standard methods fail (Lajoie et al., 2018) . Using higher order graphs helps with modeling aliased observations (Xu et al., 2016) . CSCGs also generalize beyond spatial navigation to non-spatial domains (Dedieu et al., 2019) , so schemas on CSCG can extend beyond spatial navigation. This generalization of graph schemas neatly maps to the idea of useful abstractions as template learning and structure mapping as described in Shanahan & Mitchell (2022) . Our computational model provides a concrete implementation of some of the operations described in that work. For example, their operation of seeing similarity can be interpreted as the process of finding which of the known templates best describes the observations in terms of likelihood in our model, and binding different instantiations to nodes in a structured graph can be interpreted as learning different emission bindings in a graph schema. 

2. RELATED WORK

Large language models (LLMs) have shown impressive demonstrations of few-shot generalization to new tasks without re-training, possibly using in-context learning (Brown et al., 2020; Wei et al., 



Overview of using graph schemas in CSCG model. A. Schemas serve as abstractions that can be matched with new experiences. Multiple schemas can be composed together to model observations that could be aliased. B. Schematic of cloned HMM with multiple cloned states sharing the emissions to model them in different contexts. C. A graphical model of CSCG with action conditioning. D. Schemas in CSCG are extracted from learned models with transitions and emissions, and then unbinding the specific emissions (shown as colors) but keeping the clone structure (shown as node shapes). E. In a novel environment with different observations, the agent navigates and finds the schema that best explains the observations by learning new bindings. F. Planning can be performed in a new room using the matched schema. As the plan is being executed, new observations are used to update the beliefs about current location and re-plan as needed.

