SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kevin J. Miller
Google DeepMind and University College London
kevinjmiller@deepmind.com
Maria Eckstein
mariaeckstein@deepmind.com
Matthew M. Botvinick
botvinick@deepmind.com
Zeb Kurth-Nelson
zebk@deepmind.com
Computational cognitive models are a fundamental tool in behavioral neuroscience. They instantiate in software precise hypotheses about the cognitive mechanisms underlying a particular behavior. Constructing these models is typically a difficult iterative process that requires both inspiration from the literature and the creativity of an individual researcher. Here, we adopt an alternative approach to learn parsimonious cognitive models directly from data. We fit behavior data using a recurrent neural network that is penalized for carrying information forward in time, leading to sparse, interpretable representations and dynamics. When fitting synthetic behavioral data from known cognitive models, our method recovers the underlying form of those models. When fit to laboratory data from rats performing either a reward learning task or a decision-making task, our method recovers simple and interpretable models that make testable predictions about neural mechanisms.
Fitting quantitative cognitive models to behavioral data is a fundamental tool in many areas of psychology and neuroscience [51, 12, 10, 39]. These models can be viewed as mechanistic hypotheses about the cognitive processes used by the brain. Viewed in this way, they act as an explicit software instantiation of a particular cognitive hypothesis, and can be used to make precise quantitative predictions about behavioral and neural data. A traditional modeling pipeline is for a human researcher to iterate over a three-step process: first to propose a candidate model structure (e.g. Q-learning), second to optimize model parameters (e.g. learning rate) with respect to a behavioral dataset, and finally to check whether the resulting model reproduces scientifically important features of the dataset. However, discovering an appropriate model structure is difficult (there are many possible structures to explore) as well as biased (the best structure may be one that the researcher has not thought of).
An alternative approach is fitting recurrent neural networks (RNNs) directly to behavior using supervised learning [16, 19, 46, 2]. RNNs are highly expressive and can approximate a wide variety of model structures in their weights. This is therefore a way of discovering a well-fitting model from data automatically. The drawback of this approach is that the resulting RNN is a "black box": a complex system that itself requires further analysis if it is to yield insight into cognitive mechanism.
Here, we propose a solution that aims to achieve the best of both worlds: automated discovery of human-interpretable cognitive models. Our approach, which we call disentangled RNN or "DisRNN", draws on recently developed methods from machine learning for "disentangling" [23, 24, 8, 21].
These methods encourage networks to learn representations in which each dimension corresponds to a single true factor of variation in the data [23, 8]. DisRNN encourages disentangling in two ways. The first is to separate the update rule for each element of the latent state into separate sub-networks. The second is to use information "bottlenecks", which impose a penalty on maintaining information within the network, to the inputs and outputs of these sub-networks.
We fit DisRNN on sequential behavioral datasets from rats and artificial agents performing two classic cognitive tasks: a dynamic two-armed bandit task [13, 28, 26, 43] and a "pulse accumulation" decision-making task [7, 45, 14, 41]. First, we generate synthetic datasets for the two-armed bandit task using artificial agents with known learning algorithms (Q-Learning and Actor-Critic) which have different update rules and carry different information between timesteps. When fitting these synthetic datasets, DisRNN correctly recovers the timecourses of latent state information and the structure of the update rules. Second, we fit DisRNN on large laboratory datasets generated by rats performing the same two-armed bandit task [37]. We find that DisRNN provides a similar quality of fit to the best known human-derived cognitive model of this dataset [37] as well as to an unconstrained neural network [16]. Third, we generate a synthetic dataset for the decision-making task using an agent with a known decision algorithm (bounded accumulation) and show that disRNN is able to correctly recover the structure of its decision rule. Finally, we fit disRNN to a large laboratory dataset of rats performing this task [7], and find that it recovers a strategy similar to that of the best known human-derived model for this dataset. In all cases, we find that DisRNN is able to learn simple human-interpretable cognitive strategies which can be used to make predictions about behavioral and neural datasets.
Our strategy for encouraging networks to adopt disentangled representations is directly inspired by work using variational autoencoders (VAE). Specifically,  $\beta$ -VAE [23, 8], by scaling the KL loss of a variational autoencoder's sampling step (which can be viewed as requiring information to pass through a Gaussian information bottleneck), often learns sparse, disentangled representations in which each latent variable corresponds to a single true factor of variation in the data. While  $\beta$ -VAE considered feedforward autoencoders, a wide variety of techniques have been proposed combining elements of feedforward VAEs with recurrent neural networks [20]. We adapt these ideas here in a way that emphasizes interpretability and is appropriate for cognitive model discovery.
The use of recurrent neural networks as cognitive models has a long history [5, 27], and recent work has expanded the toolkit for fitting networks to behavioral datasets and for interpreting them [16, 19, 46, 15]. This approach of fitting standard neural networks and working to interpret the resulting fits benefits from the growing toolkit for neural network interpretability [47, 35, 36]. Our approach is complementary: instead of fitting standard networks and developing tools to interpret them post hoc, we develop networks which are incentivized to learn easily-interpretable solutions. Within neuroscience, a number of methods exist which attempt to discover interpretable latent dynamics from neural recording data [40, 50, 29, 44]. To our knowledge, these ideas have not been applied in the context of cognitive model discovery from behavioral data.
The two-armed bandit task we consider here is one of a family of dynamic reward learning tasks that have been heavily studied in behavioral neuroscience. This has led to a large library of candidate cognitive models [13, 4, 17, 32, 11, 26, 34, 37, 33, 3, 42]. The development of such models typically follows a theory-first approach, beginning with an idea (e.g. from optimality [13, 42], from machine learning [33], or from neurobiology [34]). Several studies have fit behavioral dataset using highly flexible classical models [26, 37, 32, 11]. One of these has attempted to use these fits as a basis for a process of data-driven cognitive model discovery by compressing the revealed patterns into a simpler model which is cognitively interpretable [37].
Two very recent papers have proposed a very similar workflow of discovering cognitive strategies from behavioral data using constrained neural networks. The first uses "hybrid" models that combine elements of classic structured models and flexible neural networks [18]. The second constrains networks by limiting them to a very small number of hidden units [2]. Future work should compare these approaches directly on matched datasets, as well as explore other possible opportunities for adapting modern machine learning techniques as tools for cognitive model discovery.
3 Disentangled Recurrent Neural Networks
An RNN trained to match a behavioral dataset containing temporal dependencies [16, 19, 46] can be viewed as maintaining a set of latent variables which carry information from the past that is useful for predicting the future. The weights of the network can be viewed as defining a set of update rules defining how each latent variable evolves over time based on external observations as well as the previous values of the other latent variables. Standard RNN architectures (such as the GRU [9]) typically learn high-dimensional latent representations with highly entangled update rules. This makes it difficult to understand which cognitive mechanisms they have learned, limiting their usefulness as cognitive hypotheses.
Disentangled RNN Architecture
Each orange connection is an information bottleneck
Figure 1: Network Architecture. Left: Overall architecture of the DisRNN. Each latent variable is updated by a separate feedforward neural network (MLP: Multilayer Perceptron). Bottlenecks (orange connections) are imposed both on the inputs to these networks and on each latent variable to encourage interpretable representations. Right: A single information bottleneck. Output  $\tilde{z}$  is a random sample from a Gaussian distribution determined by input  $z$  and bottleneck parameters  $m$  and  $\sigma$ . The bottleneck is associated with a KL loss penalizing information transfer.
Example Information Bottleneck
3.1 Network Architecture
In order to learn an interpretable cognitive model, we encourage sparsity and independence by imposing bottlenecks which penalize the network for using excess information (Figure 1, right) [49, 1, 31, 23, 24, 8, 21]. The bottlenecks in our networks limit information flow using Gaussian noise. Our implementation uses multiple information bottlenecks, each of which is a noisy channel defined by two learned parameters: a "multiplier"  $m$  and noise variance  $\sigma$ . The output on each timestep is sampled from a Gaussian distribution:
$$
\tilde {z} _ {t} \sim \mathcal {N} (m z _ {t}, \sigma) \tag {1}
$$
where  $x$  is the scalar input to the bottleneck and  $\tilde{x}$  is the scalar output. Each bottleneck is associated with a loss which penalizes the sampling distribution for deviating from the unit Gaussian:
$$
L _ {\text {b o t t l e n e c k}} = \sum_ {t} D _ {K L} \left(\mathcal {N} \left(m _ {b} x _ {b, t}, \sigma_ {b}\right) \mid \mid \mathcal {N} (0, 1)\right) \tag {2}
$$
where  $D_{KL}$  is the Kullback-Leibler divergence, a measure of the difference between two probability distributions. This difference, and therefore the cost of the bottleneck, will be zero in the case that  $m = 0$  and  $\sigma = 1$ . In that situation, the bottleneck will output samples from the unit Gaussian. These outputs will be independent of the input, meaning that no information will flow through the bottleneck. We refer to these as "closed bottlenecks". In all other situations,  $L_{b}$  will be greater than zero, and some information will pass through. In our experiments, bottlenecks that are carrying useful information typically fit  $m \approx 1$  and  $\sigma \ll 1$ . We refer to these as "open bottlenecks".
We use these information bottlenecks to encourage cognitively interpretable models by encouraging two distinct kinds of disentangling. The first is disentangling in the latent variables themselves: we would like these to capture separable latent processes and to be relatively few in number. We
encourage this kind of simplicity by imposing a separate bottleneck on each scalar element of the network's hidden state. We refer to these as the "latent bottlenecks".
$$
\tilde {z} _ {t} ^ {i} \sim \text {b o t t l e n e c k} _ {i} \left(z _ {t} ^ {i}\right) \tag {3}
$$
The second kind of simplicity is in the update rules for the latent variables. We would like each variable to be updated by its own separate rule, and we would like each of these rules to be as simple as possible. We encourage this kind of simplicity by updating each element of the network's hidden state,  $z_{t}^{i}$ , using a separate learned update rule defined by a separate set of parameters. Each update rule consists of a multilayer perceptron (MLP) which defines a multiplicative update to be applied to its corresponding latent variable. These "Update MLPs" have access to the values of all elements of the previous timestep's hidden state,  $\tilde{z}_{t-1}^{i}$  (having passed already through its latent bottleneck), as well as to the network's current observations  $\mathbf{o}_{t}$ . Each element of the MLP input (i.e. of  $\tilde{\mathbf{z}}_{t-1}$  and of  $\mathbf{o}_{t}$ ) must pass through an additional information bottleneck, which we refer to as "Update MLP Bottlenecks". The output of each update MLP is a scalar weight  $w_{t}^{i}$  and update target  $u_{t}^{i}$ , which are used together to update the value of the associated latent variable  $z^{i}$  in a manner analogous to the Gated Recurrent Unit [9].
$$
w _ {t} ^ {i}, u _ {t} ^ {i} = \mathrm {M L P} _ {\text {u p d a t e}} ^ {i} (\text {b o t t l e n e c k} (\tilde {\mathbf {z}} _ {\mathbf {t} - \mathbf {1}}, \mathbf {o} _ {\mathbf {t}})) \tag {4}
$$
$$
z _ {t} ^ {i} = (1 - w _ {t} ^ {i}) \tilde {z} _ {t - 1} ^ {i} + w _ {t} ^ {i} u _ {t} ^ {i} \tag {5}
$$
The latent variables are used on each timestep to make a prediction about the target,  $y$ , using a separate "Choice MLP":
$$
\hat {y} _ {t} = \mathrm {M L P} _ {\text {o u t p u t}} (\bar {\mathbf {z}} _ {\mathbf {t}}) \tag {6}
$$
We train the model end-to-end by gradient descent to minimize the total loss function:
$$
L _ {\text {t o t a l}} = L _ {\text {s o f t m a x}} (y, \hat {y}) + \beta \sum_ {b \in \text {b o t t l e n e c k s}} L _ {\text {b o t t l e n e c k}} (b) \tag {7}
$$
where  $y$  are supervised targets,  $L_{\mathrm{softmax}}$  is a softmax cross-entropy loss. The hyperparameter  $\beta$  scales the cost associated with passing information through the bottleneck [23, 8]. Different values of  $\beta$  are expected to produce solutions which adopt different tradeoffs between predictive accuracy and model simplicity.
3.2 Training Details
For the two-armed bandit task, datasets consisted of sequences of binary choices made (left vs. right) and outcomes experienced (reward vs. no reward) by either a rat [37] or an artificial agent (Q-Learning or Leaky Actor-Critic, see below). Each episode corresponded to a single behavioral session. On each timestep, the observation given to the network consisted of the choice and reward from the corresponding trial, and the target was the choice on the subsequent trial. For the pulse accumulation task, datasets consisted of sequences of integer pulse counts (left vs. right) observed and binary choices made (left vs right) by either a rat [7] or an artificial agent (Bounded Accumulation, see below). Each episode corresponded to a single trial. On each timestep, the observation given to the network consisted of the left and right pulse counts in a corresponding 10ms timebin. The target was null for all timesteps during the stimulus, and on the final timestep of each trial indicated the binary choice made by the rat.
All networks trained for this paper were of the same size and trained in the same way. Each network had five latent variables. Update MLPs consisted of three hidden layers containing five units each. The Choice MLP consisted of two hidden layers of two units each. We used the rectified linear (ReLU) activation function. Networks were defined using custom modules written using Jax [6] and Haiku [22]. Network parameters were optimized using gradient descent and the Adam optimizer [30], with a learning rate of  $5 \times 10^{-3}$ . We typically trained networks for  $10^{5}$  steps, except that networks with very low  $\beta$  ( $10^{-3}$  or  $3 \times 10^{-4}$ ) required longer to converge and were trained for  $5 \times 10^{5}$  steps. Using a second-generation TPU, models required between four and fifty hours to complete this number of training steps.
4 DisRNN Recovers True Structure in Synthetic Datasets
We first apply DisRNN to synthetic datasets in which the process generating the data is known. We consider a task that has been the subject of intensive cognitive modeling efforts in psychology and neuroscience, the dynamic two-armed bandit task (Figure 2) [13, 28, 26, 43, 37]. In each trial of this task, the agent selects one of two available actions and then experiences a probabilistic reward. In our instantiation of the task, rewards are binary, and the reward probability conditional on selecting each arm drifts independently over time according to a bounded random walk [37].
We generated synthetic datasets from two reinforcement learning agents performing this task: Q-Learning and Leaky Actor-Critic. These agents have markedly different latent variables and update rules, allowing us to test whether DisRNN can recover the correct structure of each agent.
Figure 2: Dynamic Two-Armed Bandit Task. a: In each trial, the rat selects one of two actions (left or right) and receives one of two outcomes (reward or no reward). b: Example behavioral session [37]. Orange and purple lines show the drifting generative reward probabilities for each action. The placement of each tick mark (above or below) indicates the choice of the animal on that trial, while the color of the tick indicates whether or not the animal received reward. We will use this session as a running example to illustrate the dynamics of artificial agents and RNNs.
4.1 Q-Learning Agent
The Q-learning agent [48] maintains two latent variables,  $Q_{\mathrm{left}}$  and  $Q_{\mathrm{right}}$ . Each of these is associated with one of the available actions, and gives a running average of recent rewards experienced after taking that action. They are updated according to the following rules:
$$
Q _ {t + 1} \left(a _ {t}\right) = (1 - \alpha) Q _ {t} \left(a _ {t}\right) + \alpha r _ {t} \tag {8}
$$
$$
Q _ {t + 1} (a \neq a _ {t}) = Q _ {t} (a \neq a _ {t}) \tag {9}
$$
where  $\alpha$  is a learning rate parameter. We visualize these update rules by plotting the updated value of  $Q$  as a function of its initial value and of the trial type (Figure 3a). A key feature is that each  $Q$  value only changes when its corresponding action is selected.
The Q-Learning agent selects actions based on the difference in  $Q$ -values:
$$
a _ {t} \sim \operatorname {L o g i s t i c} \left(\beta \left(Q _ {t, \text {l e f t}} - Q _ {t, \text {r i g h t}}\right)\right) \tag {10}
$$
where  $\beta$  is an inverse temperature parameter. We generated a large synthetic dataset consisting of choices made and rewards received by the Q-Learning agent ( $\alpha = 0.3$ ,  $\beta = 3$ ; 1000 sessions of 500 trials each) (Figure 3b). We then trained a DisRNN to imitate this dataset. Inspecting the learned bottleneck parameters of this network (Figure 3e), we find that just two of its latents have open bottlenecks ( $\sigma \ll 1$ ), allowing them to pass information forward through time. We find that the update rules for these latents have open bottlenecks for the input from both previous choice and previous reward. We find that the three latents with closed ( $\sigma \approx 1$ ) bottlenecks take on near-zero values on all timesteps, while the two with open bottlenecks follow timecourses that are strikingly similar to those of  $Q_{\mathrm{left}}$  and  $Q_{\mathrm{right}}$  (albeit centered on zero instead of on 0.5; Figure 3c). We visualize the learned update rule associated with each latent by probing its associated update MLP (Figure 3d). We find that these update rules are strikingly similar to those for  $Q_{\mathrm{left}}$  and  $Q_{\mathrm{right}}$ . They reproduce the key feature that each is only updated following choices to a particular action.
4.2 Leaky Actor-Critic Agent
The Leaky Actor-Critic agent uses a modification of the policy gradient learning rule [48](Section 2.8). Like the Q-Learning agent, it makes use of two latent variables, but uses very different rules
Q-Learning Agent | DisRNN Trained on Q-Learning Data
a
Example Session
Update Rules
d
Example Session
e
Update Rules
.
Information Bottlenecks
Figure 3: DisRNN Recovers Latent Dynamics of Q-Learning. a: Q-Learning agent run using the choices and rewards from the example behavioral session (Figure 2) to generate timeseries for  $Q_{\text{left}}$  and  $Q_{\text{right}}$ . b: Visualization of the Q-Learning update rules. Each panel visualizes the update of either  $Q_{\text{left}}$  (top row) or  $Q_{\text{right}}$  (bottom row) following a particular combination of choice and outcome (columns). The post-update value is shown on the vertical axis, as a function of the pre-update value, on the horizontal. Dashed lines are identity. c: Dependency graph of Q-learning. Q-learning computes an update for each latent ( $Q_{\text{left}}$  and  $Q_{\text{right}}$ ) on each timestep. The update for  $Q_{\text{left}}$  depends on choice and reward (Equation 8). The new value of  $Q_{\text{left}}$  is a weighted sum of its old value (dashed line) and the update.  $Q_{\text{right}}$  is similarly updated. d: DisRNN trained on a synthetic behavioral dataset generated by the Q-Learning agent. This panel shows data from the trained model run with choice and reward input from the same example session. Latents have been ordered, signed, and colored to highlight similarities with the Q-Learning agent. e: Visualization of the learned update rules (equations 4 and 5) for the first two latents. f: Learned parameters of the information bottlenecks. Left: Transmission bottlenecks (equation 3). Right: Update MLP bottlenecks (equation 4). Darker colors indicate open bottlenecks: for example, the update for latent 1 can depend on choice and reward. g: Dependency graph of cognitive model learned by DisRNN.
to update them. The first variable,  $V$ , tracks "value". It keeps a running average of recent rewards, regardless of the choice that preceded them:
$$
V _ {t + 1} = \left(1 - \alpha_ {v}\right) V _ {t} + \alpha_ {v} r _ {t} \tag {11}
$$
where  $\alpha_{v}$  is a learning rate parameter.
The second,  $\Theta$ , is a "policy" variable that determines the probability with which the agent will select each action. The policy gradient algorithm defines two variables  $\theta_{left}$  and  $\theta_{right}$ , each associated with one of the actions and updated using the following rules:
$$
\theta_ {t + 1} \left(a _ {t}\right) = \left(1 - \alpha_ {f}\right) \theta_ {t} \left(a _ {t}\right) + \alpha_ {l} \left(r _ {t} - V _ {t}\right) \left(1 - \pi_ {t} \left(a _ {t}\right)\right) \tag {12}
$$
$$
\theta_ {t + 1} (a \neq a _ {t}) = \left(1 - \alpha_ {f}\right) \theta_ {t} (a) - \alpha_ {l} \left(r _ {t} - V _ {t}\right) \pi_ {t} (a) \tag {13}
$$
DisRNN Trained on Actor-Critic Data
Figure 4: DisRNN Recovers Latent Dynamics of Leaky Actor-Critic. a: Actor-Critic agent run on the example rat behavioral session (Figure 2) to generate timeseries for  $V$  and  $\Theta$ . b: Visualization of the Actor-Critic update equations. The update for  $v$  depends on reward but is independent of choice. The update for  $\Theta$  depends on choice and reward, as well as the value of  $V$ . c: Dependency graph of Actor-Critic learning. d: DisRNN trained on a synthetic behavioral dataset generated by an Actor-Critic agent. e: Visualization of the learned update rules for the first two latents. f: Learned bottleneck parameters. Left: Transmission bottlenecks (equation 3). The first two latents have open bottlenecks; the remaining three latents have closed bottlenecks. Right: Update MLP bottlenecks (equation 4). g: Dependency graph of DisRNN trained on Actor-Critic data.