Compositional pretraining improves computational efficiency and matches animal behavior on complex tasks

Abstract

1Recurrent neural networks (RNN) are ubiquitously used in neuroscience to capture both neural dynamics and behaviors of living systems. However, when it comes to complex cognitive tasks, training RNNs with traditional methods can prove difficult and fall short of capturing crucial aspects of animal behavior. Here we propose a principled approach for identifying and incorporating compositional tasks as part of RNN training. Taking as target a temporal wagering task previously studied in rats, we design a pretraining curriculum of simpler cognitive tasks that reflect relevant sub-computations. We show that this pretraining substantially improves learning efficacy and is critical for RNNs to adopt similar strategies as rats, including long-timescale inference of latent states, which conventional pretraining approaches fail to capture. Mechanistically, our pretraining supports the development of slow dynamical systems features needed for implementing both inference and value-based decision making. Overall, our approach is an important step for endowing RNNs with relevant inductive biases, which is important when modeling complex behaviors that rely on multiple cognitive computations.