Neural mechanisms of cognitive flexibility: Belief updating in dynamic environments with sparse rewards

John Schwarcz, Jan Philipp Bauer, Gabrielle Marmur, Haneen Rajabi, Eran Lottem, Jonathan Kadmon

Keywords: Theoretical neuroscience, neural dynamics, cognitive neuroscience, belief states, artificial neural networks, animal behavior

TL;DR: This study integrates theoretical modeling, neural network simulations, and rodent experiments to elucidate how agents can rapidly adapt to changing environments without requiring immediate rewards.

Abstract:

Humans and animals must develop adaptive strategies to optimize decision-making in dynamic and uncertain environments, often without the benefit of immediate rewards. While existing literature posits that animals use internal "belief" states as the foundation for their decision policy, the mechanism for updating them in a dynamic environment remains unclear. Furthermore, there is no known neural mechanism that can implement belief updates without the need for a reward. To address this gap, we take a multidisciplinary approach that integrates theoretical derivation, training artificial neural networks, and behavioral experiments in rodents to explore potential neural mechanisms of cognitive flexibility.

A belief state is a joint probability distribution over all relevant latent variables of the environment. Updating the joint distributions using only partial observations and marginalizing to obtain estimators is computationally demanding, in particular when latent variables are changing. Moreover, it is nontrivial for a neural network to learn how to implement this complex inference. To tackle these challenges, we introduce a novel change-detection task specifically designed to capture the complexities of partially observed dynamic environments. We formulate a Bayesian theory for sequentially updating joint probabilities and demonstrate that neural networks can accomplish the task near optimally, even in the absence of immediate rewards. We show that the network dynamics mirror the sequential update of the Bayesian latent state estimators. Furthermore, rodents trained on this task show behavior that aligns with our theoretical model and neural network simulations, suggesting that mice utilize dynamic internal state representation and inference to solve this task. Overall, our findings elucidate the computational principles behind flexible cognitive behavior that allows both biological and artificial agents to achieve zero-shot adaptation: modifying their behavior policy to reflect changes in the environment without the need for trial and error.

Primary Area: applications to neuroscience & cognitive science

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 6315