Title: Cognitive Model Discovery via Disentangled RNNs
Venue: NeurIPS 2023
Accepted: yes
Keywords: Cognitive modeling

neural networks

interpretability

disentangling

neuroscience

rodent behavior
Abstract: Computational cognitive models are a fundamental tool in behavioral neuroscience. They embody 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 excess information between timesteps, 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 choice data from rats performing a bandit task, our method recovers simple and interpretable models that make testable predictions about neural mechanisms.

Cognitive Model Discovery via Disentangled RNNs

Abstract

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.

1 Introduction

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.

2 Related Work

Our strategy for encouraging networks to adopt disentangled representations is directly inspired by work using variational autoencoders (VAE). Specifically, β -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 β -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.

DisentangledRNNArchitecture

Each orange connection is an information bottleneck

Target

ChoiceMLP

y

y

Example Information Bottleneck

SupervisedLoss

Latents:

.

UpdateMLP1

+

Bottleneck

Input:

FitParameters:m,

Sample from

Bottleneck

Output:

Zt+1

Normal(mz, α)

之t+1

UpdateMLPN→zt+

t+1

t+1

Observations:

KL LosS

Thesebottleneckscause

Thesebottleneckscause

disentangledlatents

Normal(0, 1)

disentangled updaterules

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 σ . The output on each timestep is sampled from a Gaussian distribution:

where x is the scalar input to the bottleneck and ˜ x is the scalar output. Each bottleneck is associated with a loss which penalizes the sampling distribution for deviating from the unit Gaussian:

where D KL is the Kullback-Leiber 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 σ = 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 ≈ 1 and σ ≪ 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'.

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 i t , 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, ˜ z i t -1 (having passed already through its latent bottleneck), as well as to the network's current observations o t . Each element of the MLP input (i.e. of ˜ z t -1 and of 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 i t and update target u i t , 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].

The latent variables are used on each timestep to make a prediction about the target, y , using a separate 'Choice MLP':

We train the model end-to-end by gradient descent to minimize the total loss function:

where y are supervised targets, L softmax is a softmax cross-entropy loss. The hyperparameter β scales the cost associated with passing information through the bottleneck [23, 8]. Different values of β 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 × 10 -3 . We typically trained networks for 10 5 steps, except that networks with very low β ( 10 -3 or 3 × 10 -4 ) required longer to converge and were trained for 5 × 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: QLearning 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.

A

Reward Learning

Task

B

Example

Session

1)InitiateTrial

100%

Left Choices

Probability

Reward

P(O)

P(O)

Rewarded

50%

Unrewarded

or

0%

口

Right Choices

0

100

200

300

400

500

Trial

4.1 Q-Learning Agent

The Q-learning agent [48] maintains two latent variables, Q left and Q 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:

̸

̸

where α 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:

