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) .

