SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\pi_t$  is the agent's policy defined below (equation 14),  $\alpha_l$  is a learning rate parameter, and  $\alpha_f$  is a forgetting rate parameter that we have added to the update rule. This forgetting mechanism causes the policy variables to decay towards zero over time, allowing the agent to perform in an environment with changing reward probabilities. In a setting with just two available actions,  $\theta_{\mathrm{left}}$  and  $\theta_{\mathrm{right}}$  are degenerate: it will always be the case that  $\theta_{\mathrm{left}} = -\theta_{\mathrm{right}}$ . We therefore define a single latent variable  $\Theta = \theta_{\mathrm{left}} - \theta_{\mathrm{right}}$  for the purposes of visualization. Actions are selected according to:
$$
a _ {t} \sim \operatorname {L o g i s t i c} (\Theta) \tag {14}
$$
We visualize these update rules by plotting the updated values of  $V$  and of  $\Theta$  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  $\Theta$  depends not only on choice and reward, but also on  $V$ .
We generated a large synthetic dataset consisting of choices made and rewards received by the Leaky Actor-Critic agent ( $\alpha_{V} = 0.3$ ,  $\alpha_{L} = 1$ ,  $\alpha_{F} = 0.05$ ; 1000 sessions of 500 trials each), and trained a DisRNN to imitate this dataset. Inspecting the learned bottleneck parameters of this network (Figure 4e), we find that two of its latents have open bottlenecks allowing them to pass information forward through time. One of these (latent two) has an open bottleneck to receive input from reward only. We find that this latent is strikingly similar to  $V$  from the Leaky Actor-Critic agent, both in terms
of its timecourse (Figure 4c) and update rule (Figure 4d, above). The other (latent one) has open bottlenecks to receive input from choice, reward, and from latent two. This latent is strikingly similar to  $\Theta$  from the Leaky Actor-Critic Agent, both in terms of its timecourse (Figure 4c) and update rule (Figure 4d, below).
Together with the previous section, these results indicate that DisRNN is capable of recovering the true generative structure of agents with two very different cognitive strategies to perform the two-armed bandit task. We also applied DisRNN to a synthetic behavioral dataset from an artificial performing a very different behavioral neuroscience task: decision-making by accumulation of noisy evidence (Appendix A). We found that it was similarly able to recover the true generative structure of this agent (Figure A1)
5 DisRNN Reveals Interpretable Models of Rat Behavior
Figure 5: Model Comparison on Rat Behavioral Dataset. Quality of model fit (cross-validated normalized likelihood) relative to the human-derived cognitive model from [37] for DisRNN (orange) and LSTM (blue). The x-axis for DisRNN represents different values of the hyperparameter  $\beta$ , 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 a large laboratory dataset from rats performing the drifting two-armed bandit task [37]. This dataset has previously been the subject of an intensive human effort at data-driven cognitive modeling, which resulted in a cognitive model consisting of three components each with its own latent variable: a fast-timescale reward-seeking component, a slower perseverative component, and a very slow "gambler's fallacy" component. This model provided a better fit to the dataset than existing models from the literature, and is currently the best known cognitive model for this dataset [37]. We first asked whether DisRNN could provide a similar quality of fit to this model. We fit DisRNNs using four different values of the parameter  $\beta$  ( $10^{-3}$ ,  $3\mathrm{x}10^{-3}$ ,  $10^{-2}$ , and  $3\mathrm{x}10^{-2}$ ) that controls the relative weight of predictive power and model simplicity in the network's loss function (equation 7). Following [37], we evaluated model performance using two-fold cross-validation: we divide each rat's dataset into even-numbered and odd-numbered sessions, fit a set of model parameters to each, and compute the log-likelihood for each parameter set using the unseen portion of the dataset. To compare results across different animals, we use "normalized likelihood":  $e^{\log \text{-likelihood / n}}$  trials [12]. For each rat, we compute this both for our DisRNNs and for the cognitive model from [37], and plot the differences in Figure 5. We find that DisRNNs trained with  $\beta = 3\mathrm{x}10^{-3}$  achieved a quality of fit similar to the cognitive model, and that those with  $\beta = 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 even-numbered and odd-numbered rats, identified the hyperparameters that maximized (session-level) cross-validated likelihood in each subset, and evaluated networks with those hyperparameters in the unseen subset. We found that DisRNNs with  $\beta = 10^{-3}$  achieved a very similar quality of fit to these optimized LSTMs. Together these results indicate the DisRNNs, despite their architectural constraints and disentanglement loss,
can achieve predictive performance comparable to standard RNNs and to well-fit human-derived cognitive models.
Finally, we examined the parameters of DisRNNs fit to the complete dataset for each rat. We found that these were typically low-dimensional, with only a small number of latent variables having open bottlenecks, as well as sparse, with each update MLP having only a small number of open bottlenecks. While a thorough characterization of the patterns present in DisRNNs fit to individual rats remains a direction for future work, we present a representative example in Figure 6. This figure shows three DisRNNs with different values of  $\beta$  fit to the dataset from the same example rat. The DisRNN with the largest  $\beta$  of  $3 \times 10^{-2}$  (left) adopts a strategy that utilizes only one latent variable. This strategy mixes reward-seeking and perseverative patterns into a single update rule (note that fixed points, where the line crosses unity, are more extreme for the rewarded conditions than the unrewarded ones). The DisRNN with a medium  $\beta$  of  $10^{-2}$  (middle) adds a second latent variable that is purely perseverative and operates on a considerably slower timescale. Its fast-timescale reward-seeking latent variable, however, shows a similar perseverative signature to the previous network. The network with the smallest  $\beta$  of  $3 \times 10^{-3}$  (right) retains both the fast reward-seeking and slower perseverative dynamics and adds a third, much slower-timescale, latent. The update rule for this latent depends both on reward and on the value of the perseverative latent. Interestingly, this very slow latent influences the update rules for both the reward-seeking and the perseverative latents. While the best human-derived cognitive model does include a very long-timescale component, this component simply influences choice rather than modulating the update rules of the other components. Situations where one latent variable modulates the update rule of another do sometimes occur in human-derived models for related tasks [3, 42], though it is not clear whether the pattern seen in this network is consistent with any of these. Characterizing these patterns in more detail, and especially quantifying similarities and differences among different animals, will be an important step for future work.
We also applied DisRNN to a laboratory behavioral dataset from rats performing an evidence accumulation task [7], and found that it was similarly able to identify low-dimensional models with interpretable structure (Figure A2).
Example DisRNNs: All Fit to Same Example Rat
Figure 6: DisRNN Trained on Rat Datasets. Examples of fit DisRNN networks with different values of the hyperparameter  $\beta$  controlling the tradeoff between simplicity and predictive power.
In this work, we develop a framework for discovering parsimonious cognitive models directly from behavioral data by fitting recurrent neural networks that contain structural features that encourage them to learn sparse, disentangled representations. Fit to synthetic datasets generated using known cognitive mechanisms, our method accurately recovers the structure of those mechanisms. Fit to laboratory datasets from rats performing a cognitive task, our method reveals models that are relatively simple and human-interpretable while outperforming the best known human-derived cognitive model in terms of predictive accuracy.
While the extent to which a model fit is "human-interpretable" is of course ultimately subjective, we believe that sparsity and disentangling provide benefits for at least three distinct reasons. The first is that fully understanding a model requires a human expert to inspect the update rules. The smaller the number of latents and fewer inputs to the update rule for each, the less cognitive burden will be placed on that expert, and the more likely they will be able to arrive at a satisfying human intuition about the cognitive mechanism embodied by the model. The second is that goal of discovery is to identify models that human experts will consider to be cognitively plausible. When evaluating classic handcrafted models, many experts agree that, all else being equal, simpler models (smaller number of equations, fewer terms in each equation) are more plausible. The third is that such a model is more likely to be useful for scientific tasks, such as searching for correlates in measurements of neural activity, that involve interacting with finite datasets.
The fit disRNNs have several key features than enable them to be applied immediately to standard cognitive neuroscience workflows [51, 12, 10, 39]. The first is that they generate timestep-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 mechanisms are implemented in the brain, then somewhere in the brain there are likely to be similar update rules (e.g. synaptic weights to update the ongoing neural activity, synaptic plasticity to adjust synaptic weights, etc). A third is that a trained disRNN can be used to generate predictions for experiments which alter neural activity, for example by silencing it using optogenetics. Activity within the disRNN can be altered, for example by artificially zero-ing out a particular latent on a particular subset of trials, and the model can be run to generate predictions for behavioral and neural data [38].
An appealing feature of disRNN is that, for the same dataset, models with different complexity can be discovered by training networks with different values of the hyperparameter  $\beta$ . For some applications, it may be best to select the model that achieves the highest cross-validated quality-of-fit. For others, it may be useful to work with a model that is simpler, even though it may achieve lower quality-of-fit. We expect that in practice disRNN will be most useful if researchers fit several copies, at each of several values of  $\beta$ , and inspect their fits to determine which (if any) have learned models that are useful (Appendix B).
The most important limitation of our approach is that there is no guarantee that the recovered model will necessarily correspond to the true cognitive mechanisms used by the brain. Instead, it uses the dataset to reveal a parsimonious hypothesis. This hypothesis that requires evaluation by a human scientist to determine whether, given the rest of what is known about the psychology and biology of the system, it is plausible. If plausible, it may require new experiments to test the predictions that it makes. This limitation is fundamental to any approach that seeks to make inferences about cognitive mechanisms by analyzing behavioral data, as any behavioral dataset will provide only a limited window on neural mechanisms. One view of these systems is that they perform automatic hypothesis generation, without themselves addressing the problem of hypothesis testing.
An additional limitation is that the method, because it requires fitting quite flexible models, likely is applicable only to relatively large-scale datasets. Determining how performance scales with dataset size and exploring methods of improving data efficiency may be important directions for further research. Another direction is exploring the performance of disRNN on behavioral datasets from other types of cognitive neuroscience tasks. A large number of scientifically-important tasks exist which are of similar complexity to those considered here, and we believe that for many of these disRNN would prove able to discover reasonable models. A large number of tasks also exist which are considerably more complex, and we expect that for these disRNN in its current form might struggle, perhaps by requiring prohibitively large datasets.
The ultimate goal in cognitive model discovery is to develop tools which can reveal scientifically valuable models automatically from data. We believe that disRNN represents an important step towards this goal. Its most important test will come from work applying it to novel datasets, for which no model is known, and evaluating whether the structures it learns are useful to researchers.
We would like to thank Kim Stachenfeld and Yu Jin Oh for helpful comments on the manuscript.
[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-dimensional dynamics in high-dimensional recurrent neural networks". In: Neural computation 25.3 (2013), pp. 626-649.
[48] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[49] Naftali Tishby, Fernando C Pereira, and William Bialek. "The information bottleneck method". In: arXiv preprint 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: eLife 8 (Nov. 2019), e49547. ISSN: 2050-084X. DOI: 10.7554/eLife.49547. URL: https://doi.org/10.7554/eLife.49547 (visited on 04/23/2021).
A Pulse Accumulation Task
In this appendix, we explore the generality of disRNN by applying it to datasets from a very different class of tasks: decision-making via accumulation of noisy evidence. Like the reward-learning tasks considered earlier, these tasks are heavily studied in behavioral neuroscience and have been the subject of intensive cognitive modeling efforts. We consider specifically the "pulse accumulation" tasks, which are particularly amenable to computational modeling [7, 45, 14, 41]. In each trial of these tasks, discrete "pulses" of evidence arrive over time via two parallel streams (typically "left" and "right"), and the agent is rewarded for reporting which of the two streams contained a larger total number of pulses.
A.1 Bounded Accumulator Agent
We generated synthetic data from an artificial decision-making agent performing this task by implementing a known strategy: Bounded Accumulation. The Bounded Accumulator agent maintains a single latent variable  $A$ , which indicates a running estimate of the difference in pulse counts between the two streams. At the beginning of each trial,  $A$  is initialized to 0 and then updated timestep-by-timestep as the pulses arrive according to:
$$
A _ {t + 1} = A _ {t} + \sigma_ {a} \eta_ {a} + \delta_ {t, t _ {r}} \sigma_ {p} \eta_ {r} + \delta_ {t, t _ {l}} \sigma_ {p} \eta_ {l} \tag {15}
$$
where  $\sigma_{a}$  is an accumulator noise parameter controlling noise that scales with time;  $\sigma_{c}$  is a pulse noise parameter controlling how much sensory noise is associated with each evidence pulse;  $\delta_{t,t_{R,L}}$  are delta functions are the times of the right and left pulses;  $\eta_{R,L}$  are i.i.d. Gaussian variables drawn from  $\mathcal{N}(1,1)$ ;  $\eta_{a}$  is an i.i.d. Gaussian variable drawn from  $\mathcal{N}(0,1)$ . The dynamics of  $A$  are additionally governed by a bound parameter  $B$ . If  $A$  reaches a value more extreme than  $\pm B$ , it will remain at this value for the remainder of the trial.
The agent's decision on each trial is based on the final value of  $A$  and on a lapse parameter  $l$ , which is constrained to lie between 0 and 0.5. With probability  $1 - l$ , the decision is given by the sign of  $A$  and with probability  $l$  it is inverted.
We generated a large synthetic dataset consisting of pulses observed and decisions made, using a Bounded Accumulator agent with relatively little noise ( $\sigma_{a} = 0$ ;  $\sigma_{p} = 0.01$ ;  $B = 2.9$ ;  $l = 0$ ; 200,000 trials) (Figure A1a,b). We then trained a DisRNN to imitate this dataset. We find that just one of its latents has an open bottleneck, allowing it to pass information through time, and that the timecourse of this latent is strikingly similar to that of  $A$  (Figure A1c). We visualize the learned update rules for this latent by probing its associated update MLP, and find that they are strikingly similar to those for  $A$  (Figure A1d).
A.2 Rat Behavioral Data
We next considered a large laboratory dataset from rats performing a pulse accumulation task in which the evidence streams consisted of auditory clicks delivered from a pair of speakers, one located to the left and one located to the right of the rat [7]. For the purposes of modeling, we discretized each trial into 10ms timebins, and gave as inputs to the network the integer number of clicks that were delivered from each speaker in each bin. We fit copies of DisRNN to the dataset from each individual rat, with different values of the bottleneck cost  $\beta$ .
We find that these often recover human-interpretable models which capture known features of rat behavior on this task. In the first example shown (Figure A2, left), the model has identified a single latent (green) which follows a timecourse similar to that of an accumulator variable. Unlike in the synthetic agent considered above, this accumulator does not seem to have a bound, and it does seem to have "decay" dynamics, drifting gradually towards 0 on timesteps that do not contain clicks. In the second example shown (Figure A2, left), the model has identified a similar accumulator latent (green), whose dynamics are now modified by those of a second latent (blue). The value of this second latent decreases sharply in response to clicks (from either speaker) and recovers gradually during timesteps without clicks. While it remains low, the impact of additional clicks on the accumulator variable is attenuated. This pattern is consistent with "sensory adaptation" dynamics that are often included in cognitive models for tasks of this kind [7].
Bounded Accumulator Agent
a
DisRNN Fit to Bounded Accumulator Data
Update Rule
d
Figure A1: DisRNN Recovers Latent Dynamics of Bounded Accumulator. a: Bounded Accumulator 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  $\pm 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 a synthetic behavioral dataset generated by an Bounded Accumulator agent. d: Visualization of the learned update rules for the first latent.
Poisson Clicks Task: DisRNNs Fit to Same Example Rat
Network One: Beta = 0.001
Example Trial
Network Two: Beta = 0.0001
Example Trial
Figure A2: DisRNN Trained on Rat Decision-Making Dataset. Examples of fit DisRNN networks with different values of the hyperparameter  $\beta$  controlling the tradeoff between simplicity and predictive power
B Robustness of Fits
In this appendix, we explore the robustness of disRNN fits to different values of the hyperparameter  $\beta$ , 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  $\beta$  ranging from  $10^{-5}$  to  $10^{-1}$  (Figure A3). For each DisRNN, we measure the cross-validated quality-of-fit (Figure A3, yellow, top row), as well as number of information bottlenecks that are open ( $\sigma < 0.3$ ) for both latent bottlenecks (blue, middle row) and update bottlenecks (green, bottom row). We compare the number of open latent bottlenecks in the fit DisRNNs to the number of latent variables in the true generative process (Q-Learning and Actor-Critic: two latent variables; bounded Accumulator: One latent variable), and the number of open update bottlenecks to the number of total terms across all update rules in the true generative processes (Q-Learning: four terms; Actor-Critic: five terms; Bounded Accumulator: three terms).
We find for all three datasets that DisRNNs with very high values of  $\beta$  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  $\beta$  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  $\beta$  adopted structures similar to those of the true generative agents. These DisRNNs could reliably be identified as those with the simplest structures (largest  $\beta$ ; fewest open bottlenecks) among the set of networks with high quality-of-fit. Manual inspection revealed that all DisRNNs with the correct number of open bottlenecks also adopted the correct model structure and qualitatively well-matched update rules (data not shown).
These results demonstrate that an approach of fitting multiple copies of DisRNN with different hyperparameters is able to reliably identify the true model structure in each of our synthetic datasets. They also highlight the fact that this approach results in a spectrum of models adopting different tradeoffs between simplicity and quality-of-fit. We expect that DisRNN will be most useful in practice if researchers fit several models with different values of  $\beta$  and inspect each of them to determine which (if any) have discovered structures of scientific interest.
Figure A3: Robustness of DisRNN Fits Across Different Values of  $\beta$ . See text for description
C Cognitive Model Discovery Using Classic RNNs