where β is an inverse temperature parameter. We generated a large synthetic dataset consisting of choices made and rewards received by the Q-Learning agent ( α = 0 . 3 , β = 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 ( σ ≪ 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 ( σ ≈ 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 left and Q 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 left and Q 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

b

3

4

Q-Learning Agent

DisRNN Trained on Q-Learning Data

Update Rules

Q-Value

0

0.5

1

Left Choices

Right Choices

Trial

Rewarded

Q-Left

Q-Right

Unrewarded

Update Rules

Left, Unrewarded

Left, Rewarded

Right, Unrewarded

Right, Rewarded

Latent 1

Latent 2

Updated Latent 2

Latent 2

Latent 2

Latent 2

Latent 1

Latent 1

Latent 1

Updated Latent 1

-1

0

1

-1.0

0.0

1.0

-1

0

1

-1.5

-1.0

-0.5

0.0

0.5

1.0

-1

0

1

-1.5

-1.0

-0.5

0.0

0.5

1.0

-1

0

1

-1.5

-1.0

-0.5

0.0

0.5

1.0

-1

0

1

-1.0

0.0

1.0

-1

0

1

-1.5

-1.0

-0.5

0.0

0.5

1.0

-1

0

1

-1.5

-1.0

-0.5

0.0

0.5

1.0

-1

0

1

-1.5

-1.0

-0.5

0.0

0.5

1.0

Latent One

Latent Two

Latent Three

Latent Four

Latent Five

Left Choices

Right Choices

Rewarded

Unrewarded

0

100

200

300

400

500

-2

-1

0

1

2

Trial

Information Bottlenecks

Q-Left

Left, Unrewarded

Left, Rewarded

Right, Unrewarded

Right, Rewarded

Updated Q-Left

Q-Left

Q-Left

Q-Left

Q-Right

Q-Right

Q-Right

Q-Right

Updated Q-Right

a

e

c

Choice

Reward

Latent 1

Latent 2

Latent 3

Latent 4

Latent 5

1

2

5

Latent #

Update MLP

Bottlenecks

Sigma

1

0

Bottleneck

Open

Bottleneck

Closed

1

2

5

Latent #

Latent

Bottlenecks

timestep

t

timestep

t

+1

Q-Left

Reward

Choice

Q-Right

Q-Left update

Reward

Choice

Q-Right update

d

timestep

t

timestep

t

+1

Latent 1

Reward

Choice

Latent 2

Latent 1 update

Reward

Choice

Latent 2 update

f

g

Example Session

Example Session

4

to update them. The first variable, V , tracks 'value'. It keeps a running average of recent rewards, regardless of the choice that preceded them:

where α v is a learning rate parameter.

The second, Θ , is a 'policy' variable that determines the probability with which the agent will select each action. The policy gradient algorithm defines two variables θ left and θ right , each associated with one of the actions and updated using the following rules:

̸

3

Actor-Critic Agent

DisRNN Trained on Actor-Critic Data

0

0.5

1

V

0

0.5

1

Updated V

0

0.5

1

V

Unrewarded

Rewarded

Update Rules

-

2

0

2

-

2

0

2

Theta

Updated Theta

Left, Unrewarded

-

2

0

2

Theta

Left, Rewarded

-

2

0

2

Theta

Right, Unrewarded Right, Rewarded

-

2

0

2

Theta

V

0

0.5

1

Example Session

Latent One

Latent Two

Latent Three

Latent Four

Latent Five

Left Choices

Right Choices

Rewarded

Unrewarded

0

100

200

300

400

500

-

4

-

2

0

2

4

e

Information Bottlenecks

Example Session

Left Choices

Right Choices

Trial

Rewarded

V

Theta

Unrewarded

0

100

200

300

400

500

0

1

- 5

0

5

a

b

Update Rules

Unrewarded

Rewarded

-

1

0

1

-

1

0.0

1

-

1

0

1

Latent Two

Updated

Latent Two

Latent Two

Left, Rewarded

Right, Unrewarded Right, Rewarded

Updated

Latent One

-

1

0

1

-

1

0

1

Latent One

-

1

0

1

Latent One

-

1

0

1

Latent One

Latent Two

0.5

0

1

-

1

0

1

Latent One

d

Choice

Reward

Latent 1

Latent 2

Latent 3

Latent 4

Latent 5

1

2

3

4

5

Latent #

Update MLP

Bottlenecks

Sigma

1 (Closed)

0 (Open)

1

2

3

4

5

Latent #

Latent

Bottlenecks

Θ

V

timestep

t

timestep

t

+1

Reward

Choice

Θ

update

Reward

V update

c

f

g

Latent 1

Latent 2

timestep

t

timestep

t

+1

Reward

Choice

Latent 1 update

Reward

Latent 2 update

Trial

Left, Unrewarded

where π t is the agent's policy defined below (equation 14), α l is a learning rate parameter, and α f is a forgetting rate parameter that we have added to the update rule. This forgetting mechanism causes the policy variables to decay towards zero over time, allowing the agent to perform in an environment with changing reward probabilities. In a setting with just two available actions, θ left and θ right are degenerate: it will always be the case that θ left = -θ right . We therefore define a single latent variable Θ = θ left -θ right for the purposes of visualization. Actions are selected according to:

We visualize these update rules by plotting the updated values of V and of Θ as a function of their initial value and of the trial type (Figure 4b). These update rules are quite different from those of the Q-Learning agent. One key feature is that the update rule for V depends only on reward, not on choice. Another is that the update rule for Θ depends not only on choice and reward, but also on V .

We generated a large synthetic dataset consisting of choices made and rewards received by the Leaky Actor-Critic agent ( α V = 0 . 3 , α L = 1 , α F = 0 . 05 ; 1000 sessions of 500 trials each), and trained a DisRNN to imitate this dataset. Inspecting the learned bottleneck parameters of this network (Figure 4e), we find that two of its latents have open bottlenecks allowing them to pass information forward through time. One of these (latent two) has an open bottleneck to receive input from reward only. We find that this latent is strikingly similar to V from the Leaky Actor-Critic agent, both in terms

of its timecourse (Figure 4c) and update rule (Figure 4d, above). The other (latent one) has open bottlenecks to receive input from choice, reward, and from latent two. This latent is strikingly similar to Θ from the Leaky Actor-Critic Agent, both in terms of its timecourse (Figure 4c) and update rule (Figure 4d, below).

Together with the previous section, these results indicate that DisRNN is capable of recovering the true generative structure of agents with two very different cognitive strategies to perform the two-armed bandit task. We also applied DisRNN to a synthetic behavioral dataset from an artificial performing a very different behavioral neuroscience task: decision-making by accumulation of noisy evidence (Appendix A). We found that it was similarly able to recover the true generative structure of this agent (Figure A1)

5 DisRNN Reveals Interpretable Models of Rat Behavior

RNN

FitsBetter

disRNNs vs BestKnown Cognitive Model

Difference in Quality-of-Fit

(Xval. Norm. Lik; % points)

1%

0%

-1%

-2%

disRNN:Average

LSTM:Average

-3%

disRNN:Individual RatsLSTM:Individual Rats

10-3

10-2

CognitiveModel

BottleneckCost (beta)

Fits Better

Having established that DisRNN can recover true cognitive mechanisms in synthetic datasets with known ground truth, we moved on to consider a large laboratory dataset from rats performing the drifting two-armed bandit task [37]. This dataset has previously been the subject of an intensive human effort at data-driven cognitive modeling, which resulted in a cognitive model consisting of three components each with its own latent variable: a fast-timescale reward-seeking component, a slower perseverative component, and a very slow 'gambler's fallacy' component. This model provided a better fit to the dataset than existing models from the literature, and is currently the best known cognitive model for this dataset [37]. We first asked whether DisRNN could provide a similar quality of fit to this model. We fit DisRNNs using four different values of the parameter β ( 10 -3 , 3 x 10 -3 , 10 -2 , and 3 x 10 -2 ) that controls the relative weight of predictive power and model simplicity in the network's loss function (equation 7). Following [37], we evaluated model performance using two-fold cross-validation: we divide each rat's dataset into even-numbered and odd-numbered sessions, fit a set of model parameters to each, and compute the log-likelihood for each parameter set using the unseen portion of the dataset. To compare results across different animals, we use 'normalized likelihood': e log-likelihood / n trials [12]. For each rat, we compute this both for our DisRNNs and for the cognitive model from [37], and plot the differences in Figure 5. We find that DisRNNs trained with β = 3 x 10 -3 achieved a quality of fit similar to the cognitive model, and that those with β = 10 -3 even narrowly outperformed it (Figure 5, orange curve). As an additional benchmark, we also compared our models to a widely-used RNN architecture, the LSTM [25]. We selected LSTM hyperparameters (network size and early stopping) using subject-level crossvalidation: we divided the dataset into subsets containing only even-numbered and odd-numbered rats, identified the hyperparameters that maximized (session-level) cross-validated likelihood in each subset, and evaluated networks with those hyperparameters in the unseen subset. We found that DisRNNs with β = 10 -3 achieved a very similar quality of fit to these optimized LSTMs. Together these results indicate the DisRNNs, despite their architectural constraints and disentanglement loss,

can achieve predictive performance comparable to standard RNNs and to well-fit human-derived cognitive models.

Finally, we examined the parameters of DisRNNs fit to the complete dataset for each rat. We found that these were typically low-dimensional, with only a small number of latent variables having open bottlenecks, as well as sparse, with each update MLP having only a small number of open bottlenecks. While a thorough characterization of the patterns present in DisRNNs fit to individual rats remains a direction for future work, we present a representative example in Figure 6. This figure shows three DisRNNs with different values of β fit to the dataset from the same example rat. The DisRNN with the largest β of 3 x 10 -2 (left) adopts a strategy that utilizes only one latent variable. This strategy mixes reward-seeking and perseverative patterns into a single update rule (note that fixed points, where the line crosses unity, are more extreme for the rewarded conditions than the unrewarded ones). The DisRNN with a medium β of 10 -2 (middle) adds a second latent variable that is purely perseverative and operates on a considerably slower timescale. Its fast-timescale reward-seeking latent variable, however, shows a similar perseverative signature to the previous network. The network with the smallest β of 3 x 10 -3 (right) retains both the fast reward-seeking and slower perseverative dynamics and adds a third, much slower-timescale, latent. The update rule for this latent depends both on reward and on the value of the perseverative latent. Interestingly, this very slow latent influences the update rules for both the reward-seeking and the perseverative latents. While the best human-derived cognitive model does include a very long-timescale component, this component simply influences choice rather than modulating the update rules of the other components. Situations where one latent variable modulates the update rule of another do sometimes occur in human-derived models for related tasks [3, 42], though it is not clear whether the pattern seen in this network is consistent with any of these. Characterizing these patterns in more detail, and especially quantifying similarities and differences among different animals, will be an important step for future work.

We also applied DisRNN to a laboratory behavioral dataset from rats performing an evidence accumulation task [7], and found that it was similarly able to identify low-dimensional models with interpretable structure (Figure A2).

ExampleDisRNNs:AllFit toSameExampleRat

Network One:β=0.03

Network Two:β =0.01

NetworkThree:β=0.003

Example Session

Example Session

Example Session

Right Choices

Left Choices*

Unrewarded

Rewarded

100

200

Trials

300

400

500

Latent Four

Latent Three

Latent Five

Latent Two

Latent One

Right Choices

Left Choices ***

Unrewarded

Rewarded

100

200

Trials

300

400

500

Latent Four

Latent One

Latent Two

Latent Three

Latent Five

Right Choices

Unrewarded

Rewarded

100

200

Trials

Mie babaie 1 +4 1412

300

400

500

Latent One

Latent Four

Latent Three

Latent Five

Latent Two

Update Rules

Update Rules

Jpdated Latent One

Left, Unrewarded  Lef, Rewarded Right, Unrewarded Right, Rewarded

ouo

Left, Unrewarded

Left, Rewarded

Right, Unrewarded Right, Rewarded

Updated Latent One

Left, Unrewarded

1Left, Rewarded Right, Unrewarded Right, Rewarded

Update Rule

10

-1.0

0.0

.0

-1.0

1.0

1.0

tent

epdn

Latent One

Latent One

Latent One

Latent One

Latent One

Latent One

Latent One

Latent One

Left

Right

M0

Left

Right

Unrewarded

Rewarded

1.0

-

Latent One

1

Latent One

-

Latent One

Latent One

po

Latent Two

Latent Two

D

Latent Two

Latent Two

Latent Three

1.0

Latent Three

Bottlenecks

Bottlenecks

Bottlenecks

Bottlenecks

Latent

Update MLP

Bottlenecks

Bottieneck

Open

Bottlenecks

Latent

Update MLP

Bottlenecks

Betteneck

Open

Bottlenecks

Latent

Bottlenecks

Update MLP

Bottleneck

Open

4

6 Discussion

In this work, we develop a framework for discovering parsimonious cognitive models directly from behavioral data by fitting recurrent neural networks that contain structural features that encourage them to learn sparse, disentangled representations. Fit to synthetic datasets generated using known cognitive mechanisms, our method accurately recovers the structure of those mechanisms. Fit to laboratory datasets from rats performing a cognitive task, our method reveals models that are relatively simple and human-interpretable while outperforming the best known human-derived cognitive model in terms of predictive accuracy.

While the extent to which a model fit is 'human-interpretable' is of course ultimately subjective, we believe that sparsity and disentangling provide benefits for at least three distinct reasons. The first is that fully understanding a model requires a human expert to inspect the update rules. The smaller the number of latents and fewer inputs to the update rule for each, the less cognitive burden will be placed on that expert, and the more likely they will be able to arrive at a satisfying human intuition about the cognitive mechanism embodied by the model. The second is that goal of discovery is to identify models that human experts will consider to be cognitively plausible. When evaluating classic handcrafted models, many experts agree that, all else being equal, simpler models (smaller number of equations, fewer terms in each equation) are more plausible. The third is that such a model is more likely to be useful for scientific tasks, such as searching for correlates in measurements of neural activity, that involve interacting with finite datasets.

The fit disRNNs have several key features than enable them to be applied immediately to standard cognitive neuroscience workflows [51, 12, 10, 39]. The first is that they generate timestep-bytimestep timecourses for the values of latent variables that play known cognitive roles within the model. These timecourses can be used as predictions about neural activity: if the model's mechanisms are implemented in the brain, then somewhere in the brain there is likely to be a signal that follows a similar timecourse. A second key feature is that disRNN makes explicit the rules by which latent variables are updated. These can be used as predictions about update rules within the brain: if the model's mechanisms are implemented in the brain, then somewhere in the brain there are likely to be similar update rules (e.g. synaptic weights to update the ongoing neural activity, synaptic plasticity to adjust synaptic weights, etc). A third is that a trained disRNN can be used to generate predictions for experiments which alter neural activity, for example by silencing it using optogenetics. Activity within the disRNN can be altered, for example by artificially zero-ing out a particular latent on a particular subset of trials, and the model can be run to generate predictions for behavioral and neural data [38].

An appealing feature of disRNN is that, for the same dataset, models with different complexity can be discovered by training networks with different values of the hyperparameter β . For some applications, it may be best to select the model that achieves the highest cross-validated quality-of-fit. For others, it may be useful to work with a model that is simpler, even though it may achieve lower quality-of-fit. We expect that in practice disRNN will be most useful if researchers fit several copies, at each of several values of β , and inspect their fits to determine which (if any) have learned models that are useful (Appendix B).

The most important limitation of our approach is that there is no guarantee that the recovered model will necessarily correspond to the true cognitive mechanisms used by the brain. Instead, it uses the dataset to reveal a parsimonious hypothesis. This hypothesis that requires evaluation by a human scientist to determine whether, given the rest of what is known about the psychology and biology of the system, it is plausible. If plausible, it may require new experiments to test the predictions that it makes. This limitation is fundamental to any approach that seeks to make inferences about cognitive mechanisms by analyzing behavioral data, as any behavioral dataset will provide only a limited window on neural mechanisms. One view of these systems is that they perform automatic hypothesis generation, without themselves addressing the problem of hypothesis testing.

An additional limitation is that the method, because it requires fitting quite flexible models, likely is applicable only to relatively large-scale datasets. Determining how performance scales with dataset size and exploring methods of improving data efficiency may be important directions for further research. Another direction is exploring the performance of disRNN on behavioral datasets from other types of cognitive neuroscience tasks. A large number of scientifically-important tasks exist which are of similar complexity to those considered here, and we believe that for many of these disRNN would prove able to discover reasonable models. A large number of tasks also exist which are considerably more complex, and we expect that for these disRNN in its current form might struggle, perhaps by requiring prohibitively large datasets.

The ultimate goal in cognitive model discovery is to develop tools which can reveal scientifically valuable models automatically from data. We believe that disRNN represents an important step towards this goal. Its most important test will come from work applying it to novel datasets, for which no model is known, and evaluating whether the structures it learns are useful to researchers.

Acknowledgments and Disclosure of Funding

We would like to thank Kim Stachenfeld and Yu Jin Oh for helpful comments on the manuscript.

References

Alexander A Alemi et al. 'Deep variational information bottleneck'. In: arXiv preprint arXiv:1612.00410 (2016).

Li Ji-An, Marcus K Benna, and Marcelo G Mattar. 'Automatic Discovery of Cognitive Strategies with Tiny Recurrent Neural Networks'. In: bioRxiv (2023), pp. 2023-04.

Timothy EJ Behrens et al. 'Learning the value of information in an uncertain world'. In: Nature neuroscience 10.9 (2007), pp. 1214-1221.

Celia C Beron et al. 'Mice exhibit stochastic and efficient action switching during probabilistic decision making'. In: Proceedings of the National Academy of Sciences 119.15 (2022), e2113961119.

Matthew Botvinick and David C Plaut. 'Doing without schema hierarchies: a recurrent connectionist approach to normal and impaired routine sequential action.' In: Psychological review 111.2 (2004), p. 395.

James Bradbury et al. JAX: composable transformations of Python+NumPy programs . Version 0.3.13. 2018. URL: http://github.com/google/jax .

Bingni W Brunton, Matthew M Botvinick, and Carlos D Brody. 'Rats and humans can optimally accumulate evidence for decision-making'. In: Science 340.6128 (2013), pp. 95-98.

Christopher P Burgess et al. 'Understanding disentangling in beta-VAE'. In: arXiv preprint arXiv:1804.03599 (2018).

Kyunghyun Cho et al. 'On the properties of neural machine translation: Encoder-decoder approaches'. In: arXiv preprint arXiv:1409.1259 (2014).

Greg Corrado and Kenji Doya. 'Understanding neural coding through the model-based analysis of decision making'. In: Journal of Neuroscience 27.31 (2007), pp. 8178-8180.

Greg S Corrado et al. 'Linear-nonlinear-Poisson models of primate choice dynamics'. In: Journal of the experimental analysis of behavior 84.3 (2005), pp. 581-617.

Nathaniel D Daw et al. 'Trial-by-trial data analysis using computational models'. In: Decision making, affect, and learning: Attention and performance XXIII 23.1 (2011).

Nathaniel D Daw et al. 'Cortical substrates for exploratory decisions in humans'. In: Nature 441.7095 (2006), pp. 876-879.

Ben Deverett et al. 'Cerebellar involvement in an evidence-accumulation decision-making task'. In: elife 7 (2018), e36781.

Amir Dezfouli, Richard Nock, and Peter Dayan. 'Adversarial vulnerabilities of human decisionmaking'. In: Proceedings of the National Academy of Sciences 117.46 (2020), pp. 2922129228.

Amir Dezfouli et al. 'Models that learn how humans learn: The case of decision-making and its disorders'. In: PLoS computational biology 15.6 (2019), e1006903.

R Becket Ebitz, Eddy Albarran, and Tirin Moore. 'Exploration disrupts choice-predictive signals and alters dynamics in prefrontal cortex'. In: Neuron 97.2 (2018), pp. 450-461.

Maria K. Eckstein et al. 'Predictive and Interpretable: Combining Artificial Neural Networks and Classic Cognitive Models to Understand Human Learning and Decision Making'. In: bioRxiv (2023). DOI: 10.1101/2023.05.17.541226 . eprint: https://www.biorxiv. org/content/early/2023/05/17/2023.05.17.541226.full.pdf . URL: https: //www.biorxiv.org/content/early/2023/05/17/2023.05.17.541226 .

Matan Fintz, Margarita Osadchy, and Uri Hertz. 'Using deep learning to predict human decisions and using cognitive models to explain deep learning models'. In: Scientific reports 12.1 (2022), p. 4736.

Laurent Girin et al. 'Dynamical Variational Autoencoders: A Comprehensive Review'. In: CoRR abs/2008.12595 (2020). arXiv: 2008.12595 . URL: https://arxiv.org/abs/2008. 12595 .

Anirudh Goyal et al. 'Infobot: Transfer and exploration via the information bottleneck'. In: arXiv preprint arXiv:1901.10902 (2019).

Tom Hennigan et al. Haiku: Sonnet for JAX . Version 0.0.9. 2020. URL: http://github.com/ deepmind/dm-haiku .

Irina Higgins et al. 'Beta-V AE: Learning basic visual concepts with a constrained variational framework'. In: International conference on learning representations . 2017.

Irina Higgins et al. 'Darla: Improving zero-shot transfer in reinforcement learning'. In: International Conference on Machine Learning . PMLR. 2017, pp. 1480-1490.

Sepp Hochreiter and Jürgen Schmidhuber. 'Long short-term memory'. In: Neural computation 9.8 (1997), pp. 1735-1780.

Makoto Ito and Kenji Doya. 'Validation of decision-making models and analysis of decision variables in the rat basal ganglia'. In: Journal of Neuroscience 29.31 (2009), pp. 9861-9874.

MI Jordan. Serial order: a parallel distributed processing approach. technical report, june 1985-march 1986 . Tech. rep. California Univ., San Diego, La Jolla (USA). Inst. for Cognitive Science, 1986.

Hoseok Kim et al. 'Role of striatum in updating values of chosen actions'. In: Journal of neuroscience 29.47 (2009), pp. 14701-14712.

Timothy D Kim et al. 'Inferring latent dynamics underlying neural population activity via neural differential equations'. In: International Conference on Machine Learning . PMLR. 2021, pp. 5551-5561.

Diederik P Kingma and Jimmy Ba. 'Adam: A method for stochastic optimization'. In: arXiv preprint arXiv:1412.6980 (2014).

Diederik P Kingma and Max Welling. 'Auto-encoding variational bayes'. In: arXiv preprint arXiv:1312.6114 (2013).

Brian Lau and Paul W Glimcher. 'Dynamic response-by-response models of matching behavior in rhesus monkeys'. In: Journal of the experimental analysis of behavior 84.3 (2005), pp. 555579.

Daeyeol Lee, Benjamin P McGreevy, and Dominic J Barraclough. 'Learning and decision making in monkeys during a rock-paper-scissors game'. In: Cognitive Brain Research 25.2 (2005), pp. 416-430.

Germain Lefebvre et al. 'Behavioural and neural characterization of optimistic reinforcement learning'. In: Nature Human Behaviour 1.4 (2017), p. 0067.

Niru Maheswaranathan et al. 'Universality and individuality in neural dynamics across large populations of recurrent networks'. In: Advances in neural information processing systems 32 (2019).

Vladimir Mikulik et al. 'Meta-trained agents implement bayes-optimal agents'. In: Advances in neural information processing systems 33 (2020), pp. 18691-18703.

Kevin J Miller, Matthew M Botvinick, and Carlos D Brody. 'From predictive models to cognitive models: Separable behavioral processes underlying reward learning in the rat'. In: bioRxiv (2018).

Kevin J Miller, Matthew M Botvinick, and Carlos D Brody. 'Value representations in the rodent orbitofrontal cortex drive learning, not choice'. In: Elife 11 (2022), e64575.

John P O'Doherty, Alan Hampton, and Hackjin Kim. 'Model-based fMRI and its application to reward learning and decision making'. In: Annals of the New York Academy of sciences 1104.1 (2007), pp. 35-53.

Chethan Pandarinath et al. 'Inferring single-trial neural population dynamics using sequential auto-encoders'. In: Nature methods 15.10 (2018), pp. 805-815.

Lucas Pinto et al. 'An accumulation-of-evidence task using visual pulses for mice navigating in virtual reality'. In: Frontiers in behavioral neuroscience 12 (2018), p. 36.

Payam Piray and Nathaniel D Daw. 'A model for learning based on the joint estimation of stochasticity and volatility'. In: Nature communications 12.1 (2021), p. 6587.

Kazuyuki Samejima et al. 'Representation of action-specific reward values in the striatum'. In: Science 310.5752 (2005), pp. 1337-1340.

Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. 'Learnable latent embeddings for joint behavioural and neural analysis'. In: Nature (2023), pp. 1-9.

Benjamin B Scott et al. 'Sources of noise during accumulation of evidence in unrestrained and voluntarily head-restrained rats'. In: Elife 4 (2015), e11308.

Mingyu Song, Yael Niv, and Mingbo Cai. 'Using Recurrent Neural Networks to Understand Human Reward Learning'. In: Proceedings of the Annual Meeting of the Cognitive Science Society . Vol. 43. 43. 2021.

David Sussillo and Omri Barak. 'Opening the black box: low-dimensional dynamics in highdimensional recurrent neural networks'. In: Neural computation 25.3 (2013), pp. 626-649.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Naftali Tishby, Fernando C Pereira, and William Bialek. 'The information bottleneck method'. In: arXiv preprint physics/0004057 (2000).

Alex H Williams et al. 'Unsupervised discovery of demixed, low-dimensional neural dynamics across multiple timescales through tensor component analysis'. In: Neuron 98.6 (2018), pp. 1099-1115.

Robert C Wilson and Anne GE Collins. 'Ten simple rules for the computational modeling of behavioral data'. In: eLife 8 (Nov. 2019), e49547. ISSN: 2050-084X. DOI: 10.7554/eLife. 49547 . URL: https://doi.org/10.7554/eLife.49547 (visited on 04/23/2021).

A Pulse Accumulation Task

In this appendix, we explore the generality of disRNN by applying it to datasets from a very different class of tasks: decision-making via accumulation of noisy evidence. Like the reward-learning tasks considered earlier, these tasks are heavily studied in behavioral neuroscience and have been the subject of intensive cognitive modeling efforts. We consider specifically the 'pulse accumulation' tasks, which are particularly amenable to computational modeling [7, 45, 14, 41]. In each trial of these tasks, discrete 'pulses' of evidence arrive over time via two parallel streams (typically 'left' and 'right'), and the agent is rewarded for reporting which of the two streams contained a larger total number of pulses.

A.1 Bounded Accumulator Agent

We generated synthetic data from an artificial decision-making agent performing this task by implementing a known strategy: Bounded Accumulation. The Bounded Accumulator agent maintains a single latent variable A , which indicates a running estimate of the difference in pulse counts between the two streams. At the beginning of each trial, A is initialized to 0 and then updated timestep-by-timestep as the pulses arrive according to:

where σ a is an accumulator noise parameter controlling noise that scales with time; σ c is a pulse noise parameter controlling how much sensory noise is associated with each evidence pulse; δ t,t R,L are delta functions are the times of the right and left pulses; η R,L are i.i.d. Gaussian variables drawn from N (1 , 1) ; η a is an i.i.d. Gaussian variable drawn from N (0 , 1) . The dynamics of A are additionally governed by a bound parameter B . If A reaches a value more extreme than ± B , it will remain at this value for the remainder of the trial.

The agent's decision on each trial is based on the final value of A and on a lapse parameter l , which is constrained to lie between 0 and 0 . 5 . With probability 1 -l , the decision is given by the sign of A , and with probability l it is inverted.

We generated a large synthetic dataset consisting of pulses observed and decisions made, using a Bounded Accumulator agent with relatively little noise ( σ a = 0; σ p = 0.01; B =2.9; l =0; 200,000 trials) (Figure A1a,b). We then trained a DisRNN to imitate this dataset. We find that just one of its latents has an open bottleneck, allowing it to pass information through time, and that the timecourse of this latent is strikingly similar to that of A (Figure A1c). We visualize the learned update rules for this latent by probing its associated update MLP, and find that they are strikingly similar to those for A (Figure A1d).

A.2 Rat Behavioral Data

We next considered a large laboratory dataset from rats performing a pulse accumulation task in which the evidence streams consisted of auditory clicks delivered from a pair of speakers, one located to the left and one located to the right of the rat [7]. For the purposes of modeling, we discretized each trial into 10ms timebins, and gave as inputs to the network the integer number of clicks that were delivered from each speaker in each bin. We fit copies of DisRNN to the dataset from each individual rat, with different values of the bottleneck cost β .

We find that these often recover human-interpretable models which capture known features of rat behavior on this task. In the first example shown (Figure A2, left), the model has identified a single latent (green) which follows a timecourse similar to that of an accumulator variable. Unlike in the synthetic agent considered above, this accumulator does not seem to have a bound, and it does seem to have 'decay' dynamics, drifting gradually towards 0 on timesteps that do not contain clicks. In the second example shown (Figure A2, left), the model has identified a similar accumulator latent (green), whose dynamics are now modified by those of a second latent (blue). The value of this second latent decreases sharply in response to clicks (from either speaker) and recovers gradually during timesteps without clicks. While it remains low, the impact of additional clicks on the accumulator variable is attenuated. This pattern is consistent with 'sensory adaptation' dynamics that are often included in cognitive models for tasks of this kind [7].

BoundedAccumulatorAgent

DisRNNFittoBoundedAccumulatorData

a

ExampleTrial

C

Example Trial

Decision

Variable

Left Clicks

三

Latent Activity

2

Left Clicks

11

11

LatentOne

LatentTwo

0

LatentThree

LatentFour

Right Clicks

-2

Right Clicks I1

———

LatentFive

0

10

20

40

50

0

10

20

30

40

50

Timesteps

Timesteps

b

UpdateRule

p

UpdateRule

No Click

Left Click

Right Click

Left+Right Clicks

No Click

Left Click

Right Click

Left+Right Clicks

DecisionVariable

Updated Latent

3

Updated

2

2

One Activity

2

2

1

2

1

0

0

0

0

0

T

-1

-1

1

2

-2

2

2

2

Decision Variable

2

0

2

DecisionVariable

4

-2

Decision Variable

2

0

Decision Variable

2

3

2

Latent One

0

2

3

-2

Latent One

0

2

3

2

Latent One

0

2

3

2

Latent One

0

Activity

Activity

Activity

Activity

Poisson Clicks Task: DisRNNs Fit to Same Example Rat

NetworkOne:Beta=0.001

NetworkTwo:Beta=0.0001

Example Trial

ExampleTrial

Latent Activity

Left Clicks I| II II

LatentOne

Activity

Left Clicks |||IIII

LatentOne

1

Latent Two

LatentTwo

0

LatentThree

Latent

LatentThree

LatentFour

LatentFour

LatentFive

3Right Clicks]|||

-2

LatentFive

4

500

Time (ms)

1000

-4

500

Time (ms)

1000

UpdateRules

No Click

Left Click

Right Click

Left +Right Clicks

Laten

Update Rule

odated

No Click

Left Click

Right Click

Left + Right Clicks

Jpdated Latent

One Activity

Latent One

'Latent One

Latent One

'Latent One?

-1.5

Activity

Activity

Activity

Activity

Updated Latent

Two Activity

2

Latent One

2

-2

Latent One

Q

2

Latent One

0

-2

Latent One

Activity

Activity

Activity

Activity

Latent Two

Latent Two

-2

Latent Two

2

Latent Two

Activity

Activity

Activity

Activity

B Robustness of Fits

In this appendix, we explore the robustness of disRNN fits to different values of the hyperparameter β , which controls the contribution of the information bottleneck cost to the loss function (equation 7). We fit three copies of DisRNN to each of our three synthetic datasets (Q-Learning, Actor-Critic, and Bounded Accumulation) at each of nine values of β ranging from 10 -5 to 10 -1 (Figure A3). For each DisRNN, we measure the cross-validated quality-of-fit (Figure A3, yellow, top row), as well as number of information bottlenecks that are open ( σ < 0 . 3 ) for both latent bottlenecks (blue, middle row) and update bottlenecks (green, bottom row). We compare the number of open latent bottlenecks in the fit DisRNNs to the number of latent variables in the true generative process (Q-Learning and Actor-Critic: two latent variables; bounded Accumulator: One latent variable), and the number of open update bottlenecks to the number of total terms across all update rules in the true generative processes (Q-Learning: four terms; Actor-Critic: five terms; Bounded Accumulator: three terms).

We find for all three datasets that DisRNNs with very high values of β adopted structures simpler than those of the true generative agents (fewer open bottlenecks), and that these DisRNNs earned lower quality-of-fit. We find for all three datasets that DisRNNs with very low values of β adopted structures more complex than those of the true generative agents (more open bottlenecks), and that these typically earned a high quality-of-fit. In between the extremes, DisRNNs with medium values of β adopted structures similar to those of the true generative agents. These DisRNNs could reliably be identified as those with the simplest structures (largest β ; fewest open bottlenecks) among the set of networks with high quality-of-fit. Manual inspection revealed that all DisRNNs with the correct number of open bottlenecks also adopted the correct model structure and qualitatively well-matched update rules (data not shown).

These results demonstrate that an approach of fitting multiple copies of DisRNN with different hyperparameters is able to reliably identify the true model structure in each of our synthetic datasets. They also highlight the fact that this approach results in a spectrum of models adopting different tradeoffs between simplicity and quality-of-fit. We expect that DisRNN will be most useful in practice if researchers fit several models with different values of β and inspect each of them to determine which (if any) have discovered structures of scientific interest.

Cross-Validated

Normalized

Likelihood

60%

DisRNNsFittoQ-LearningData

Cross-Validated

Normalized

Likelihood

65%

DisRNNsfittoActor-CriticData

Cross-Validated

Normalized

Likelihood,

DisRNNsfittoBoundedAccumulatorData

100%

50%

60%

50%

#Latent

Bottlenecks

5

#Latent

Bottlenecks

5

#Latent

Bottlenecks

5

0

0

0

#Update

Bottlenecks

#Update

Bottlenecks

35

#Update

Bottlenecks

20

20

0

10-5

10-4

10-3

102

101

10-5

10-4

10-3

102

10-1

0

10-5

10-4

103

10-2

10-1

Penaltyscale(beta)

Penaltyscale(beta)

Penaltyscale(beta)

C Cognitive Model Discovery Using Classic RNNs

In this appendix, we explore whether similar results might have been obtained using classic neural network architectures, either by constraining the size of the network [2] or by fitting a larger network and considering only the top few principal components of its activity space. We consider networks using the popular Gated Recurrent Unit (GRU) architecture [9], because related work has reported good performance with this architecture on similar datasets [2, 16] and because having only a single type of recurrent unit makes them simpler to analyze.

First, we explore quality of fit for GRU networks of different sizes, ranging from one to fourteen hidden units, fit to each of the five datasets (table 1). We find that for the synthetic datasets, 'tiny' networks with just one or two recurrent units do provide a quality-of-fit that is only slightly less than that of larger networks. This raises the possibility that these networks might be able to discover the true model structure for these synthetic datasets. For the rat datasets, we find that larger networks provide a meaningfully better quality-of-fit. This indicates that very small networks do not quantitatively fully capture the structure present in the rat datasets, though it does not rule out the possibility that they might nevertheless discover scientifically-useful model structures.

Next, we inspected the fit of the 'tiny' two-unit GRU networks to our three synthetic datasets, examining plots of their update rules (Figure A4). We see that these typically do not have a 1:1 relationship with the true generative latent variables. The exception is the Q-Learning agent, for which two-unit GRUs do often discover a disentangled solution. Solutions found for the Actor-Critic and Bounded Accumulator datasets are fully entangled, with each unit's update dependent on the value of the other unit and on both input variables. We interpret this to mean that very small conventional networks can discover dynamics that recapture those of certain generative processes, but that they do not do so reliably.

Finally, we inspected the fit of larger ten-unit GRU networks to our three synthetic datasets (Figure A5). We summarize their dynamics by plotting update rules for the first two principal components. While some interpretable patterns are apparent, there is still not a 1:1 mapping between PCs and

the latent variables of the generative process. The dynamics are entangled, with each PC's update depending on all inputs and on the value of the other PC.

Taken together, we interpret these results to indicate that conventional neural networks like GRUs can be a viable route to cognitive model discovery in some circumstances, but also that they have important limitations. One limitation is that, while task training ensures that the dynamics they contain are sufficient to solve the task, nothing ensures that all aspects of these dynamics are necessary (they are free to retain null-space dynamics). Another limitation is that, while the number of latent variables can be constrained by limiting network size or by only considering the top few PCs, nothing ensures that the update rules for these variables are sparse, and nothing encourages them to be 'axis aligned', mapping 1:1 onto the true generative dynamics. Note that the 'tiny RNNs' approach of Ji-An et al., [2] introduced additional architectural changes to the networks, which may help address some of these limitations.

difference between the average score for that combination of dataset and network size and a reference size for each dataset (two units for Q-Learning, Actor-Critic, Bandit Rats, and Clicks Rats; one unit for Bounded Accumulator)

GRUTrainedonQ-LearningDataset

GRUTrainedonActor-CriticDataset

Update Rule:Unit One

Update Rule:Unit One

Left, Unrewarded

Left, Rewarded

Right, Unrewarded

Right, Rewarded

Left, Unrewarded

Left, Rewarded

Right, Unrewarded

Right, Rewarded

1.U

0.5

0.5

0.5

Unit Two

0.5

0.5

Unit Two

0.0

0.0

0.0

0.5

0.5

0.5

0.5

0.5

0.5

-0.5

1.0

1.0

1.0

1.0

1.0

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Update Rule: Unit Two

Update Rule: Unit Two

1.0

Left, Unrewarded

1.0

Left, Rewarded

1.0

Right, Unrewarded

1.0

Right, Rewarded

Left, Unrewarded

Left, Rewarded

1.0

Right, Unrewarded

1.0

Right, Rewarded

0.5

0.5

0.5

0.5

Unit One

0.5

0.5

0.5

0.5

Unit One

0.0

0.0

0.5

0.5

0.5

0.5

0.5

0.5

0.5

-1.0

1.0

-1.0

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Previous Activity

GRUTrainedonBounded AccumulatorDataset

Update Rule:Unit One

No Clicks

Left Click

Right Click

Left + Right Clicks

1.0

Unit Two

0.5

0.5

1.0

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Update Rule: Unit Two

No Clicks

Left Click

Right Click

Left + Right Clicks

0.5

0.5

Unit One

0.5

Previous Activity

Previous Activity

Previous Activity

Previous Activity

GRU Trained on Q-Learning Dataset

GRU Trained on Actor-Critic Dataset

Update Rule:PC One

Update Rule: PC One

Left, Unrewarded

Left, Rewarded

Right, Unrewarded

Right, Rewarded

ewarded

Left, Rewarded

Right, Unrewarded

Right, Rewarded

PC Two

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Previous Activity

=1

Previous Activity

Update Rule:PC Two

Update Rule: PC Two

Left, Unrewarded

Leff, Rewarded

Right, Unrewarded

Right, Rewarded

Left, Unrewarded

Left, Rewarded

Right, Unrewarded

Right, Rewarded

auo

Previous Activity

0

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Previous Activity

GRU Trained on Bounded Accumulator Dataset

Update Rule:PC One

No Clicks

Left Click

Right Click

Left + Right Clicks

PC

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Update Rule: PC Two

No Clicks

Left Click

Right Click

Left + Right Clicks

1

Previous Activity

Previous Activity

Previous Activity

Previous Activity

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents an RNN with bottlenecks (both a sampling and a transition bottleneck) and applies the RNN to behavioural data from simple RL models, as well as rodent data. It is shown that the model learns interpretable latents, and ‘rediscovers’ q-learning as well as actor-critic learning and is able to account for the rodent data as well as the current best cognitive models.\n\nSTRENGTHS:\nThis direction is interesting, and relevant to the animal behaviour and neuroscience literature. The technical details are sound.\n\nWEAKNESSES:\n1)\tThe model is bespoke – a separate MLP for each element of z, as well as the bottleneck. It would be helpful to have an empirical understanding of how important these elements are to learn a disentangled RNN (e.g. via ablation). \n\n2)\tIt would also be helpful to understand the benefit of disentanglement, i.e. if you just trained a standard RNN then looked at the first few principle components do they correspond to the disentangled latents?\n\n3)\tThe benefit of this approach is that it can offer an understanding of how animals learn. But there is only one analysis of a single rodent experiment, which does not unveil much beyond existing cognitive models.\n\n4)\tNo analysis of model latent dynamics. Or learning dynamics – does the model go through several stages of understanding? Does this map onto animal behaviour?\n\nQUESTIONS:\nI’m surprised that the models take up to 40 hours to train on a TPU?! The models are pretty small…\n\nSee weaknesses for other questions.\n\nLIMITATIONS:\nSee weaknesses for limitations.\n\nI'm not sure if this counts as a limitation or not, but I found it hard to place this work with other concurrent-ish preprints that are conceptually similar. E.g. Li et al., 2023 (cited in this work) that came out 1-2 months prior which trains a RNN on cognitive tasks (but without disentanglement) and is interpretable, but is somewhat more complete in terms of model analysis and relationships to animal data.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a method for discovering cognitive models automatically by fitting them to behavioral data. They use a recurrent neural network and pass its output through a bottleneck (implemented using variational autoencoders) which is expected to extract relevant cognitive variables. They evaluate the model on synthetic data and on behavioral data from rats performing a reward learning task. The dataset consisted of sequences of binary choices made (left vs. right) and outcomes experienced (reward vs. no reward). The authors evaluated agents based on Leaky Actor-Critic and Q-learning. They used supervised learning and trained the networks to `` imitate this dataset''. The network had a total of five bottlenecks, while the task had two latent variables. After training, the authors found that only two bottlenecks were ``open'' and corresponded to the latent variables. This demonstrated the ability of the network to discover the latent variables.\n\nSTRENGTHS:\nThis is quite an original effort to automatically learn cognitive models from the data. The use of bottlenecks was reasonable and well-justified. The results are consistent with the hypothesis as the model was able to recover the two latent variables and use only the necessary bottlenecks.\n\nWEAKNESSES:\nThe model is fitted only to a single task, which makes it difficult to evaluate whether it would scale and how useful it would be for the broader community. \n\nIt is not clear how robust the approach is - while in the particular setups, VAEs converged to the expected values, it is not known if that would happen for different hyperarameters (for example, what if the number of bottlenecks was different or if learning rate or some other hyperparameter was chosen differently).\n\nQUESTIONS:\nCan you identify other tasks where this approach could be used? \n\nHow should one go about choosing the number of bottlenecks? Should it always be larger than what is expected that the model will need and if yes, how larger?\n\nLIMITATIONS:\nThe authors discussed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors introduce a novel recurrent architecture that potentially learns more interpretable strategies than an LSTM while achieving roughly similar performance in fitting. They investigate the performance of their model on three separate datasets (two synthetic and one rat behavioral dataset), and offer qualitative evidence that the strategies learned by the model are easy to read out by first finding the \"open bottlenecks,\" then looking at their activities. This is an exciting direction and I found the work very interesting and appropriate for NeurIPS. \n\n**Update**\n\nIncreasing my score due to the updates provided by the authors. This is a very intriguing paper and a refreshing approach that could be especially useful for Neuro/Cog scientists.\n\nSTRENGTHS:\n- I like the motivation. The algorithmic approach is also very interesting as an extension of Beta-VAEs to RNNs.\n\n- This is really cool how the open vs. closed-ness of bottlenecks is indicative of whether or not units are being used for solving a task! Sorry for the naive question, but would this provide more insight than if you were to simply regularize units for sparsity? Can you clarify on what this would tell above and beyond that?\n\nWEAKNESSES:\n- I think this is a potentially fantastic direction for cog and neuroscience and interpretability. I just wish there was more validation of the proposed architecture, and stronger evidence that the learned strategies are representative of those used by the rats in that dataset + that the model achieves similar performance as LSTMs as task complexity scales.\n\n- The figures need work. Graphics are too small, there's awkward layouts, and color schemes need more contrast. I will list some detailed comments below.\n\n- I find the Architecture figure very difficult to understand. I have faced the daunting task of making RNN figs in the past so believe me I understand the challenge here. But I don't find the relationship between left and right panels intuitive, understand how this could be used for computing, or what the goal is of this architecture. Just a thought, but maybe something more high-level? Or even consider removing this fig as you do a good job of explaining the model in the text.\n\n- Fig 2b is hard to see. Could you stack the two parts of A on top of each other, then expand B?\n\n- Figure 3 is difficult to understand. The text describing this is super cool and makes sense. But I don't get much from this figure. It feels more like SI to me.\n\n- Figure 4, everything is too small. Am I correct in that 4C/G (and 3C/G for that matter) are cartoons of the data generating process? There's no clear correspondence between them and the data. I would either remove them or figure out a way to make the correspondence clearer with the data.\n\n- Figure 5, The light orange and light blue are hard to read. Also this could be bigger.\n\n- Figure 6, This one describes the most interesting dataset in the paper, but is hardest of all to read! I also don't understand whats going on by looking at it. I see the bottlenecks are changing as you reduce beta. But what does that mean? I know this is a tough ask but you could maybe focus on just one of these models and make it very intuitive to show how it is revealing an insight into the animals' strategies for solving the task.\n\n- I am glad the authors brought up Ji et al in their related work. But I am also confused why they didn't compare to that approach. I assume this work is in progress? If these are networks with a \"very small number of hidden units\" as the authors wrote then it should be straightforward to do.\n\n- The biggest limitation I see here is that the interpretability work is all postdictive. The model is used to explain existing generated datasets which is super cool, then fit to an animal dataset, which is even cooler. But the interpretability of this fit is purely qualitative. This would be a slam dunk if it could be validated experimentally in animal or human neural recordings (predicting neural activity after fitting to behavior) or behavioral data (e.g., identifying and testing biases, potentially).\n\nQUESTIONS:\n- Can you reduce the number of hidden units in the LSTM to roughly match the number of open bottlenecks in the DisRNNs, and get similar performance + interpretability? For interpretability, I guess you could look at the cell state of the circuit? It may not work of course because of the complexity of LSTMs.\n\n- Is there a code release?\n\nLIMITATIONS:\nSee weaknesses. This is tantalizingly close to a great paper. More evidence is needed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a methodology to learn (or system-identify) parsimonious cognitive models directly from data. Specifically, they introduce the idea of disentangled RNNs (DisRNNs). DisRNNs are gated recurrent neural networks with additional bottleneck constraints. These (learnable) bottlenecks bound the information transfer capacity by parametrically controlling the signal-to-noise ratio per latent unit. Moreover, they propose that each latent unit in the DisRNN is updated independently per its own learning rule. The authors demonstrate the effectiveness of DisRNNs on a dynamic two-arm bandit task using synthetic and animal behavioral data.\n\nSTRENGTHS:\nThe authors raise important points about data-driven models usually being inscrutable. System identification is a long-standing problem of interest over the years that several studies have aimed to address. Combining system identification with data-driven approaches is a certainly promising direction. This information-theoretic approach to controlling signal quality through noise induction is interesting. Finally, though speculative, the authors provide a sense of how testable hypotheses for neuroscience can be obtained from their model.\n\nWEAKNESSES:\nMy main concern about this manuscript is its limited scope (in the formulation and experiments). The authors motivate their approach by stating that \"discovering an appropriate model structure ...\" (L21 Pg. 1). However, the model structure of the DisRNN has several carefully chosen inductive biases that closely mimic the model tested in this paper (including the linear update terms and logistic readout). The fact that the DisRNN is able to learn exact parameter specifications is not necessarily surprising. The interpretability aspect of the learned DisRNN also relies heavily on the apriori known ground truth model. Can the authors test this model on other cognitive process models, even within the scope of decision-making? \n\nThe way the DisRNNs are set up, it is unclear to me how they'll scale -- both as a function of the complexity of the underlying model governing data and in terms of parameter and sample complexity. Particularly since each latent dimension has its own associated MLP parameters, it will be very beneficial for the manuscript to have extended numerical evaluations on this front. \n\nWriting style:\nThe general clarity of the article in a few places can be improved. Here are some of my suggestions. I'd encourage the authors to pay more attention to language of this kind throughout the manuscript.\n\nFig. 1 must be improved. The arrow marks (and corresponding colors) do not have a clear legend. The orange text indicates that these lines are modulated by a bottleneck but what does the intersection of orange and blue lines mean? Similarly, the \"Update rule X\" can be depicted as an MLP while clearly showing the dimensionality of $z$. It is unclear if the indices authors are using denote time (which seems to be the case for observations) but not for $z$. \n\nClarity: \"In order to learn an interpretable cognitive model, we encourage sparsity.\" Can the authors expand and clarify why sparsity implies interpretability?\n\nClarity: \"Synthetic datasets from two reinforcement learning agents performing this task..\" It's perhaps better to say that data was generated using ground truth update equations since the \"agents\" themselves here were not trained but rather specified.\n\nQUESTIONS:\nPlease refer to the weaknesses section above.\n\nLIMITATIONS:\nPlease refer to the weaknesses section above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis study introduces \"Disentangled RNNs\" (DisRNNs), a type of interpretable RNN designed with sparse latent variables and simple update rules. The model's interpretability is achieved by utilizing noisy channels both to maintain the state of each latent variable and to read out these states into the update rules. An auxiliary loss function penalizes channels with non-zero Signal-to-Noise Ratio, balancing model fit with sparseness and simplicity. \n\nThe authors trained the DisRNNs directly on behavioral data and then examined the resulting latent representation and learned update rules. First, the authors trained DisRNNs on synthetic behavioral data sampled from either a Q-learning agent or an actor-critic agent, successfully replicating the generating latent variables and their corresponding update rules. Subsequently, the authors trained DisRNNs on real behavioral data from mice performing a dynamic two-armed bandit task. The DisRNNs achieved a competitive cross-validated fit to the behavioral data. Upon analysis, a DisRNN fitted to a specific mouse exhibited a strong resemblance to the best-known human-derived model, with the model's complexity varying based on the weight assigned to the auxiliary loss.\n\nSTRENGTHS:\n1. This manuscript presents a compelling approach to interpretable machine learning. The use of noisy channels is well justified from a computational neuroscience perspective, and the validation of the model on synthetic data from artificial agents is logically sound. Overall, this work may provide a promising foundation for further research in neuroAI using interpretable models.\n\n2. The manuscript is exceptionally well-written. The foundational principles of the reinforcement learning agents are elucidated with the clarity of an outstanding textbook, and the algorithm's description is concise yet clear and accessible.\n\nWEAKNESSES:\n1. Upon reading the paper, I found it somewhat disappointing that the interpretation of the DisRNNs trained on real behavioral data was quite cursory, focusing on a \"typical-best\" example without a systematic analysis of the learned representations across the trained DisRNNs. The study would have been more comprehensive and impactful had it demonstrated how applying this method to behavioral data could facilitate neuroscientific discovery.\n\n2. The manuscript lacks explicit details regarding the selection of the parameter $\\beta$ in section 4 (i.e., experiments with synthetic data). Consequently, the robustness of the procedure to the choice of $\\beta$ remains unclear; it is uncertain whether there was a need for considerable adjustment of $\\beta$ to achieve a good fit or to yield a sensible latent representation. This omission, coupled with the absence of open-source code and data, somewhat undermines my confidence in the applicability of this approach \"out of the box\".\n\nQUESTIONS:\nWhat should guide the selection of $\\beta$ for an unknown system? Is it appropriate to determine $\\beta$ based on cross-validated fit? In essence, how can we interpret the continuum of models obtainable along the bias-complexity tradeoff?\n\nLIMITATIONS:\nThe limitations are adequately discussed.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-26"}