Cognitive Model Discovery via Disentangled RNNs
Authors: Kevin J Miller
Maria K Eckstein
Matthew Botvinick
Zeb Kurth-Nelson
Venue: NeurIPS 2023 poster (2023)
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
Kevin J. Miller
Maria Eckstein Google DeepMind London, UK mar maribackksteinddeepmmddd com
Google DeepMind and University College London
London, UK
kevinjmiller&ddepmmni com
Matthew M. Botvinick
Zeb Kurth-Nelson
Google DeepMind London, UK
Google DeepMind and University College London
London, UK zebx3deepmind.com
botvinicck&deepmind.com
Abstract
Computational cognitive models fundamental tool in behavioral neuroscience. They instantiate in software precise hypotheses about the cognitive mechanisms underlying partieular behavior. Constructing these models is typically difficult iterative process that requires both inspiration from the literature and the creativity of individual researcher. Here, we adopt an alternative approach learn parsimonious cognitive models directly from data. We fit behavior data using recurrent neural network that penalized for carrying information forward 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 reward learning task recision-making task, our method recovers simple and interpretable models that make testable predictions about neural mechanisms.
Introduction
Fitting quantitative cognitive models to behavioral data is a fundamental tool in many areas of psychology and neuroscience 12, To 39]. These models can be viewed as mechanistic hypotheses about the cognitive processes used by the brain. Viewed this way, they act as an explicit software instantiation of particular cognitive hypothesis, and can be used to make precise quantitative predictions about behavioral and neural data. traditional modeling pipeline is for human researcher to iterate over three-step process: first to propose candidate model structure (e.g. Q-learning), second 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) 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) 2]. RNNs are highly expressive and can approximate 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 a "black box": complex system that itself requires further analysis if it is to yield insight into cognitive mechanism.
Here, propose solution that aims 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
g 21
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

These methods encourage networks to learn representations in which each dimension corresponds to a single true factor of variation in the data [3 8.. DisRNN encourages disentangling 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 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 [3]
26 43) and a "pulse accumulation" First, we generate synthetic datasets for the two-armed bandit task using artificial agents with known learning algorithms (Q-Learning and Actor- 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 similar quality of fit to the best known human-deriven cognitive model of this dataset [37] as well as to an unconstrained neural network [16]. Third, we generate synthetic dataset for the decision-making task using an agent with 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 77), and find that it recovers a strategy similar to that of the best known numan--eriver model for this dataset. In all cases, find that DisRNN is able to learn simple human- human-interprette interpretable cognitive strategies which can be used to make predictions about behavioral and neural datasets.
28
preision-making task 7 45 14 4T).
2
Related Work
Our strategy for encouraging networks to adopt disentangled representations is directly inspired by work using variational autoencoders (VAE). Specifically, 3-VAE [23 8. by scaling the KL loss of variational autoencoder's sampling step (which can be viewed requiring information to pass through a Gaussian information bottleneck), often learns sparse, disentangled representations in which each latent variable corresponds to single true factor of variation in the data. While 3 3-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 way that emphasizes interpretability and is appropriate for cognitive model discovery.
The use of recurrent neural networks as cognitive models has a long history [7]) 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
solutions. Within
neuroscience, number of methods exist which attempt to discover interpretable latent dynamics from neural recording data 40 50 29
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 one of a family of dynamic reward learning tasks that have been heavily studied in behavioral neuroscience. This has led to large library of candidate
cognitive models [13 4 17,37 32 11
26
34
37
33 3 422 The development of such models typically
follows theory-first approach, beginning with an idea (e.g. from optimality [13 42). from machine
learning [33), or from neurobiology (34] D. Several studies have fit behavioral dataset using highly flexible classical models 226 37
32
I1I. One of these has attempted use these fits as basis for
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 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.
2

---

Disentangled Recurrent Neural Networks
An RNN trained to match a behavioral dataset containing temporal dependencies [16] 19 46] can be viewed maintaining 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 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 on bottleneck Network Output Target
Choice MLP
9.
Example Information Bottleneck
Supervised Loss
Latents: z! 2
Update MLP1
Bottleneck Input:
Fit Parameters: m.
Bottleneck Output: 2
Sample from Normal(mz, Ø)
Update
2
Zi+3
Observations:
1 Loss
ol....o
These batilenecks cause disentang latents
NormalO,
These bolllenecks
disentangled update rules
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 ă is random sample from Gaussian distribution determined by input and bottleneck parameters m and a.. The bottleneck is associated with KL loss penalizing information transfer.
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 I right) [49]
11 23 24 8 21). The bottlenecks our networks limit information flow using Gaussian noise. multiple information bottlenecks, each of which is noisy channel defined by two learned parameters: a "multiplier" m and noise variance σ. The output on each timestep is sampled from Gaussian distribution:
Our implementation
3 ~ ~~((m22,0)
where x is the scalar input to the bottleneck and is the scalar output. Each bottleneck is associated with loss which penalizes the sampling distribution for deviating from the unit Gaussian:
Lbottleneck (N(mmphn) (M) 1))
where DKL is the Kullback- divergence, 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 = and a 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". zero, and some information will pass through. In experiments, bottlenecks that carrying useful information typically fit m x 1 and a < 1. We refer to these as "open bottlenecks".
all other situations, L3 will be greater than
We use these information bottlenecks to encourage cognitively interpretable models by encouraging two distinct kinds of disentangling. The first disentangling in the latent variables themselves: would like these to capture separable latent processes and to be relatively few in number. We

---

encourage this kind of simplicity by imposing separate bottleneck on each scalar element of the network's hidden state. We refer these as the "latent bottlenecks".
bottleneck,
(3)
The second kind of simplicity is in the update rules for the latent variables. We would like each variable be updated by its own separate rule, and would like each of these rules be simple as possible. We encourage this kind of simplicity by updating each element of the network's hidden state, zi, using a separate learned update rule defined by a separate set of parameters, Each update rule consists of a multilayer perceptron (MLP) which defines multiplicative update to be applied to its corresponding latent variable. These "Update MLPs" have access the values of all elements of the previous timestep's s hidden state, _ (having passed already through its latent bottleneck), as well as to the network's 's current observations o₁. Each element of the MLP input (i.e. of ž and of о₂) must pass through an additional information bottleneck, which we refer to as "Update MLP Bottlenecks". The output of each update MIP is scalar weight we and update target uf, which are used together to update the value of the associated latent variable 2f in manner analogous to the Gated Recurrent Unit (9.
Wip ut MLP -MMPuudttee
(bP update (bottleneck(24- (&.
(4)
= (1 - wf) + wiuf
(5)
The latent variables are used on each timestep to make a prediction about the target, using separate "Choice MLP""
Ú. MLP outpur (Ze)
(6)
We train the model end-to-end by gradient descent to minimize the total loss function:
Lotal Lsftmma (y.
Σ
bottleneck (b)
be bettleensscks
where y are supervised targets, Tsooftmax is softmax cross-entropy loss. The hyperparameter B scales the cost associated with passing information through the bottleneck [23 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 rat (37] or an artificial agent (Q-Learning or Leaky Actor-Critic, see below). Each episode corresponded 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 network consisted the left and right pulse counts 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.
single behavioral
single trial. On each timestep, the observation given to
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 learning rate of 1 We typically trained networks for 10 steps, except that networks with very low 3 (10-3 or 3 x 10-") required longer to converge and were trained for 5 x 10 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 synthetie 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 [3 28 26 37 each trial of this task, the agent selects one of two available actions and then experiences 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 bounded random walk (37]
We generated synthetic datasets from two reinforcement learning agents performing this task: Q- Learning and Leaky Actor-Crite Cator- These agents have markedly different latent variables and update rules, allowing to test whether DisRNN can recover the correct structure of each agent.
Reward Learning Task
Example Session
1) Initiate Trial
2) Select Choice Port
100%
Left Choices
Probaniity
P
P())
Roward
Rewarded Unrewarded
50%
70
0%
Right Choices
300
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 337. 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 not the animal received reward. We will use this session as running example to illustrate the dynamics artificial agents and RNNs.
4.1 Q-Learning Agent
The Q-learning agent [48) maintains two latent variables, Qlett and right Each of these is associated with one of the available actions, and gives running average of recent rewards experienced after taking that action. They are updated according to the following rules:
Qc+i(a:) = (1 - (1-a))r(a:) arr,
Q441(a a,) = Qela aa)
where a is learning rate parameter. We visualize these update rules by plotting the updated value of Q as a function of its initial value and the trial type (Figure Ba). 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₂ Logistic (3(Q,,left Qeighht))
(10)
where 9 is inverse temperature parameter. We generated a large synthetic dataset consisting of choices made and rewards received by the Q-Learning agent (a = 0.3, 3 = 3; 1000 sessions 500 trials each) (Figure 3D). We then trained 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 (a < 1), allowing them 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 (o 1) bottlenecks take near-zero values all timesteps, while the two with open bottlenecks follow timecourses that are strikingly similar to those of Qleft and Qright (albeit centered update rule associated with each latent by probing its associated update MLP (Figure 3). We find that these update rules are strikingly similar to those for Qleft and Qright They reproduce the key feature that each only updated following choices to particular action.
instead of on 0.5; Figure 3D). We visualize the learned
4.2 Leaky Actor-Critic Agent
The Leaky Actor-Critic agent uses a modification of the policy gradient learning rule (Section 2.8). Like the Q-Learning agent, it makes use of two latent variables, but uses very different rules
5

---

Q-Learning Agent
DisRNN Trained Q-Learnng Data
Example Session
Example Session
Left Choices
Left Choioss
Latent One
Q-iah G-Right
Rewarded Unrewarded
Labent Two
Unrowarded
Latent
Right Choices
Right Choiss
Latent Five
Update Rules
Update Rules
Unrewarded
Left, Rewarded
Rewarded
Left, Unrowarded
Lett Rewarded
Right Unrewarded
Rewarded
Updated O-Lah
Upcate Latore
O-ean
12-Pant 13
Lasont 2
Latoni 2
Latent
Latont
Upslated D.Rig
10
Updond Latent
0-Right
Q-Right
G-Rihh
Rriight
Latent 1
Latent
Latent 1
Lanent
Information Bottlenecks
Latent
MLP Bottlenecks
Bottlenecks
Cheinin Rwort
updatt
up
Mevand-
Right
updaoe
Reward
Aweadd
Liton
Litan
Larest
Close
Latert up
Figure 3: DisRNN Recovers Latent Dynamics of Q-Learning. a: Q-Learning agent using the choices and rewards from the example behavioral session (Figure 2 to generate timeseries for Qleft and right b: Visualization of the Q-Leaniing update rules. Each panel visualizes the update of either Qleft (top row) or rrigtt (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, the horizontal. Dashed lines identity. Dependency graph of Q-learning. Q-learning computes an update for each latent (Qiet and &right) on each timestep. The update for Qlet depends on choice and reward (Equation S). The new value of Q left is weighted sum of its old value (dashed line) and the update. Right is similarly updated. d: DisRNN trained on 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 (5) for the first two latents. Learned parameters of the information bottlenecks. Left: Transmission bottlenecks (equation 3) Right: Update MI.P bottlenecks (equation Darker colors indicate open bottlenecks: for example, the update for latent can depend on choice and reward. g: Dependency graph of cognitive model learned by DisRNN.
update them. The first variable, V, tracks "value". It keeps a running average of recent rewards, regardless of the choice that preceded them:
V = (1- - a. + a.Tt
(11)
where Og a learning rate parameter.
The second, O, is a "policy" variable that determines the probability with which the agent will select each action. The policy gradient algorithm defines two variables Otefe and Oright, each associated with one of the actions and updated using the following rules:
(ar) = (1 - (1-aa)0((a)) + cu(r: - Vi) - (((a)) (
(12)
00+1(a # az) (1 ay (z((a) ai (re (()
(13)

---

Actor-Critic Agent
DisRNN Trained Actor-Critic Data
Example Session
d
Example Session
Choloes
One
Rawerrded Uneowarded
Thats
Rewarded Unowoorded
Latent
Latent
Latent Four
Right Choloos
Right Cholows
Latent Fivo
Trial
Update Rules
Update Rules
Unrewarded
Rewarded
Unrowarded
Reworded
Updated Later Two
Updati
Lotent Two
Latent Two
Lefh, Unowwerded
Le0, Rewarded Right. Unowarded Right. Resserded
Left, Unrowarded Loft. Rewarded Right, Unrrwanted Right, Reserded
Updated Latent Onna
Labend Two
Updated
Thets
Theta
Thets
Thets
Lasent Oni
Labent One
Lutent One
Latand One
Information Bottlenecks
g
Update MIP
Bottlesocks
lenerks
Latent
Chic Reward
Latene
Latue
Chrlke
Reward
Lalent 2
V ubdato
Rewant
Latent 2 update
Rewerd
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 and e. b: Visualization update equations. The update for depends reward but is independent of
of the Actor-Crite choice. The update for θ depends on choice and reward, as well as the value of V. c: Dependency graph of Actor- Critic learning, d: DisRNN trained 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 33. The first two latents have open bottlenecks; the remaining three latents have closed bottlenecks. Right: Update MLP bottlenecks (equation g: Dependency graph of DisRNN trained on Actor-Critic data.
where Ne is the agent's policy defined below (equation 14). Og is learning rate parameter, and a, 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 perform in an environment with changing reward probabilities. In a setting with just two available actions, Oleft and right are degenerate: it will always be the case that leett right We therefore define single latent variable O feet righht for the purposes of visualization. Actions are selected according to:
Gr Logistic(E))
(14)
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-Crite agent (ay = 0.3, az 1, a = 0.05; 1000 sessions of 500 trials each), and trained DisRNN to imitate this dataset. Inspecting the learned bottleneck parameters of this network (Figure (b), 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-Criti agent, both in terms
7

---

of its timecourse (Figure 4b) and update rule (Figure 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 4e) and update rule (Figure fi below).
4
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 perform the two-armed bandit task. We also applied DisRNN to synthetic behavioral dataset from an artificial performing a very different behavioral neuroscience task: decision-maie 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 A
5 DisRNN Reveals Interpretable Models of Rat Behavior
RNN
Model Comparison: disRNNs vs Best Known Cognitive Model
Fin Batter
Difference in Quality-of-Fil (Xval, Norm. Lik: % points)
1%
0%
-1%
-2%
diRNN: Average
LSTM: Average
-3%
diRNN: Individual Rats
LSTM: Individual Rats
10-3
10-2
Cognitive Model
Better
Bottleneck Cost (beta)
Figure 5: Model Comparison on Rat Behavioral Dataset. Quality of model fit (cross-validated
normalized likelihood) relative the human-derived cognitive model from [37] for DisRNN (orange) and LSTM (blue). The x-axis for DisRNN represents different values of the hyperparameter 3, which controls the tradeoff between model simplicity and predictive performance. LSTM hyperparameters were selected by subject-level cross-validation. Error bars indicate standard error over N = 20 rats.
Having established that DisRNN can recover true cognitive mechanisms in synthetic datasets with known ground truth, we moved on to consider 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 data-driven cognitive modeling, which resulted in cognitive model consisting of three components each with its own latent variable: a fast-timescale reward-seeking component, slower perseverative component, and very slow "gambler's '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 3x10-1 3, 10 2 and 3x10 2) that controls the relative weight of predictive power and model simplicity in the network's 's loss function (equation 77). Following (37). we evaluated model performance using two-fold cross-validation: we divide each rat's dataset into even-nembered and odd-numbert sessions, fit 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": llo-- leglliiklihood//nn /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 = 3x10-3 achieved a quality of fit similar to the cognitive model, and that those with 3 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 cross- validation: we divided the dataset into subsets containing only weennnumereed and odd-numbered rats, identified the hyperparameters that maximized (session-level) cross-validated likelihood each subset, and evaluated networks with those hyperparameters in the unseen subset. We found that DisRNNs with 3 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,
8

---

can achieve predictive performance comparable to standard RNNs and to well-fit Humann-dervve 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 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 thorough characterization of the patterns present in DisRNNs fit to individual rats remains direction for future work, we present a representative example in Figure 6 This figure shows three DisRNNs with different values of [3 fit to the dataset from the same example rat. The DisRNN with the largest 3 of 3x10 (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 medium 3 of 10-2 (middle) adds a second latent variable that is purely perseverative and operates on 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 3 of 3x10-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 seeuing and the perseverative latents. While the best human-derived cognitive model does include very long-timi -imessale 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 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
Example DisRNNs: All Fit Same Example Rat
Network One: f= 0.03
Natworks Two: ( 0.01
Network Threa: A 0.003
Fac role Sendion
Ead Seesion
Faaople Session
Pay Cholse
Updera
Usdar
righ Rowvebed
Ros aoada
See Setfenaco
Figure 6: DisRNN Trained on Rat Datasets. Examples of fit DisRNN networks with different values of the hyperparameter 3 controlling the tradeoff between simplicity and predictive power.
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 cognitive task, our method reveals models that are relatively simple and human-interpretablle while outperforming the best known human-derived cognitive model in terms of predictive accuracy.
9

---

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 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 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-by- timestep 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 's mechanisms are implemented in the brain, then somewhere in the brain there are likely to be similar update rules (c.g. synaptic weights to update the ongoing neural activity, synaptic plasticity to adjust synaptic weights, etc). A third is that 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 particular latent 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 B3. For some applications, it may be best to select the model that achieves the highest cross-validatt it may be useful to work with model that is simpler, even though it may achieve lower quality-oft -otty--offff fit. We expect that in practice disRNN will be most useful if researchers fit several copies, at each of several values of B. and inspect their fits determine which (if any) have learned models that are useful (Appendix B).
quality-of-fit For others,
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 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 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 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 ceerntfficallyimportte 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 considerably more complex, and we expect that for these disRNN in its current form might struggle, perhaps by requiring prohibitively large datasets.
tasks exist
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.
10

---

Acknowledgments and Disclosure of Funding
We would like to thank Kim Stachenfeld and Yu Jin Oh for helpful comments the manuscript.
References
[1] Alexander A Alemi et al. "Deep variational information bottleneck". In: arXiv preprint :tti::/612.004100 (2016).
[2] 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.
[3] Timothy EJ Behrens et al. "Learning the value of information in uncertain world". In: Nature neuroscience 10.9 (2007), pp. 1214-1221.
[4] Celia C Beron et al. "Mice exhibit stochastic and efficient action switching during probabilis- tic decision making". In: Proceedings of the National Academy of Sciences 119.15 (2022),
e21 1111611199
[5] Matthew Botvinick and David C Plaut. "Doing without schema hierarchies: recurrent connec- tionist approach to normal and impaired routine sequential action." " In: Psychological review 111.2 (2004), p. 395.
[6] James Bradbury et al. JAX: composable transformations of Python+ programs. Ver- sion 0.3.13. 2018. URL: http: http:://gihhbb wtt/ comm/oogle/jax
[7] Bingni W Brunton, Matthew M Botvinick, and Carlos D Brody. "Rats and humans can optimally accumulate evidence for decision decision-making" makinn"". In: Science 340.6128 (2013), pp. 95-98. In: arXiv preprint
[8] Christopher P Burgess al. "Understanding disentangling in arXiv:1804.0333 (2018).
[9] Kyunghyun Cho et al. "On the properties of neural machine translation: Encoder-decoder
approaches". In: arXiv preprint arXiv: 1259 (2014).
[10] 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-81 810.
[11] Greg Corrado et al. monnearronninearrare Poisson models of primate choice dynamics". In: Journal of the experimental analysis of behavior 84.3 (2005). pp. 581-617.
[12] 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).
[13] Nathaniel Daw al. "Cortical substrates for exploratory decisions in humans". In: Nature 441.7095 (2006), pp. 876-879.
[14] Ben Deverett al. "Cerebellar involvement in an evidence-accumulltion decision-making task". In: elife 7 (2018), e36781.
[15] Amir Dezfouli, Richard Nock, and Peter Dayan. "Adversarial vulnerabilities of human decision- making". In: Proceedings of the National Academy of Sciences 117.46 (2020), pp. 29221- 29228.
[16] 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.
[17] R Becket Ebitz, Eddy Albarran, and Tirin Moore. "Exploration disrupts choice preecc-pre signals and alters dynamics in prefrontal cortex". In: Neuron 97.2 (2018), pp. 450-461.
[18] Maria K. Eckstein et al. "Predictive and Interpretable: Combining Artificial Neural Networks and Classic Cognitive Models to Understand Human Learning and Decision Making". In:
www biorxiv bioRxiv (2023). DO1: 10. 1101/ 0 2023 05. 17 541226 eprint: https: :/www.
org/ htg/content// early/ 2023 / 17 2023, 05. 17 541226 full pdf URL: https
//www.biorxix org/ www.oooottnttcccccc wwweenottttttttetnnnneeeeeeeeeee htr///e21lyyyy213 00/////5/11/// 005//7//222... 05. 17 541222
[19] Matan Fintz, Margarita Osadchy, and Uri Hertz. "Using deep learning predict human decisions and using cognitive models to explain deep learning models". In: Scientific reports
12.1 (2022). p. 4736.
[20] Laurent Girin et al. "Dynamical Variational Autoencoders: A Comprehensive Review" In:
CoRR abs/2008. 12595 (2020). arXiv: 2008. 12595 URL: 12595
[21] Anirudh Goyal et al. "Infobot: Transfer and exploration via the information bottleneck"". In: arXiv preprint arXiv:1901.10902 (2019).

---

[22] Tom Hennigan et al. Haiku: Sonnet JAX. Version 0.0.9. 2020. URL: http:: com/
deepmind/ dm- haiku
[23] Irina Higgins et al. "Beta-VAL Learning basic visual concepts with a constrained variational
framework". In: International conference learning representations. 2017. [24] Irina Higgins et al. "Darla: Improving zero-shot transfer in reinforcement learning". In: Inter- national Conference on Machine Learming. PMLR. 2017, pp. 1480-1490.
[25] Sepp Hochreiter and Jürgen Schmidhuber. "Long short-term memory". In: Neural computation 9.8 (1997), pp. 1735-1780.
[26] Makoto Ito and Kenji Doya. "Validation of macisinn-making models and analysis of decision variables in the rat basal ganglia". In: Journal of Neuroscience 29.31 (2009), pp. 9861-9874.
[27] 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.
[28] Hoseok Kim et al. "Role striatum in updating values of chosen actions". In: Journal of neuroscience 29.47 (2009), pp. 14701-14712.
[29] Timothy D) Kim al. "Inferring latent dynamics underlying neural population activity via neural differential equations". In: International Conference on Machine Learning. PMLR. 2021, pp. 5551-5561.
[30] Diederik P Kingma and Jimmy Ba. "Adam: A method for stochastic optimization", In: arXiv
preprint arXiv: arXiiv:4422 14.66988 (2014).
[31] Diederik P Kingma and Max Welling. "Auto-encoding variational bayes". In: arXiv preprint
arXiv:1312.61 14 (2013).
[32] Brian Lau and Paul W Glimcher. "Dynamic response-by
sssoose-by-response response models of matching behavior in rhesus monkeys". In: Journal of the experimental analysis of behavior 84.3 (2005), pp. 555- 579.
[33] Daeyeol Lee, Benjamin P McGreevy, and Dominic J Barraclough. "Learning and decision making in monkeys during a rock rock-paper paper-scissors scissors game". In: Cognitive Brain Research 25.2 (2005), pp. 416-430.
[34] Germain Lefebvre al. "Behavioural and neural characterization of optimistic reinforcement learning". In: Nature Human Behaviour 1.4 (2017), 0067.
[35] Niru Maheswaranathan et al. "Universality and individuality in neural dynamics across large populations of recurrent networks". In: Advances neural information processing systems 32 (2019).
[36]
Vladimir Mikulik al. "Meta-trained agents implement bayes-optimal agents". In: Advances in neural information processing systems 33 (2020), pp. 18691-18703.
[37] 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).
[38] Kevin Miller, Matthew M Botvinick, and Carlos D Brody. "Value representations in the rodent orbitofrontal cortex drive learning, not choice".. In: Elife 11 (2022), e64575.
[39] John P O' 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 104.1 (2007), pp. 35-53.
[40] Chethan Pandarinath al. "Inferring single-trial neural population dynamics using sequential auto-encoders" " In: Nature methods 15.10 (2018), pp. 805-815.
[41] Lucas Pinto et al. "An accumulation-of-evideeee
task using visual pulses for mice navigating in virtual reality". In: Frontiers in behavioral neuroscience 12 (2018), p. 36.
[42] 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.
[43] Kazuyuki Samejima et al. "Representation of action-specite reward values in the striatum". In: Science 310.5752 (2005), pp. 1337-1340.
[44] Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. "Learnable latent embed- dings for joint behavioural and neural analysis". In: Nature (2023), pp. 1-9.
[45] Benjamin B Scott et al. "Sources of noise during accumulation of evidence unrestrained and voluntarily head-restrained rats". In: Elife 4 (2015), el 1308.
12

---

[46] 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.
[47] David Sussillo and Omri Barak. "Opening the black box: low- www-dimensional dynamics in high- dimensional recurrent neural networks". In: Neural computation 25.3 (2013), pp. 026-649. [48] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press,
2018.
[49] Naftali Tishby, Fernando Pereira, and William Bialek. "The information bottleneck method".. In: arXiv preprini physics/0004057 (2000).
[50] 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.
[51] Robert C Wilson and Anne GE Collins. "Ten simple rules for the computational modeling of behavioral data". In: eLiee 8 (Nov. 2019), e49547. ISSN: 2050-084X. DOr: 10 7554/ eLife URL: https: /tops:://do...o otgp://///oo.../1.//.o (1554/eeife 49547
49547
(visited on 04/23/2021).
13

---

A Pulse Accumulation Task
In this appendix, we explore the generality of disRNN by applying to datasets from a very different class of tasks: decision-making via accumulation of noisy evidence. Like the reward-learnin 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 im- plementing known strategy: Bounded Accumulation. The Bounded Accumulator agent maintains single latent variable A, which indicates running estimate of the difference in pulse counts between the two streams. At the beginning of each trial, A is initialized and then updated timestep-by-tinee bytimmesepp as the pulses arrive according to:
A A. + Jalaa ôt,t, JpTr + de,t,
(15)
where oa is an accumulator noise parameter controlling noise that scales with time; de is a pulse noise parameter controlling how much sensory noise is associated with each evidence pulse; 6t,tn, are delta functions are the times of the right and left pulses; JR.LL are i.idd Gaussian variables drawn from N11 1): 7a an i.i.d. Gaussian variable drawn from (0, 1). The dynamics of A are additionally governed by 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 I, which constrained to lie between 0 and 0.5. With probability 1, L, the decision is given by the sign of A, and with probability/ / it is inverted.
We generated a large synthetic dataset consisting of pulses observed and decisions made, using Bounded Accumulator agent with relatively little noise (o, = O: оμ = 0.01: B=2.9; (=0; 200,000 trials) (Figure bl.. We then trained 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 ATe). 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 Al).
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 33.
We find that these often recover -uman-interpretabbee 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 have a bound, and does seem to have "decay" dynamics, drifting gradually towards 0 on timesteps that do not contain clicks. In the second example shown (Figure [22] 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
14

---

Bounded Accumulator Agent
DisRNN Bounded Accumulator Data
Example Trial
Example Trial
Decision Variable
||I
Letere Actett
Left Clides
IIll
Latent
Latent
Latent
Latent Four
light Chicks II
III EIITI
Latent Five
Right Cleck
Timesisss
Timestops
Update Rule
Update Rule
Click
Left
No Click
Click
Left Rig Cliks
Updorod Lacon: One Artitty
Oecision teribbe
Decidoon Variebe
Decsion Varatle
Dension Varatle
Cne
Latent
Figure A1: DisRNN Recovers Latent Dynamics of Bounded Accumulator. a: Bounded Ac- cumulator agent run on an example trial of the click accumulation task to generate timeseries for the accumulator variable A. b: Visualization of the Bounded Accumulator update equation. A is incremented on timesteps containing a left click only and decremented on trials containing a right click only, unless it has already reached one of the bounds (here =3). in which case it remains at that bound. On timesteps with no click or with clicks on both sides, A remains the same. The dynamics of A beyond the bound are undefined, as these values are never reached c: DisRNN trained on synthetic behavioral dataset generated by an Bounded Accumulator agent. d: Visualization of the learned update rules for the first latent.
Poisson Clicks Task: DisRNNs
Example Beta 0.0001 Example Trial
Network
0.001
Network
Example Trial
Laft Cido
Latent Latant Two
Left Clicsl
I
Latent O.no Latent Latent Latent Five
Latent Three
WWW W
Latent Your
Right Chica
Latent Five
Rigt Dide
(ms)
Update Rules
Letss Right
Lift Right Clds
Update Rule
Rightc
Let: Right Cldss
Laked Ond Actvity
"Laeri One Auadity
Luain Ord Arsvly
"uteri One Aubdyyy
spctee upot
1ood Dcc
Loon Drl Adt N/
1.00000 60
100000 6a0r
Jendy
AdW
AcEvit
Figure A2: DisRNN Trained on Rat Decision-Making Dataset. Examples of fit DisRNN networks with different values of the hyperparameter /3 controlling the tradeoff between simplicity and predic- tive power
B Robustness of Fits
In this appendix, we explore the robustness of disRNN fits to different values of the hyperparameter 3, 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 3 ranging from 10 10 (Figure A3}. For each DisRNN, we measure the cross-validated quality-of-fi -f- (Figure A3 yellow, top row), as well as number of information bottlenecks that are open (a < 0.3) for both latent bottlenecks (blue, middle row) and update bottlenecks (green, bottom row). We compare the number of open latent bottlenecks the fit DisRNNs to the number of latent variables the true generative process (Q-Learning and Actor-Critie 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-Lennnng:: four terms; Actor Critic: five terms; Bounded Accumulator: three terms).
15

---

We find for all three datasets that DisRNNs with very high values of 3 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 (3 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 |9 adopted structures similar to those of the true generative agents. These DisRNNS could reliably be identified those with the simplest structures (largest (3 fewest open bottlenecks) among the set of networks with high quality-of-fi fit. of--fit. Manual inspection revealed that all DisRNNs with the correct number of open bottlenecks also adopted the correct model structure and qualitatively well- cell-matched update rules (data not shown).
These results demonstrate that an approach of fitting multiple copies of DisRNN with different hyperparameters able to reliably identify the true model structure 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
We expect that DisRNN will be most useful practice if researchers fit several models with different values of 3 and inspect each of them to determine which (if any) have discovered structures of scientific interest.
Cowooatere Laetood
DisRNNs Fit to Q-l D-Laannng Dala
Creasta bates Nonokood
DisRNNs fit to Actor-Oriiic Data
DisRNNs to Bounded Accumuletor Data
Crescia
82%
30%
F-3000
10:3
10
-0
10
10
10:3
10
10
10
10:
10
10
13
10:2
10
Ponalty scala (bota)
Panally saale (bota)
Ponaty scale (bota)
Figure A3: Robustness of DisRNN Fits Across Different Values of B. See text for description
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 two recurrent units do provide quality 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, a meaningfully better quality-of-fe -ff-f.. 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
find that larger networks provide
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 [44). 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- and Bounded Accumulator datasets are fully entangled, with each unit's update dependent 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 1:1 mapping between PCs and
16

---

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 be viable route to cognitive model discovery in circumstances, but also that they have important limitations. One limitation is that, while task training ensures that the dynamics they contain 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 Jatent 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 dynamies. 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.
Hidden
Q-Learning
Actor-Critic
Bandit Rats
Accumulator
Clicks Rats
1
-0.72
-1.38
-2.16
-0.36
2
-0.008
3
0.002
0.01
0.34
0.001
0.08
4
0.003
0.02
0.26
0.003
0.07
5
0.002
0.02
0.39
0.004
0.08
6
0.002
0.02
0.38
0.004
0.09
7
0.002
0.02
0.35
0.004
0.08
8
0.002
0.02
0.46
0.004
0.10
9
0.002
0.02
0.44
0.004
0.11
10
0.002
0.02
0.38
0.004
0.05
11
0.002
0.02
0.35
0.004
0.10
12
0.002
0.02
0.51
0.004
0.10
13
0.001
0.02
0.53
0.004
0.10
14
0.001
0.02
0.50
0.004
0.13
Table 1: Model Comparison: GRUs of Different Sizes Difference in cross-validated normalized likelihood (percentage points) between GRU networks with different numbers of hidden units. Score for each synthetic dataset is the average of three networks fit with different random seeds. Score for the rat two-armed bandit dataset is the average three networks for each 20 rats. Score for the rat Poisson clicks dataset is the average of three networks for each of 19 rats. Scores are reported the difference between the average score for that combination of dataset and network size and a reference size for each dataset (two units for Q-Leaming, Actor-Critic, Bandit Rats, and Clicks Rats; one unit for Bounded Accumulator)
17

---

Trained on Q-Learning Dataset
Trained Actor-Critic Dataset
Uns Teep
Never
Updats Rulo:
Updato Ruie Two Deyarder
right Liwwarte
Paght Heearie
und Onee
AAn Cren
GRU Trained Bounded Accumulator Dataset
Updase Ruri: UnB One
Figure A4: Update Rules for Two-Unit GRUs Fit Synthetic Datasets. The update rules for the GRU fit to the Q-Learning dataset are strikingly similar to the true update rules of the Q-Learning agent (Figure 3). The update rules for GRUs fit to the Actor-Critic and Bounded Accumulator datasets quite different from the true update rules of these agents.
GRU Trained on Q-Learning Dataset
GRU Trained Actor-Critic Dataset
Updase Rule: One
Macaat Rado Oree
PC
PC
Updato Rale: PC Two
Upcats Rute: PC Two
PC G
PC
GRU Trained on Bounded Accumulator Dataset
No
Hii Chule
PC Twe
HIMM Chate
PC Ore
Figure A5: Update Rules for Top Two PCs of Large GRUs to Synthetic Datasets. For all three networks, the update rules for the top two PCs are quite different from the true update rules of the corresponding generative agents.
18


Reviews:
[{"review_id": "iJ3q0tApOm", "reviewer": "Reviewer_M5Mx", "summary": "This 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.", "questions": "What 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?", "limitations": "The limitations are adequately discussed.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses**\n> 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\nWe agree with the reviewer that a key future direction for this line of work will be to analyze in detail the fits of the disRNN to the laboratory datasets, and to relate them in detail to existing cognitive models as well as to neuroscientific data. In the current manuscript, we hope to establish that applying this method to behavioral data results in models that are suitable for standard neuroscientific workflows that currently rely on handcrafted models. For example, the timecourses of their latent variables might be used as regressors for analyzing neural activity, or they might be run to generate synthetic datasets that make predictions about animal behavior in new situations. We will revised the manuscript to be more clear about this. \n\n> 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\nWe have added an analysis sweeping a range of $\\\\beta$s for each of our synthetic data fits (rebuttal pdf, bottom row). We find that good predictive fit can be found over a wide range of values of $\\\\beta$, and that sensible latent representations can be found at the largest values of $\\\\beta$ which also produce good fit. We will include this analysis in our revised manuscript, and also revise the text to be more clear about this. \n\nWe promise to open-source our code and synthetic datasets as soon as possible, and definitely well in advance of the conference. The laboratory datasets we used are already publicly available.\n\n**Questions**\n> What should guide the selection of Beta for an unknown system? Is it appropriate to determine based on cross-validated fit? In essence, how can we interpret the continuum of models obtainable along the bias-complexity tradeoff?\n\nFor our synthetic datasets, ground-truth involved only a small number of latent variables, and a procedure of selecting from among the models with good predictive performance the one with the smallest number of open bottlenecks would have been sufficient to identify it. For discovering cognitive models using laboratory datasets, we expect that the best $\\\\beta$ will depend on the scientific use-case of the model, and that it may sometimes be useful to consider a model which does not achieve the best cross-validated fit, for example because the disRNN has been fit to a large behavioral dataset, but its latents are being used as regressors for a much dataset of neural recordings.\n\nAlthough this does mean that disRNN can produce multiple models from the same dataset, it is worth noting two things: 1) Over several orders of magnitude of beta, it is typical that only a handful of distinct models emerge, and they are often closely related, for example with one new latent appearing as beta is reduced. 2) These models can often be thought of as different levels of resolution on the cognitive mechanism. Each can be useful depending on the level of resolution needed for the research question.\n\nIf the fit disRNN is to be used as a cognitive model, a key question is whether the mechanistic claims it embodies are psychologically and biologically plausible about the system being studied. Evaluating this necessarily requires considering not only predictive performance and model simplicity, but also expert domain knowledge. We expect that disRNN will be most useful in a workflow that involves fitting different networks to discover a number of different models,  then applying domain expertise as a filter to identify which (if any) are cognitively plausible and practically useful for the user’s current scientific goals.\n\nWe have expanded our discussion of these issues in the manuscript. \n"}, {"review_id": "p6KU1Wa3JJ", "reviewer": "Reviewer_9K9A", "summary": "The 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.", "questions": "Please refer to the weaknesses section above.", "limitations": "Please refer to the weaknesses section above.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The 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.", "weaknesses": "My 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses**\n> My 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\nWe thank the reviewer for pointing out this concern. We agree that the multiplicative update rules used in the Q-Learning agent and by the Actor component of the Actor-Critic are very similar to the multiplicative update we built into disRNN, and that their logistic decision rule is similar in form to the softmax cross-entropy loss function. We have now added a new set of experiments training the same disRNN architecture on synthetic datasets generated by a bounded accumulation agent. This agent shares neither of these features: the update rule for its decision variable is additive, with a nonlinearity at the sticky bound; its choice rule is binary with a lapse rate. We find that disRNN is able to recover the structure of this agent. \n\nWe hope that this goes some way to reassure the reviewer that the structure of disRNN is suitable for fitting generic dynamical systems. It is of course also possible to construct disRNN using additive update rules $ z_i^{t+1} = z_i^t + \\\\text{MLP}_i (\\\\mathbf{z}^t, \\\\mathbf{o}^t) $, which more-obviously express generic dynamical systems. In informal exploration we found that these required more time to train but ultimately converged on similar solutions. Exploring more rigorously whether there are advantages to one variant or the other may be a useful question for future work. \n\n> The 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\nWe agree with the reviewer that discovering cognitive models that are more complex, for example using behavior from more complex tasks, will likely require larger networks and therefore larger datasets, and that it seems likely that this will prove impractical for some applications, at least without introducing additional regularization. We note that in cognitive neuroscience, a large fraction of the literature focuses on tasks similar in complexity to those we examine here. These tasks are chosen because they are thought to isolate, in an experimentally tractable way, key building blocks of cognition. We believe that disRNN and methods like it will be useful for more complex tasks. But even if they are not, accelerating discovery in domains like trial-by-trial reward learning and like decision-making is in itself an impactful contribution. We have expanded the discussion of these issues in our manuscript. \n\n> Fig. 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\nWe have improved this figure following these suggestions, and those of the other reviewers (rebuttal pdf, top row).\n\n> Clarity: \"In order to learn an interpretable cognitive model, we encourage sparsity.\" Can the authors expand and clarify why sparsity implies interpretability?\n\nWe thank the reviewer for calling our attention to this important point. We have expanded our discussion to address it in detail. The relevant passage now reads:\n\n*Limiting the number of latent variables provides three distinct benefits. The first 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 second is that interpreting a fit disRNN 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 human 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 third 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.*\n\n> Clarity: \"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\nWe thank the reviewer for pointing out this issue of vocabulary. In cognitive science, the term “agent” is frequently used to denote software modules that take “actions”, interacting in closed-loop with an “environment”, regardless of whether these are hand-crafted or themselves the result of machine learning. We do feel this usage is appropriate here, but have added clarifying language to our manuscript in several places to assist readers from other backgrounds."}, {"review_id": "mhr4jYVtQL", "reviewer": "Reviewer_SAvA", "summary": "The 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.", "questions": "- 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?", "limitations": "See weaknesses. This is tantalizingly close to a great paper. More evidence is needed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- 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?", "weaknesses": "- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses**\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\nThank you for these detailed and helpful suggestions!\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\nWe have substantially revised this figure, and believe that it is now much more clear (rebuttal pdf, top row). We welcome additional thoughts and suggestions. \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\nWe have increased the size of this figure\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. 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\nWe have increased the size of these figures. We have removed these cartoons. We agree that they don’t add much that is not already present in the “bottlenecks” plot.\n\n> Figure 5, The light orange and light blue are hard to read. Also this could be bigger.. \n\nWe have increased the contrast and the size of this figure\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\nWe have revised this figure and our description of it in the text, in an attempt to be more clear. \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\nWe agree with the reviewer that this will be an important direction for future work. In informal exploration we have found that LSTMs with a very small number of hidden units typically underperform those with larger numbers of units. It will be important to compare directly using matched models (GRUs and their \"S-GRU\"), and datasets. \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\nWe agree with the reviewer that a key future direction for this line of work will be to analyze in detail the fits of the disRNN to the laboratory datasets, and to relate them in detail to existing cognitive models as well as to neuroscientific data. \n\n**Questions**\n\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\nWe informally explored this in previous work, and found that shrinking the LSTM down to just a few units typically results in a loss of predictive performance (the \"best LSTM hyperparameters\" used in figure 5 have 7-9 hidden units). We also informally explored fitting larger LSTMs and analyzing the dynamics of the first few principal components – this typically reveals retention of some dynamics that are not induced by training but instead leftover from network initialization. Frustration with these approaches is in fact what led us to the disRNN ideas of 1) explicitly penalizing the network for retaining information in its activations that it was not actually using for anything and 2) allowing the update rule for those activations to be free-form, parameterized by a feedforward sub-network, rather than constrained by particular update equations as in an LSTM or GRU. \n\nWe will add new supplemental information investigating this in our datasets\n\n>Is there a code release?\n\nWe promise to open-source our code and synthetic datasets as soon as possible, and definitely well in advance of the conference. The laboratory datasets we used are already publicly available from the authors of the papers in which they were originally reported.\n"}, {"review_id": "6Mu8qbb9II", "reviewer": "Reviewer_Zh1M", "summary": "The 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. ", "questions": "Can 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? ", "limitations": "The authors discussed limitations. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This 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. ", "weaknesses": "The 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). ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses**\n> The 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\nWe have now added similar results from an additional task: sensory decision-making via accumulation of noisy evidence. In this task, as we showed in the RL task, disRNN successfully recovers the structure of a handcrafted cognitive model, and reveals a plausible human-interpretable model when fit to a large rat dataset. This task is representative of a very heavily-studied class of decision-making tasks that is widely used in behavioral neuroscience because it isolates an important cognitive process thought to be a building block of cognition. The two-armed bandit task is likewise representative of a very heavily-studied class of “reward learning” tasks thought to isolate a different building block. We expect the method to be immediately useful for other tasks from these domains and others of similar complexity, make up a large fraction of behavioral neuroscience research. We agree that exploring how it will scale to tasks of much greater complexity, especially in the face of limited dataset sizes, will be an important question for future research. \n\n> It 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\nWe have added simulations exploring different values of the weighting parameters $\\\\beta$ (rebuttal pdf, bottom row), which impacts the complexity of the models that are discovered.  \n\n**Questions**\n> Can you identify other tasks where this approach could be used?\n\nSee response above\n\n> How 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\nFor a network to learn disentangled representations, the number of latents available in the network structure does need to be at least as large as the number of true factors of variability in the generative process: if fewer are available, the network will necessarily either fail to learn some factors, or will entangle information about multiple factors into a single latent. In our simulations we have typically chosen a number of latent variables that was about double the number that we expected the fit model to need. While we have not explored this formally, informal experiments allowing larger numbers of latent variables resulted in networks that ultimately converged to similar solutions, though at the expense of longer training time and larger dataset size requirements.\n"}, {"review_id": "GIh6pfY2yL", "reviewer": "Reviewer_u93Y", "summary": "This 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.", "questions": "I’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", "limitations": "See 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.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "This direction is interesting, and relevant to the animal behaviour and neuroscience literature. The technical details are sound.", "weaknesses": "1)\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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses**\n> The 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\nWe will add an ablation analysis as new supplemental information, removing the disentanglement loss for either the latent bottlenecks only or the update rule bottlenecks only. \n\n> It 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\nWe informally explored this in previous work, and found that examining the dynamics of the first few principal components of activity in an LSTM typically does not reveal disentangled dynamics. One reason for this is that the LSTM typically will retain dynamical modes that are present prior to training, just from its random initialization, because it is under no pressure to unlearn them. Frustration with this approaches is in fact what led us to the disRNN ideas of 1) explicitly penalizing the network for retaining information in its activations that it was not actually using for anything and 2) allowing the update rule for those activations to be free-form, parameterized by a feedforward sub-network, rather than constrained by particular update equations.\n\nWe will add new analysis demonstrating this in our synthetic datasets as new supplemental information. \n\n> The 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\nWe have added analysis of a new experiment: decision-making via accumulation of evidence (rebuttal pdf second and third row).\n\n> No analysis of model latent dynamics. Or learning dynamics – does the model go through several stages of understanding? Does this map onto animal behaviour?\n\nThe question of the model's learning dynamics is an interesting one to explore in future work. We have focused on using disRNN as a tool for discovering the asymptotic dynamics that govern behavior, but it is true that it could itself be considered as a hypothesis about how animals meta-learn these kinds of tasks. This would involve analyzing in detail the trajectories by which disRNN discovers the asymptotic dynamics, and comparing them to animal's learning trajectories. With respect to analysis of latent dynamics of the trained model, this is a central focus of our work: DisRNN facilitates understanding of these dynamics by allowing us to visualize the update rules that govern them. \n\n**Questions**\n\n> I’m surprised that the models take up to 40 hours to train on a TPU?! The models are pretty small…\n\nThis surprised us as well! Typically the networks achieve good predictive performance very quickly (minutes to tens of minutes), but then require a much larger number of training steps to identify disentangled solutions. We expect that suitable schedules of learning rate and bottleneck penalty might substantially reduce training time, and plan to explore this in future work. \n"}]

Author rebuttal:
We thank the five reviewers for their insightful and helpful feedback on our manuscript. Three major concerns stood out to us as raised in similar ways by multiple reviewers. We summarize our responses to these three concerns here. We have also responded separately to each reviewer’s individual concerns. 

1) Reviewers felt that the manuscript did not sufficiently make the case that disRNN is likely to be useful across a variety of different cognitive neuroscience tasks. They raised the possibility that it might contain structural biases that tailor it specifically to the domain of dynamic reward-learning tasks. 

    We have addressed this concern by adding new experiments testing disRNN in a very different domain: decision-making via accumulation of noisy evidence. This domain, like that of reward-learning, is heavily studied in behavioral neuroscience and has been the subject of intensive cognitive modeling efforts. We consider specifically the “Poisson clicks” task (Brunton, Botvinick, and Brody, 2013). In each trial of this task, rats are presented with a series of auditory clicks delivered from a pair of speakers, one to the left and one to the right of the rat, and they are rewarded for reporting which speaker delivered a larger number of clicks. 

    We first considered a synthetic behavioral dataset (click times and choices) generated by a “bounded accumulator” agent performing this task (rebuttal pdf, second row). This agent keeps track of the relative number of clicks on each side, and commits to a decision when the absolute difference crosses a bound, ignoring any clicks that occur after this bound has been crossed. We fit a copy of DisRNN to this synthetic behavioral dataset, and find that it is able to recover the structure of the bounded accumulator agent. 

    We then considered an open-source laboratory dataset from rats performing this task (Brunton, Botvinick, and Brody, 2013). We fit several copies of DisRNN to behavioral data from individual rats, using a variety of different complexity penalties. We find that these often recover human-interpretable models which capture known features of rat behavior on the task. In the examples shown (rebuttal pdf, third row), the green latent is tracking the relative number of clicks on each side, while the blue latent is tracking a short-term sensory adaptation effect (clicks that shortly follow another click have a reduced impact on decision-making) that is known to play an important role in this task. 

2) Reviewers wondered how robust the fits to synthetic data were with respect to various hyperparameters, especially the information penalty \beta. 

    We have added simulations sweeping this hyperparameter in each of the three synthetic datasets (Q-Learning, Leaky Actor-Critic, and Bounded Accumulation) and report both cross-validated quality-of-fit and the number of open information bottlenecks (rebuttal pdf, bottom row). We find in each case that the true model structure (dotted lines) is reliably recapitulated by the networks which are the simplest (smallest number of open bottlenecks) that also achieve good predictive performance (comparable to that of the best models). 


3) While reviewers felt that the clarity of the writing was high, they felt that the clarity of the figures could be improved, especially Figure 1, which is critical for readers to understand as it explains the structure of the model.

    We have substantially improved this figure (rebuttal pdf, top row) and made a number of edits to the other figures as well, relying heavily on the detailed and generous advice of reviewer SAvA. We believe this has greatly improved the clarity of the figures and of the manuscript as a whole. 

Meta-review: There was general agreement that the work described is significant for the brain sciences and for X-ai. The reviewers raised some concerns regarding the experimental validation of the approach including more substantial evidence that the learned strategies are representative of those used by animals. Overall, many of the initial concerns seem to have been well addressed in the rebuttal. The paper received moderate but unanimous support from the reviewers. The AC recommends the paper to be accepted.
Decision: Accept (poster)