Discovering Symbolic Cognitive Models from Human and Animal Behavior
Authors: Pablo Samuel Castro
Nenad Tomasev
Ankit Anand
Navodita Sharma
Rishika Mohanta
Aparna Dev
Kuba Perlin
Siddhant Jain
Kyle Levin
Noemi Elteto
Will Dabney
Alexander Novikov
Glenn C Turner
Maria K Eckstein
Nathaniel D. Daw
Kevin J Miller
Kim Stachenfeld
Venue: ICML 2025 spotlightposter (2025)
Keywords: cognitive modeling
cognitive science
neuroscience
psychology
Primary area: applications->neuroscience_cognitive_science
TL;DR: Automated program discovery powered by FunSearch finds novel, interpretable cognitive models that capture human and animal behavior better than classic cognitive model baselines.

Abstract: Symbolic models play a key role in cognitive science, expressing computationally precise hypotheses about how the brain implements a cognitive process. Identifying an appropriate model typically requires a great deal of effort and ingenuity on the part of a human scientist.
Here, we adapt FunSearch (Romera-Paredes et al. 2024), a recently developed tool that uses Large Language Models (LLMs) in an evolutionary algorithm, to automatically discover symbolic cognitive models that accurately capture human and animal behavior.
We consider datasets from three species performing a classic reward-learning task that has been the focus of substantial modeling effort, and find that the discovered programs outperform state-of-the-art cognitive models for each.
The discovered programs can readily be interpreted as hypotheses about human and animal cognition, instantiating interpretable symbolic learning and decision-making algorithms. Broadly, these results demonstrate the viability of using LLM-powered program synthesis to propose novel scientific hypotheses regarding mechanisms of human and animal cognition.

Discovering Symbolic Cognitive Models from Human and Animal Behavior
Pablo Samuel Castro Nenad Tomasev' Ankit Anand Navodita Sharma' Rishika Mohanta 23
Aparna Dev?
Kuba Perlin
Siddhant Jain' Kyle Levin Noémi Éltetō 14 Will Dabney Alexander Novikov
Glenn C T
Maria K Eckstein'
Nathaniel D Daw 15
Kevin J Miller 16 Kimberly L Stachenfeld 117
Abstract
Better
Difference Normalized
Likelhood Points)
Symbolic models play a key role in cognitive sci- ence, expressing computationally precise hypothe- ses about how the brain implements a cognitive process. Identifying an appropriate model typi- cally requires a great deal of effort and ingenuity the part of a human scientist. Here, adapt FunSearch (Romera-Paredes, et al., 2024), re- cently developed tool that uses Large Language Models (LLMs) in an evolutionary algorithm, to automatically discover symbolic cognitive mod- els that accurately capture human and animal be- havior. We consider datasets from three species performing classic reward-learning task that has been the focus of substantial modeling effort, and find that the discovered programs outperform state-of-the-ar cognitive models for each. The discovered programs can readily be interpreted as hypotheses about human and animal cogni- tion, instantiating interpretable symbolic learning and sreciiion-mainn algorithms. Broadly, these results demonstrate the viability of using LLM- powered program synthesis to propose novel sci- entific hypotheses regarding mechanisms of hu- man and animal cognition.
bodistitat
T
Exsting Mood Setter
Bandiz
Rly Bandit Datase:
Dataset
Dataset
Figure 1. Discovered models outperform human- mod- els. We evaluate the best program discovered by CogFunSearch for each dataset, using average normalized likelihood of the choices made by held-out test subjects, and it to the best existing model from the neuroscience and psychology literature (all p 0.002, signed-rank test) Left: Human dataset and model from Eckstein et (2024) Center: Rat dataset and model from Miller et (2021) Right: Fruit fly dataset from (Mohanta, 2022; Rajagopalan et al., 2023), model from Ito Doya (2009).
1. Introduction
Symbolic cognitive models are used in neuroscience and psychology to instantiate precise, mechanistic hypotheses about the processes used by the brain to control behavior (Daw et al., 2011; Corrado & Doya, 2007: O' D'Doherty et al., 2007). These models readily afford interpretation: inter- nal variables like "prediction error" or "forgetting rate" are meaningful in their mentations in the brain. This ease of interpretation has made symbolic models a key tool in understanding how brains produce behavior.
right, and suggest possible imple-
'qqual contribution 'Google DeepMind "Janelia Farm Re- search Campus, Howard Hughes Medical Institute, Ashburn, VA, USA 'Laboratory of Neurophysiology and Behavior, The Rock- efeller University, New York, NY, USA 'Max Planck Institute for Biological Cyhernetics, Tübingen, Germamy 'Princeton Neu- roscience Institute, Princcton University, Princeton, NJ, USA "Sainsbury Wellcome Centre, University College London, United Kingdom 'Center for Theoretical Neuroscience. Columbia Univer- sity, New York, NY, USA. Correspondence Kevin Miller <kev- injmiller @ deepmind.com>. Kimberly Stachenfeld <stachen- feld@ deepmind.com.
The development of such models has historically been a hypothesis-fire process, in which researchers draw on in- spiration from the literature and individual human creativity to specify a model, and then refine it match the idiosyn- crasies of behavior. A limitation of this approach is that the space of symbolic models is vast, and it is far from cer- tain that the best possible model for a dataset is one that researchers will have considered (Daw et al., 2011; Wilson & Collins, 2019). Indeed, recent work has used comparisons to flexible recurrent neural networks (RNNs; Dezfouli et al.
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Discovering Symbolic Cognitive Models from Bchavior
A.Automated Cognltive Model Discovery with FunSearch
B. Predicting Behavior with Cognitive Model Program
Scarinng Cogn tive Model Programs Data
Datasel
LLM
progremm
state
next state P,,
pregram
Fit parameters
Predict
Dataset Lkelihood
param
Choice ukalhood (R..
next choice prob P.,
optimize params O
choice
00d
Prompt
progrem +
Validate params
score 0
reward
score
Datasel
Predict
Dataset Likellhood
Programs datalase
Figure Overview of CogFunSearch. A: CogFunSearch uses LLMs to evolve Python programs that are predictive of behavioral dataset, and maintains database of candidate programs which are ranked by their model scores. B: Programs Φ receive as input the previous choice es and reward rs. an evolving hidden state he, and a set of trained parameters 0. They output predictive probabilities over the next choice and updated hidden state. C: The model score for programs is computed by fitting parameters subset of the datasct and evaluating its normalized likelihood on a separate validation set.
2019b; Song et al. 2021; Ger et al. 2023; Eckstein et al. 2024) to show that commonly-used symbolic models tend to dramatically underfit their datasets.
in higher-scorine and more interpretable discovered pro- grams. Discovered programs are often surprisingly readable, for example containing informative variable names and com- ments. Several unexpected and intriguing motifs are appar- ent: complex exploration strategies, unconventional value updates, and idiosyncratic patterns of reward- iewarddddddeeddeent choice sequences. Each run of CogFunSearch generates large number of programs, which show a clear tradeoff between quantitative performance and program complexity. While the highest-scoring programs are more complex than our baseline models, possible to identify models for each dataset that are both higher-performing as well as sim- pler. Broadly, these results validate the of LLM-based program synthesis for data-driven discovery of cognitive models, and suggest novel hypotheses about soward-guided learning in humans and animals.
An appealing alternative adopt a data-driven approach, considering very large space of possible models and al- lowing quality-of-fit to guide model selection (Brunton & Beyeler, 2019; Miller et al., 2021). As applied to human and animal behavior, this approach has been most broadly successful with model spaces that are differentiable, such as neural networks (Dezfouli et al., 2019a; Miller et al., 2023).
In this work, we demonstrate method which is able to auto- matically identify symbolic models that are consistent with given dataset. We build on advances in program synthesis, set of methods for automatically identifying computer pro- grams that optimize some objective function. Specifically, we adapt FunSearch (Romera-Paredee et al., 2024), a recent tool that uses Large Language Models (L.LMs) to mutate programs within an evolutionary optimization process. Our approach, which we term "CogFunSearch"" augments Fun- Search with an additional level of optimization. In the outer optimization loop, FunSearch evolves programs (Fig. 2A), while in the inner loop, model parameters are fit to data (Fig. 2C). We apply CogrchSearch to datasets from three species (humans, rats and fruit flies) performing a classic reward-guidee guided decision-maaee focus of substantial human cognitive modeling effort (Fig. 1; Miller et al. 2021; Eckstein et al. 2024; Mohanta 2022; Rajagopalan et al. 2023). Our discovered programs reliably outperform the best numan-discoveree cognitive models of which we are aware.
2. Related Work
Data-driven Cognitive Modeling A number of papers have attempted to invert the traditional theory-first approach to cognitive model building and adopt a data-driven ap- proach. Some of these use statistical tools to identify pat- terns in behavioral data. These include explicitly characteriz- ing patterns in choice and reward sequences using statistical models (Lau & Glimcher, 2005: Sugrue et al., 2004; Ito & Doya, 2009), well adding additional degrees of free- dom to classic cognitive models (Le et al., 2023; Venditto et al., 2024; Roy et al., 2021). One recent paper has at- tempted process of successive model reduction, beginning with flexible statistical models and ending with a model that can be interpreted as a cognitive model (Miller et al., 2021). Related work has applied methods for symbolic regression (Brunton al., 2016: Landajuela et al., 2022) to discover equations describing behavior. One of these
task which has been the
Because CugaucSe uses LLMs for program generation, human-proviee provided information in the prompt
is able use human- and seed program, with more informative prompts resulting
2

---

Discovering Symbolic Cognitive Models from Behavior
(Musslick, 2021) considered synthetic data in a variety of decision-making tasks. Another (LaFollette al., 2023) modeled human behavior in an RL task, but, due to limita- tions of the technique, did not model choices directly but instead asked subjects to report their internal value estimates and modeled the dynamics of these reports.
the chosen action and set of dynamic environment parameters that are independent of choice. Such tasks are often called "Dynamic Multiarmed Bandit Tasks", and have been the focus of a great deal of computational modeling effort (Corrado & Doya, 2007; O'Doherty et al., 2007; Daw et al., 2011; Wilson & Collins, 2019). We selected diverse datasets from three different fruit flies-which share two important properties. First, each dataset is unusually large, consisting of many individ- uals, behavioral sessions, or both. This reduces the risk of overfitting that is inherent in considering large and flexible space of possible models. Second, each dataset has been the focus of past computational modeling efforts, providing a strong numan-discovered baseline against which to compare our discovered models.
rats, and
Interpretable Deep Learning Building on the observa- tion that neural networks often fit data better than symbolic models, several recent papers have introduced methods that attempt to add constraints to deep learning methods to ren- der them more readily interpretable for the purposes of cognitive modeling. One approach is to integrate neural net- works into theoretically-motivatee cognitive architecture (Peterson et al., 2021; Eckstein et al., 2024). Another is to constrain information flow within the network, either using explicit information bottlenecks (Miller et al., 2023) or by simply using a very small network (Ji-An et al., 2023). Like this work, our approach optimizes models directly to data, considering a large model space with structural constraints that we hope will afford interpretability. Importantly, unlike network weights, the model space we search (Python pro- grams) is explicitly symbolic, thereby facilitating analysis and interpretation by humans or indeed LLMs.
Human Dataset (Fig. 3A; Eckstein et al. 2024) considers human participants performing four-alternative task with graded rewards. Participants performed the task online, and indicated their choice on each trial by pressing either 'D', 'F,, 'J' or 'K' on their keyboard. Reward was indicated by displaying an integer number of 'points' between 0 and 100, which subjects were asked to maximize. Available rewards followed independent bounded random walks with additional trial-unique noise. Each participant performed up to five back-to-back sessions of up to 150 trials each. The dataset contains choices from 862 participants performing 4, 1. 34 total sessions and 617,871 total trials.
Program Synthesis for statistical and agent modeling, A large body of previous work explores evolutionary meth- ods for discovering computer programs that optimize some objective. A family of work in machine learning uses ap- proaches like these to discover useful algorithms (Real et al., 2020; Co-Reyes al., 2021; Ellis al., 2023; Chen al., 2024). Within cognitive science, programs have previously been proposed as the representational format for learning and storing complex sequences (Planton et al., 2021), geo- metric shapes (Sablé-Meyer et al., 2022), concepts (Lake et al., 2015), rules (Rule et al., 2024) and behavioral strate- gies (Correa et al., 2024). This body of work highlights that program synthesis can yield programs that show quali- tative commonalities with the average human participants' response patterns. However, to our knowledge, the current work is the first to synthesize programs that maximize the predictive fit of individuals" behavior. Perhaps most simi- lar to our methodology. Li et al. (2024) and Shojace et al. 2025 use L.MSs to propose statistical models, then fit the parameters of these models to data.
Eckstein et al. (2024) performed an extensive comparison of a wide variety of computational cognitive models on this dataset. The model that performed best in this comparison was one we refer to as "Perseverative Forgetting Q-I e-raennnng" (PFQ, see Appendix D and Appendix E for details). We adapt PFQ as the human bandit dataset.
baseline model for the
Rat Dataset (Fig. 3B; Miller et al. 2021) considers rats performing a two-armed bandit task with binary rewards. Rats indicated their choice on each trial by entering one of two available nose ports, which were equipped de- liver small liquid rewards. Reward probabilities followed independent bounded random walks. Rats performed daily sessions of approximately one hour. The dataset contains choices from 20 rats performing 1, 1,946 total sessions and 1,087,140 total trials.
Miller et al. (2021) performed an intensive human process of data-driven model discovery on this dataset. This process began with set of highly flexible models and several iter- ations of successive model reduction: identifying patterns in the fit parameters of a more-flexible model, and propos- ing llss-flexible model that embeds those patterns into its structural assumptions. This resulted in a model that we refer to as Reeward-Seeking/Habi/GGG. Fallacy" (RHG,
3. Methods
3.1. Datasets and Sluman-Discoveed Baseline Models
We consider datasets from humans and other species per- forming reward-learning tasks in which the subject selects repeatedly between several discrete actions and receives a reward whose magnitude and probability depends only
3

---

Discovering Symbolic Cognitive Models from Behavior
3.2. Problem Formulation
Example Session: Human Dataset
Choose Choose
koy
We focus on developing models that predict the choice made by the subject on each trial, using information about the previous choices made and rewards received. We for- malize this as follows: On each trial, the subject selects between n options and receives a reward. Each subject (indexed i) performs one or more sessions (j) each consist- ing of a series of trials (t). For each trial & we consider the discrete choice e E {1..... n made by the subject and the reward received, which may be binary (r E {0. 1), rat and fruit fly datasets) or graded (r [0, 100). hu- man dataset). A sequence of T trials constitutes a ses- sion Sn := {(c1, (cr. rr)} Conditional on the choice, the delivery of rewards is stochastic according to some experimenteer-contollle distribution; in particular, choice c's reward is governed by an independent distribution (Bernoulli or truncated Gaussian, depending on the dataset), with a mean Pi.j,,c that changes over trials according to dataset-dependent process. For each subject i set of Ns sessions are collected: d: := {ss,1, Nt lis with testing conditions reset between sessions. A dataset collects this for ND subjects: D = {dz, d2..... ..... dND The values of T, Ns, and ND will vary based on the dataset in question, as we specify below. See Figure for an illustration of this.
Chrose key
Megnttues
Revard
50
60
202
120
140
Example Session: Rat Dataset
Chooce Chocke
1000%
Beward Pretabilins
50%
1.00
200
300
405
500
660
Triel
Example Session: Fruit Fly Dataset
Thooss A
100%
Probabitas
Reward
0%
:0c
200
300
400
500
Figure 3. Illustration of datasets Example behavioral sessions, showing choices and rewards received, as well as reward contin- gencies, from A the Human Dataset, B the Rat Dataset, C the Fruit Fly Dataset (bottom).
We consider models in the form of functions that receive choice and reward (ct, and produce a probability distri- bution over the next choice Ct+1, denoted as Pr+ Models typically also maintain a hidden state that is updated iter- atively across a sequence of trials, and are finally parameter- ized by a vector 6 which is fixed over trials. If we consider the state as an external object h that the model updates, we may formalize a model as ő(cc, Te, h, 0) + (8r41, h1+1).
see Appendix D and Appendix E for details). Miller al. (2021) also performed model comparison between the RHG model and a wide variety of alternative models from the literature, and confirmed that it provides a better fit to data. We adopt RHG as the -isaa-discovered baseline model for the rat dataset.
Fruit Fly Dataset (Fig. 3C; Mohanta 2022) considers fruit flies performing two-armed bandit task with binary rewards. Flies performed the task in a three-armed "Y-maze" setup in which separate odors could be delivered to each of the three arms (Rajagopalan et al., 2023). Flies indicated their choice on each trial by selecting an odor and walking to the end of the associated arm. Reward was delivered via a brief pulse of red light which activated flies' sugar-sensing neurons (Haberkern et al., 2019). Reward probabilities fol- lowed a random block structure with randomly sampled reward probabilities and block lengths. Each fruit fly per- formed one behavioral session. The dataset contains choices from 347 flies performing 68,000 total trials.
Following conventions in the literature (which often capture individual differences via a handful of subject- specific parameters, such as learning rate), we assume the model's parameters 0 may vary over subjects, but are fixed over trials and sessions. Thus, for evaluation we optimize the parameters per-subject and score the model by cross pooss-alidating across sessions.' In particu- lar, for each subject i, we split its sessions into even and odd sets dfven := {st.0, S1,,8..,,,...MM.. [St.1, 81.3....,8..... 8.,M }. respectively.
and di :=
For subject i, iterating a model & with parameters O over session ((c,o::r,)) .......... re.r), will yield the se- quence (0o, (C2.o, 72,0; hr.o) = hi,,,, hi,1), ... (Pa.r, harr);; for each trial the accuracy of de, (.) is measured via the likelihood of the data under output distribution: ((c..elp..). We use cross-validation to evaluate the overall accuracy of
Rajagopalan et al. (2023) and Mohanta (2022) have per- formed extensive model comparison on similar datasets, and identified a popular model known as "Differential For- getting Q-Learning" (DFQ; Ito & Doya 2009; see Appendix D and Appendix E for details) as performing at least as well as any other. We adopt DFQ as the human disaa-discovered baseline model for the fruit fly dataset.
For the fruit fly dataset, since we have only one session per subject, forego this additional level of variation and treat the dataset as though it were multiple sessions from a single subject with a single 0.
4

---

Discovering Symbolic Cognitive Models from Behavior
within-subject patterns of behavior in an implicitly hierar- model follows. We fit two sets of parameters, O,eevnn and uddd on even and ddd respectively. All the even parameters O 0i,even are evaluated on the odd dataset as:
chical way. To the best of our knowledge, this is novel application of FunSearch, which has thus far been used as single-level evolutionary optimization process.
2 MT
12 (0,0, 0,"00) exp
Σ sEdtie t=0
Σ
log ((crelir)))
3.4. Configuring FunSearch
FunSearch is initialized with prompt that contains an ini- tial seed program do as well as documentation about the problem and scoring process. We explored four different initial prompts for CogFunSearch that varied the informa- tiveness of the seed program from mere "hints" as to what the programs should do in code and in natural language documentation-to more substantive starter code (see Ap- pendix F).
The score S2 ( is computed analogously, and we define the overall subject score (1) as the average of the two. The dataset score over the entire dataset D is computed by averaging over all subject scores in the dataset. See Sec. A and B for further specific details of scoring and how the datasets were divided for cross-validation.
LowInfo contains function and variable names (model, x, y) and documentation that describes fitting a generic statistical model. Structured1 includes informative function and vari- able names (agent, choice, reward) and documentation describing the data and modeling framework, but no explicit suggestions to the model's structure. Structured2 mod- ifies Structured/ to include hints on how to structure and style the program. Finally, FullModel initializes seed pro- gram do to the baseline program itself, permitting a "warm start." " These represent different levels of information with which a researcher might approach a modeling problem in different settings.
3.3. Evolving Cognitive Models With FunSearch
We use the above score function to evolve programs us- ing FunSearch RRomera-Pareder et al., 2024). Briefly, Fun- Search uses large-language models ( to create "muta- tions" of existing programs. To generate each new program d', the L.M is prompted with a pair of existing programs dn and di, along with information about how the programs will be evaluated and a directive generate an improved program ø'. The new program is then scored using the pro- cess described above (Section 3.2) and is added, along with its score, M(o'), to program database. To select "parent" programs for each LLM sample, FunSearch stochastically samples programs from this database based on their scores. We refer to the combination of FunSearch with our scoring function and framework for evolving cognitive models "CogFunSearch" (Figure 2).
4. Results
We performed three independent runs of CogFunSearch on each of our three datasets (Human, Rat, Fruit Fly) using each of the four seed programs described above (LowInfo, Structured1, Structured2, FullModel; Sec. 3.4). We ran each of these until it had sampled at least 1 150,000 new programs. This left with very large library of cvolved programs for each dataset, each of which can in principle be considered a hypothesis about the cognitive mechanisms at play in that dataset. Whether a program makes a good cognitive model is typically evaluated according to two criteria: how well they match data and how complex they are. Below, we consider each of these questions in turn.
Each CogFunSearch experiment is run with three separate seeds, each running for seven days, generating 100,000s to 1,000,000s of programs depending on the dataset. The produced models have up to ten unique parameters, and no bound on the hidden state size.
To validate the discovered models, we hold out additional test set containing subjects not seen by the FunSearch pro- cess. The programs that emerge from CogFunSearch are fit to data from these subjects, and again evaluated using an even-odd split. Unless otherwise indicated, all reported scores are performed on these held out test subjects, More details on how data were structured can be found in Ap- pendix B.
4.1. Best Discovered Programs Fit Data Well
We first considered the best programs discovered for each dataset and asked how well they perform fitting data. To identify single best program for each dataset, selected the program, from among all CogFunSearch runs across all seed programs, which earned the highest score on the training dataset. For the rat and fruit fly datasets the best performing programs came from runs using the FullModel seed program, and for the human dataset the best perform- ing program came from a run using the Structured2 seed
Our optimization thus effectively nests two levels: the outer- level CogFunScarch optimizes over Python programs (Fig- ure 2A), whereas the inner level fits program parameters over a dataset to evaluate performance (Figure 2C). The result is that CogFunSearch produces model "templates" that can be fit on a dataset, capturing across-subject and
ESpecifically, we use Gemini 1.5 Flash (Gemini Team, 2024).
5

---

Discovering Symbolic Cognitive Models from Behavior
Best Program Overr the Course of Evolution:
program outperformed its baseline model when evaluated held-out data (Figure 4). The performance of the final programs was broadly similar between runs using the three informative seed programs (Structured1, Structured2, and FullModel), but yielded lower-performing programs with LowInfo (p=0.05, 0.08, 0.05, 0.27, 0.05, 0.05, 0.65, 0,05, and 0313, for the nine comparisons between LowInfo and the other seed programs for the three species; Wilcoxon rank- sum tests across runs). Runs with different seed programs also differed in how many eycles of evolution were required to identify well-fit models (Fig. 4, rat dataset; Fig. 10, hu- man and fly datasets). These results demonstrate the ability of CogFunSearch to improve upon a human- -peeinnee model (FullModel), to synthesize a model that fits well given no information about the specific problem being solved (Low- Info), as well as to make use of intermediate amounts of information (Structured I and Structured2).
Average Norm. Lik. Score
All Rums on Rot Datasct
67%
FullModel
66%
Structured2 Structured1
Lowinfo Existing Madel (Baseline)
65%
0
50k
100k
150k
Steps of Program Evolution
Figure 4. Quality-of-fit over the course of program evolution. Results averaged over held-out test subjects. computed at each timepoint evolution for the then-best program of each CogFun- Search run. Individual runs are plotted as light lines, while dark lines represent the average across all runs.
program. The complete code for these "best programs" can be found in Appendix G.
To further test whether the best discovered programs pro- vided a quantitatively good match to data, we compared their fit to held-out data with that of recurrent neural networks. Specifically, we train GRU model (Cho et al., 2014) over dvvn run a sweep on the number of hidden units (over (1,2, 4, 8, 16, 32, 64, 128)), and use early-stopping to se- lect the best parameters. All the variants were trained with the Adam optimizer (Kingma & Ba, 2015) with a learning rate of le - 4. We found that the performance of the best discovered programs is broadly similar to that of the best RNNs, suggesting that they are capturing nearly all of the structure that exists in our datasets (Figure 5).
We find that the best discovered programs (maximizing over runs and seed programs) substantially outperform the human-discovered models for each dataset (Wilcoxon signed rank tests on score improvement by subject: W = 1772 (Human), 0 (Rat), 1479 (Fruit Fly); all with p < 0.002). Figure 1 shows relative improved likelihood for each held-out test subject as percentage points between the top-scoring CogFunSearch program (FS) and the species- specific baseline model (BL) (2fs- ). ). Across datasets, the majority of test subjects are better fit with the best Cog- FunSearch program than by the baseline (95.6% for Human, 100% for Rat, and 83.8% for Fruit Fly), with average im- provements in normalized likelihood score of 5.17 + 0.15 (Human), 0.54 + 0.06 (Rat), 0.48 + 0.05 (Fruit Fly).
A second measure of the quality of model is its ability to generate synthetic datasets which reproduce the scien- tifically important features of the real datasets (Wilson & Collins, 2019; Palminteri et al., 2017). To test our best pro- grams, we used them to generate synthetic datasets, match- ing the real datasets for number of trials per session, number of sessions, and reward contingencies. To characterize pat- terns present in each dataset, applied a trial-history logis- tic regression analysis common in behavioral neuroscience (Lee et al., 2005; Lau & Glimcher, 2005), which quantifies the extent to which rewards are followed (at various trial lags) by repeated choices treward-seekingg and the extent to which choices tend to be repeated ("choice persevera- tion"). The Rat and Fruit Fly dataset exhibit different pat- terns of reward-seeeking and choice perseveration, and we find that the patterss revealed by this analysis match closely between the synthetic and the real datasets (Fig. 6). This indicates that our discovered programs are able to match in detail the relationship between past choices and rewards, and future choice in our datasets.
Decomposing these results by seed program type and mul- tiple runs (and aggregating over subjects) investigate consistency, we found that CogFunSearch yielded well-fit models reliably: in 35 of our 36 runs the highest-scoring
Human
Fly
(Deemeeeed Modk 20-
Better
Difference Normalized
Likelhond (% Points
co
TWN indiiiaaia AMO Aeange Copfandeerct Keerage
Extsiig Model Better
-20
Figure 5. CogFunScarch closes the gap with neural networks. We compare the difference in normalized likelihood between the RNN and best discovered CogFunSearch programs. Across all species, the RNN narrowly outperforms best CogFunSearch pro- gram, although CogFunSearch has bridged the majority of the gap. Scores can be found in Table 1.
A final, very practical feature is that a model should be easy to work with in the sense of having a likelihood landscape that facilitates parameter optimization via gradient descent.
6

---

Discovering Symbolic Cognitive Models from Bchavior
Rat Dataset vs. Discovered Program Reward-Seekee
Quantitative Score v. Complexity
Choice Perseveration
Rat Dataset, FullModel seed program
Trial Hist-Histrry Regression
1.0
Example Rat Best Rat Program
LLM L.M-based pair wise Complexity Score
1500
Halstead Difficulty
Complexity Score
Esample Progrim
Highest program
Example Progrm 2 Example Progrum
Weights
1000
40
SHG Baseline
0.5
500
20
64.5
65.0
65.5
64.5
65.0
65.5
0.0
Average Normalized Likelihood Score (%)
10
Trials the Past
Trials in the Post
Fly Dataset vs. Discovered Program
0.15
Reward-Seeking
Choice Perseveration
Figure 7. Quality-of-Fit trades off with Complexity. Two com- plexity measures (LLM-based, Halstead Difficulty) both show a tradeoff between program complexity and normalized likelihood on training data. See Appendix for example programs from along the frontier.
Tril--Hictory Ragression
Fruit Flies
Best Fruit Fly Program
0.10
Weights
0.05
Examining these complexity measures (Figure 7 for Rat Dataset, Structured2 configuration; Figure 12,13 for all other datasets and configurations), we see a clear tradeoff be- tween likelihood score and complexity, with higher scoring programs tending have higher complexity scores. These plots also allow us identify example programs on the "effficiett frontier", earning the lowest complexity for their likelihood score. Researchers interested in identifying mod- els with different tradeoffs between fit and simplicity might choose to examine programs at different positions along this frontier. Comparing the frontier to the baseline program for each dataset, we see that the CogFunSearch programs include some that are improvements both in quality-of- and in complexity (e.g. "example program 3" in Figure 7), and that this is true across most seed programs and datasets (Figures 12 and 13).
0.00
10
10 Trials in the Post
20
Trials in the Past
Figure 6. Synthetic datascts generated by discovered programs match patterns in real data. Red lines show lagged regression model fit to behavioral data. Blue lines ("Best Program") ") show fits ten synthetic datasets each generated using the best program, model parameters fit to the data, and task conditions matched to the real dataset. Multiple synthetic datasets are shown to illustrate the variability that results from stochasticity in the model.
Hluman-designed models often lack this feature, and improv- ing it is the target of effortful reparameterization (Wilson & Collins, 2019). We find that our CogFunSearch programs return stable parameter estimates when multiple times, in- dicating that they have favorable likelihood landscapes (Fig. 17). They also find similar parameters across different splits of the data, suggesting stable solutions given the variability in the data (Fig. 18).
5. Discovered programs exhibit novel strategies
Despite not having been explicitly optimized for inter- pretability, we find that the best discovered CogFunSearch programs exhibit intriguing and often interpretable strate- gies. Moreover, many of the programs involve parameters, learning updates, and decision pociiign-making rules that were not present in our baselines. Here we aim to understand what computational motifs characterize our highest-scoring pro- grams (see Appendix G for full programs).
4.2. Quality-of-Fit Trades Off With Complexity
To evaluate the complexity of our discovered programs, we computed two distinct measures for each. The first is one of the Halstead (1977) measures, used in software engineering to quantify code complexity. We focus on the "difficulty" measure, which relates to how difficult the code for a human to read and understand, for example during code review. The second is a novel measure based on prompting an LLM to rate the relative complexity of our programs (details in C.3). We have found this measure to align well with our subjective sense of how easy it is to understand these programs. We compute these measures both for the programs returned by CogrhhSearch as well as our human- discovered baseline models.
Top Scoring Human Program A salient feature of this program that the bulk of the internal state-was devoted to choice history rather than re- ward tracking. In addition to variables tracking the expected values of the four actions, the best program introduced a number of novel variables that each track different reward- independent statisties of previous choices:
of the agent's s
7

---

Discovering Symbolic Cognitive Models from Behavior
baseline. The additional state variables implement a learning rule similar those found in "Q-learning" models (e.g. Ito & Doya 2009), updating the value of the chosen action towards the reward, but which also updates it towards the value of the unchosen arm:
avvalues sgent_ataceq:d] old_choico sgtnr_statc[4]
triai_eince__oee wtt__ece__s___tttttoo stant__state|s]] exploration=eess aqent_stata[e] cunchoice stent_state[/::1]]
These variable names are largely indicative of their actual function. ("q values" is standard terminology in reinforce- ment learning models for tasks of this kind to represent the reward an agent expects to receive following an action e.g. Ito & Doya 2009). Together, these variables can cap- ture behavior motif's on different timescales, such as long runs of the same choice, and other more complicated choice patterns such as cycles or switching within vs. between hands. Such patterns (several of which were also docu- mented by Eckstein et al. 2024) are evidently prominent in four-alternative bandit datasets, and not wholly captured by the simple stay-switch perseverative mechanisms included in the baseline model, which were inherited from models of two-alternative bandit tasks (Lau & Glimcher, 2005). This previously unappreciated aspect of the task may account for the relatively weak performance of baseline models in the human dataset relative to the other species.
state_w [choice]
alpie. state_x (chniice a alphaay) + swaard_frr_.pdee alpas biss garme x state wil-ohosce
(gamma.w and alpha are constrained to be posi- tive). This mechanism somewhat echoes the decay-toward- average motif discovered above. Other top programs im- plemented term tracking recent average reward, which modulated the exploration rate, reminiscent of several exist- ing cognitive models (Aston-Jones & Cohen, 2005; Eldar et al., 2016; Palminteri et al., 2015).
Top Scoring Fruit Fly Program The best Fruit Fly pro- grams were also derived from FullModel. Similar compu- tational motifs from the original baseline were present, but often in altered form. "Forgetting". which decays the value of the unchosen action, was modified to apply only when positive reward prediction error was present. Differential learning rates for rewards versus omissions were replaced by the differential forgetting of values. This program again exhibited mixture of exploration strategies that combined biased softmax and e-greedy.
Three independent runs with the Structured2 seed program, (highest scoring program for Human), separately discovered common (but, to our knowledge, novel) motif whereby the learned values were decayed, at each step, toward their average:
Best Runan Bandie Progran
laes exploration rate) valoes explosesion_rate jap.meen (gleel
Subjective Code Quality Broadly, we noted that param- eters were often assigned names that were meaningful in the context of RL (c.g. lapse.rate, initial.valuue These names were often indicative of their role in the pro- gram, and a sense of the program's function could be gleaned by browsing assigned parameters. Moreover, as indicated above, many lines of code were sufficiently com- prehensible that they could, with moderate effort, be mean- ingfully summarized.
Prograa 2
Vvviiss (1-biac) g_yclues bias top.ncan
Program 3
vaalue 11 choicce + vovluce sipha_choice choice ing.meen (gyuuss))
All three of these top programs also introduced ceiling on the logit-derived choice probabilities, which was not in- cluded in the bascline model but has been used in earlier published models of similar tasks (Shteingart et al., 2013). Suggestively, variable names from all three programs re- ferred this parameter as a "lapse" rate, presumably re- flecting the prompted LLM's contextual sensitivity and pre- training the psychology literature, where that term is commonly used (Green et al., 1966). In the reinforcement learning literature, in contrast, that rule is usually known as "c-greedy".
However, we note clear room for improvement in terms of readability and complexity. While variable names tended to be informative, for all seed programs besides LowInfo, number were vague (bias2), unnamed (params [0]), or misleadingly named (e.g. exploration rate in the best Human program snippet above actually modulated value for- getting). Programs frequently attempted to index more pa- rameters than were available (which does not throw an error in compiled JAX code, but instead defaults to indexing the final array element), and occasionally defined state variables they never used (e.g. the best Rat program defines 18 dimensional hidden state, and proceeds to update 7 of them). Programs had occasional tautological lines (agent state = agent state), or terms that seemed meaningful but were in fact multiplied by a parameter that always fit to a value of zero (gamma q in best Rat program). For examples
Top Scoring Rat Program The best rat program was derived from the FullModel seed program. This was the only dataset for which we found the state variables from the baseline model still present in the final program, and supports the notion that RHG (which was derived by a manual data-driven approach, Miller et al. 2021) is a strong
8

---

Discovering Symbolic Cognitive Models from Behavior
of programs that strike different balance of interpretability and performance (as in Figure 7, see Appendix I).
eters hidden states? Further LLM-based analysis (like the first steps we take with interpretability ranking) may help to address these questions. Also, to the extent different expressions of a model are behaviorally equivalent in some regime, they might suggest additional experimental tests or measurably different neural implementations.
6. Conclusion
We introduce CogFunSearch, an extension of FunSearch, which enables data-driven discovery of cognitive models of human and animal behavior. We find that CogFunSearch can discover programs that outperform state-of-the afrt base- lines for predicting animal behavior, while remaining largely interpretable. By sampling different programs from within the database generated by CogFunSearch, we can identify programs with different interpretabilityy performance trade- offs.
This work demonstrates the potential of LLM-guided pro- gram synthesis for discovering novel models of human and animal behavior. There are exciting opportunities to ap- ply this to other, potentially more complex behaviors and cognitive processes. More broadly, these results provide a promising indication that LLM-guided discovery can be a generally useful scientific tool in finding predictive models that are interpretable as well as predictive.
Future work will aim to close the remaining quantitative gap between the discovered programs and RNN, and to improve model interpretability. Hybrid neurosymbolic architectures with evolutionary architecture search and model distillation comprise two promising approaches. While we focus on bal- ancing quality-of offlit with interpretability, future work may also look at trading off quality-of- with other desirable model qualities, like runtime, optimizability, data efficiency, or generizability.
Impact Statement
This work introduces an Al tool for discovering interpretable predictive models of human and animal cognition. Mod- els that explain behavior and can offer hypotheses about underlying mechanisms are useful for psychologists and neuroscientists, and tools for improving this practice stand to advance the field considerably. Long-term, tools for build- ing better predictive models of behavior and cognition also stand to have a positive impact in the clinical psychology setting. as they might support diagnosis or intervention.
An important area for future work is to improve the ef- ficiency of (Cog))unSearch, which required hundreds of thousands of LLM calls to identify the best programs that we reported here. A simple idea we explored in this direction is to do form of rejection sampling during program genera- tion, which focused the optimization process on promising programs. Although this approach risks converging to local optima (due to a lack of exploration), our results indicate that we were still able to obtain interpretable programs that out- performed the baseline with only tens of thousand samples. More sophisticated approaches can potentially improve fur- ther on this afficiency-accurace tradeoff. We provide more details in Appendix C.6. It would also be worthwhile to ex- plore the efficacy of other LLM models, other AI discovery methods in the literature, well as non-differentiable ap- proaches (Acerbi & Ma, 2017; Ma et al., 2024; Zhou et al., 2024; Misra& & Kim, 2024; Ye et al., 2023; Lu et al., 2024; Rmus et al., 2025).
Modeling human behavior is also associated with a degree of risk, as models that effectively describe and predict be- havior could also be deployed to disrupt human behavior in targeted, adversarial ways. Dezfouli et al. (2020) illustrate how predictive models of learning and decision-making can be used to bias decision decision-making, and similar risk applies to our setting as well.
An additional consideration is the compute cost associated with using LLMs. In discovering these programs, Cog- FunSearch placed hundreds of thousands of LLM calls per run. We explore potentially more efficient approach in Appendix C.6; however, this remains limitation of LLM- based methods.
Acknowledgements We would like to thank Doina Pre- cup, Joelle Barral, Zoe Ashwood, the FunSearch commu- nity at Google DeepMind (Romera-Parrede et al., 2024; Novikov et al., 2025), and the rest of Google DeepMind for providing support and feedback during the preparation of this work.
Relatedly, the sheer density of programs output by CogFun- Search also highlights that choice predictiveness alone is likely not sufficient to identify single "best" model. This contrasts against the implicit assumption of much previ- ous work in this area, which has often focused on ranking smaller, curated sets of hand-designed models. An impor- tant question for future work will be systematically explor- ing similarities and differences between groups of models with similar predictive power: to what extent are they dif- ferent expressions of the same input-output relationship, or notational variants that reparamaterize equivalent param-
We would also like to thank the anonymous reviewers for their useful suggestions for improvement. Finally, the authors would also like to thank the Python community (Van Rossum & Drake Jr, 1995; Oliphant, 2007) for devel- oping tools that enabled this work, including NumPy (Haris et al., 2020), Matplotlib (Hunter, 2007), Jupyter (Kluyver

---

Discovering Symbolic Cognitive Models from Behavior
al., 2016), Pandas (McKinney, 2013) and JAX (Bradbury al., 2018b).
Co-Reyes, D., Miao, Y., Peng, D., Real, E., Levine, S., Le, Q. V., Lee, H., and Faust, A. Evolving reinforcement learning algorithms. arXiv preprint 1xXi::201100395. 2021.
The work at Janelia Reseach Campus was supported by fund- ing from the Howard Hughes Medical Institute. We thank Kaitlyn N. Boone (Janelia Project Technical Resources) for assistance with fly behavior assays.
Corrado, G. and Doya, K. Understanding neural coding through the model-based analysis of decision making. Journal of Neuroscience, 77311)8177-811 2007.
References
Correa, C. G., Griffiths, T. L.,, and Daw, N. D. Program- based strategy induction for reinforcement learning, 2024. URL https: ht/paaxxivvvi.... org abs /0lps/240...m 10.40...66666
Acerbi, L. and Ma, W. J. Practical bayesian optimization for model fitting with bayesian adaptive direct search. In Guyon, Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Ad- vances in Neural Information Processing Systems, vol- ume 30. Curran Associates, Inc., 2017.
Daw, N. D. al. Trial-by-trial data analysis using compu- tational models. Decision making. affect, and learning: Attention and performance XXIII, 23(1), 2011.
Dezfouli, A., Ashtiani. H., Ghattas, O., Nock, R., Dayan, P., and Ong, C. S. Disentangled behavioural representations. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buse Buch F., Fox. E., and Garnett, R. (eds.), Advances in Neural Information Process- ing Systems, volume 32. Curran Associates, Inc., 2019a. URL htttps://proceeddigga neurips. cc/paper Sicapee_fiees/pppeee/////////// 4077e1554eaccc333777556622222442 pdf.
Aston-Jones, G. and Cohen, J. D. An integrative theory of locus coeruleus nhrrpneernnn function: adaptive gain and optimal performance. Annu. Rev. Neurosci., 28(1): 403-450, 2005.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J.. Wanderman- Milne, S., and Zhang, Q. JAX: compos- able transformations of Python+NumPP programs, 2018a. jax.
URL htth::ttt///ttio
com
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., et al. Jax: composable transfor- mations of python+ numpy programs. 2018b.
Dezfouli, A., Griffiths, K., Ramos, F,, Dayan, P., and Balleine, B. W. Models that learn how humans learn: The case of decision-making and its disorders. PLoS computational biology, 15(6)ee1006003, 2019b.
Brunton, B. W. and Beyeler, M. Data-driven models in human neuroscience and neuroengineering. Current Opinion Neurobiology, 58:21-29, 2019. ISSN 0959- 4388. doi: https://bo..rr/11.0006...ooooob URL
Dezfouli, A., Nock, R., and Dayan, P. Adversar- ial vulnerabilities of human decision-making. Pro- ceedings of the National Academy of Sciences, 117(46): 2020. doi: 10.1073/pnas. 2016921117, URL https: /www, pnas.org/doi. Abs/10.1073//
com/
201 6921117.
Computational Neuroscience.
Eckstein, M. K., Summerfield, C., Daw, N., and Miller, K. J. Hybrid neural-cognitive models reveal how memory shapes human reward learning. 2024.
Brunton, S. L., Proctor, J. L.,, and Kutz, J. N. Dis- covering governing equations from data by sparse identification of nonlinear dynamical systems. Pro- ceedings of the National Academy of Sciences, 113(15):3932 -937, 2016. doi: 10.1073/pnas. 1517384113. URL htt.t://////// //wp::://wwwww. http://wwww........ pnas. org.ddorg/////// 0os/10.1071 11...//paa..1113811111..
Eldar, E., Rutledge, R. B., Dolan, J., and Niv, Y. Mood representation of momentum. Trends Cogn. Sci., 20 (1):15-24,, January 2016.
Ellis, K., Wong, L., Nye, M., Sable-Meyer, M., Cary, L., Anaya Pozo, L., Hewitt, L., Solar- A., and Tenen- baum, J. B. Dreamcoder: growing generalizable, inter- pretable knowledge with wake-sleep bayesian program learning. Philosophical Transactions of the Royal Society A, 381(2251):0020055, 2023.
Chen, X., Liang, C., Huang, D., Real, E., Wang, K., Pham, H., Dong. X., Luong, T., Hsich, C.-J., Y., et al. Sym- bolic discovery of optimization algorithms. Advances in neural information processing systems, 36, 2024.
Cho, K., van Merrienboer, B., Bahdanau, D., and Ben- gio, On the properties of neural machine transla- tion: Encoder-decoder approaches, 2014. URL https: 1259.
Gemini Team. Gemini 1.5: Unlocking multimodal under- standing across millions of tokens of context, 2024. URL https://aree /ttp:////iiiv// http:////xxxvv..r/bbbbbsss www./org////..... 240.b.///44... 05530.
//a www...oog/vaaaaa.
10

---

Discovering Symbolic Cognitive Models from Behavior
Ger, Y., Shahar, M., Shahar, N., and Ger, Y. Using recurrent neural network to estimate irreducible stochasticity in human choice-behavior. 2023.
Lake, B. M., Salakhutdinov, R., and Tenenbaum, B, Human-level concept learning through probabilistic pro- gram induction. Science, 350(6266): 1332-133333333 2015.
Green, D. M., Swets, J. A., et al. Signal detection theory and psychophysics, volume 1. Wiley New York, 1966.
Landajuela, M., Lee, C. S., Yang. J., Glatt, R., Santiago, C. P. Aravena, I., Mundhenk, T., Mulcahy, G., and Pe- tersen, B. K. unified framework for deep symbolic regression. Advances in Neural Information Processing Systems, 35:33985-33998 2022.
Haberkem, H., Basnak, M. A., Ahanonu, B., Schauder, D., Cohen, J. D., Bolstad, M., Bruns, C., and Jayaraman, V. Visually guided behavior and optogenetically induced learning in wead-fixed flies exploring virtual landscape. May 2019.
Lau, B. and Glimcher, W. Dynamic response-by geppons--y-rrepponse models of matching behavior rhesus monkeys. J. Exp. Anal. Behav., 84(3):555-577 November 2005.
Curr. Biol., 29(10):1677-100
Halstead, M. H. Elements of Software Science (Operating and programming systems series). Elsevier Science Inc., 1977.
Le, M., Yildirim, M., Wang, Y., Sugihara, H., Jazayeri, M., and Sur, M. Mixtures of strategies underlie rodent behavior during reversal learning. PLOS Computational Biology, 190)xx1011430, 2023.
Harris, C. R., Millman, J. Van Der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, F., Taylor, J., Berg, S., Smith, N. J., et al. Array programming with numpy. Nature, 585(7825);357-362 2020.
Lee, D., McGreevy, B. P., and Barraclough, D. J. Learning and decision making in monkeys during rock ces--ppper-sscse scissors game. Cognitive Brain Research, 26(2):416-430, October 2005.
Hunter, J. D. Matplotlib: A 2d graphies environment. Com- puting in science & engineering, ((03):90-55, 2007.
Li, M. Y.. Fox, E. B., and Goodman, N. D. Automated statistical model discovery with language models, 2024. URL wwthtttt//// //atps://arivv..o wwwpeiiiiorggaeoooooooe orb/ wttaaaa///222..
Ito, M. and Doya, K. Validation of decision-making mod- cls and analysis of decision variables in the rat basal ganglia. Journal of Neuroscience, 19881)4981199844 August 2009.
Lu, C., Lu, C., Lange, R. T., Foerster, J., Clune, J., and Ha, D. The Al Scientist: Towards fully automated open-ended scientific discovery. arXiv preprint arXiv:2408.06292. 2024.
Ji-An, L., Benna, M. K., and Mattar, M. G. Automatic discovery of cognitive strategies with tiny recurrent neural networks. bioRxiv, URL
2023. doi: 10.11012203.04412.556222
Ma, P., Wang, T.-H., Guo, M.. Sun, Z., Tenenbaum, J. B., Rus, D., Gan, C., and Matusik, W. LLM and sim- ulation as bilevel optimizers: A new paradigm to ad- vance physical scientific discovery. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp, F. (eds.). Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 33940-33962. PMLR, 21-27 Jul 2024. URL https: proceedings mlr www.llrp..ees/e/eeee /v2csss///255. ma2 4m. html.
https: //wp:://wwww./ www.:/iiorivvvvvooo wwwp:bboor..i..oo//coteetett early/2023/04/13///
536629.
Kingma, D. and Ba, J. Adam: A method for stochas- tic optimization. In Bengio, Y. and LeCun, Y. (eds.), 3rd International Conference on Learning Represen- tations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Truck Proceedings, 2015. URL http: // //arxiv. org/ abs / 1112. 6980.
Kluyver, T., Ragan-Kelley, B., Pérez, F., Granger, B., Bus- sonnier, M., Frederic, J., Kelley, K., Hamrick, J., Grout, J., Corlay, S., Ivanov, P., Avila, D., Abdalla, S., Will- ing, C., and Jupyter Development Team. Jupyter Note- books-a publishing format for reproducible computa- tional workflows. In IOS Press, pp. 87-90. 2016. doi: 10.3233/978-1-665 1-1119788----4499-999--9--9
McKinney, W. Python for Data Analysis: Data Wrangling with Pandas, NumPy, and IPython. O' O'Reilly Media, edition, February 2013. ISBN 9789351 100065, www.:/wwwwaaaaaaaaaa/a URL http: wwww://wwwwwc.. amazon com http///eee////ioo..///o/o/ soeelreetttaagrcttenneeee------- AS ASIN/1449319792
Miller, K. J., Botvinick, M. M., and Brody, C. D. From predictive models to cognitive models: Separa- ble behavioral processes underlying reward learning in the rat. BioRXiv, 2021. doi: https://do.orr/100 1101/461129, URL ht http///w./...c www.biorxiv.crg content/10. 1100
LaFollette, K. J., Yuval, J., Schurr, R., Melnikoff, D., and Goldenberg, A. Data driven equation discovery reveals non-linear reinforcement learning in humans, Jul 2023. URL osf. htt.://pppppinssssssaaeeeeo 65 15jeh_vi

---

Discovering Symbolic Cognitive Models from Behavior
Miller, K. J., Eckstein, M., Botvinick, M. M., and Kurth-Nelson, Z. Cognitive model discovery BioRXiv, 2023. URL
S., and Dehaene, S. theory of memory for binary se- quences: Evidence for mental compression algorithm in humans. PLoS computational biology, :7(1):e 1008598, 2021.
via disentangled rnns. https:// /www www. www.xiv.org/ /orgg//ooeeeet00 1101 /2013/203.. 06. 23. 546250v1.
Rajagopalan, A. E., Darshan, R., Hibbard. L., Fitzgerald, J. E., and Turner, G. C. Reward expectations direct leam- ing and drive operant matching in drosophila. Proc. Natl. Acad. Sei. U. S. A., 120639):e2221115122 September 2023.
Misra, K. and Kim, N. Generating novel experimental hypotheses from language models: case study on cross- dative generalization, 2024. URL https: arriv..
org/ /org/abs/20 2408. 05086.
Mohanta, R. Deciphering value learning rules in fruit flies using a model-driven approach. Master's thesis, Indian Institute of Science Education and Research Pune, Maha- rashtra, India 41 1008, 2022.
Real, E., Liang, So, D., and Le, Q. Automl-zero: Evolv- ing machine learning algorithms from scratch. In Interna- tional conference machine learning, pp. 8007-8019. PMLR, 2020.
Musslick, S. Recovering quantitative models of human information processing with differentiable architecture search. arXiv preprint arXiv:2103. (tti::2003.193330 13939, 2021.
Rmus, M., Jagadish, A. K., Mathony, M., Ludwig, T., and Schulz, E. Generating computational cognitive models using large language models, 2025. URL https:// arxiv. arxiv.org/ httgabs/////eeocceeegggggeeaaer/ / 00879.
Novikov, A., Vü, N., Eisenberger, M., Dupont, E., Huang, P.-S., Wagner, Z., Shirobokov, S., Kozlovskii, B., Ruiz, F. J. R., Mehrabian, A., Kumar, M. P., See, A., Chaud- huri, S., Holland, G., Davies, A., Nowozin, S., Kohli, P., and Balog, M. Alphaevolve: A coding agent for scien- tific and algorithmic discovery. Technical report, Google DeepMind, 2025.
Romera-Pareeer B., Barekatain, M., Novikov, A., Balog, M., Kumar, M. P., Dupont, E., Ruiz, F. J. R., Ellen- berg, J. S., Wang, P., Fawzi, O., Kohli, P., and Fawzi, A. Mathematical discoveries from program search with large language models. Nature, 625(7995) 2024. doi: 10. 0001:///1115600030002222222 URL https: / //dbi. /10. 1038 441586 023 06924 4
O'Doherty, J. P., Hampton, A., and Kim, H. Model-base fmri and its application reward learning and decision making. Annals the New York Academy of sciences, 1104(1):35 (1431):5--3,, 2007.
Roy, N. A., Bak, J. H., Akrami, A., Brody, C. D., and Pillow, J. W. Extracting the dynamics of behavior in sensory decision-making experiments. Neuron, 109(4):59--610.e6 2021. ISSN 0896- 6273. doi: https://oi..or///.1116.....mm 004. URL WWRL:ttsss/www.cceendde.e.eeeeee sttenee/artccce///////9622222222
Oliphant, T. E. Python for scientific computing. Computing in Science Engineering, 9(3):10-20, 2007. doi: 10. 1109/CSE.2007..55
Palminteri, S., Khamassi, M., Joffily, M., and Coricelli, G. Contextual modulation of value signals in reward and punishment learning. Nature Communications, 6(1): 8096, 2015. doi: 10. 1010///mmmmm9996.. URL https: //doi, org/ / 10. 1038/ 1occmmmm9999
Rule, J. S., Piantadosi, S. T., Cropper, A., Ellis, K., Nye, M., and Tenenbaum, J. B. Symbolic metaprogram search improves learning efficiency and explains rule learning in humans. Nature Communications, 15(1 (5(1):6847, 2024.
Palminteri, S., Wyart, V., and Koechlin, The impor- tance of falsification in computational cognitive modeling. Trends Cogn. Sci., 21(6):425-43) June 2017.
Sablé-Meyer, M., Ellis, K., Tenenbaum, J., and Dehaene, S. A language thought for the mental representation geometric shapes. Cognitive Psychology, 139:101527, 2022.
Peterson, J. C., Bourgin, D. D., Agrawal, M., Re- ichman, D., and Griffiths, T. L. Using large- scale experiments and machine learning to discover theories of human decision-making. Science, 372 (6547):1209-1214, 2021. doi: 10.1 abe2629. URL htttps:/www.scieee doi/abs// 110.111//sccincee.ee abe2629.
Shojaee, P., Meidani, K., Gupta, S., Farimani, A. B., and Reddy, C. K. Llm-sr: Scientific equation discovery via programming with large language models, 2025. URL https://aaxive wrg. absg///b//// abaaaba/2440 /244 8400.
org/
Shteingart, H., Neiman, T., and Loewenstein, Y. The role first impression in operant learning. Journal of Experi- mental Psychology: General, 142(2):476, 2013.
Planton, S., van Kerkoerle, T., Abbih, L., Maheu, M., Meyniel, F., Sigman, M., Wang, I.., Figueira, S., Romano,
12

---

Discovering Symbolic Cognitive Models from Behavior
Song, M., Niv, Y., and Cai, M. Using recurrent neural networks understand human reward learning. In Pro- ceedings of the Annual Meeting of the Cognitive Science Society, volume 43, 2021.
Sugrue, L. P., Corrado, G. S., and Newsome, W. T. Matching behavior and the representation of value in the parietal cortex. Science, 304(56788
June 2004.
Van Rossum, G. and Drake Jr, Python reference man- ual. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
Venditto, S. J. C., Miller, K. J., Brody, C. D., and Daw, N. D. Dynamic reinforcement learning reveals time-dependent shifts in strategy during reward learning. eLife, 2024.
Wilson, R. C. and Collins, A. G. Ten simple rules for the computational modeling behavioral data. eLife, 8:e49547, nov 2019. ISSN 2050-084X. doi: 10.7554/ 49ife.49547. URL https: http://i.ooogg////////// 7554/ eLife. 49547.
Ye, S., Lauer, J., Zhou, M., Mathis, A., and Mathis, M. W. Amadeusgpt: natural language interface for interactive animal behavioral analysis. In Proceedings of the 37th In- ternational Conference on Neural Information Processing Systems, NIPS 23, Red Hook, NY, USA, 2023. Curran Associates Inc.
Zhou, Liu, H., Srivastava, T., Mei, H., and Tan, C. Hypothesis generation with large language models. In Proceedings the /st Workshop on NLP for Science (NLP4Science), pp. 117-139, Association for Compu- tational Linguistics, 2024. doi: 10. 10.18653/v120224 11p4science-1.10. URL 18653/v1/2024. (6553//112222.npppacccceee----.. m1p4science111.1..
13

---

Discovering Symbolic Cognitive Models from Behavior
A. Extended Methods
A.1. Evaluation
Our models are discovered using a bilevel optimization procedure in which programs are optimized across- using an LLM-powered evolutionary algorithm in the outer loop, and parameters of the model are per-subject using gradient descent in the inner loop (Fig. 8). This necessitates two levels of training and validation test sets: one (for the outer loop) that splits the data into train and test subjects Dunin and Dees and one that splits the data of subject into sessions diankl and even used for crossvalidation. For clarity, we refer to validation at the level of the outer loop as "evaluation" and validation at the level of the inner loop as "scoring."
In Sec. 3.2, we describe the procedure by which each program is scored on data from each subject using two-fold cross validation across sessions. For each subject i, we divide sessions into equal splits (even and odd sessions), fit the parameters of the model e separately to each split, and evaluate the fit parameters evenn and 0 aadd on data from the opposite split. This gives us two scores per subject: even (Ø, (o.dd) and (o, 0 .evee) - the likelihood of data from even sessions given parameters fit odd data, and vice versa. We average these scores together to produce single score (2(()) per subject. These per-subject scores are then aggregated across the data in the training by computing the average across all subjects in the dataset (())..
During scoring, parameters for the program are fit to the training data with gradient descent. CogFunSearch'' programs must be implemented Jax (Bradbury et al., 2018a) so that they are differentiable. This is communicated to the LLM by including type hints and Jax imports in CogFunScarch'ss 'ss specification. Non-jittable functions will usually throw an error and not get added to the program database.
We use the AdaBelief optimizer with learning rate 5 x 10-2 which is run until convergence until 10,000 steps of gradient descent are reached. In order to test convergence, we compare the current score at iteration k., S2 to the previously recorded is less than a convergence threshold
score - every 100 steps. If the relative change in score [O2 10-2 2. we conclude that parameter fitting has converged. In order to protect against local minima, this parameter fitting process is repeated from different initial parameters up to 10 times, or until programs have converged to the current best.
We maintain a group of held-out subjects, Dieat in order to validate our discovered programs. Data from these subjects was not seen at all during the evolution scoring of FunSearch programs. In order to score programs on this held out data, parameters are fit to data from each subject using the cross validation procedure described above.
Bilevel optimization Finding programs scrost- and fitting programs subjen
Evolve program on training dataset
D
train
Fit params
score
model <p
trans
optimize model @
D
Score on test dataset
test
score Opert
piedict
model sp*
Figure 8. Organizing data for train and test. The full dataset is split into Duain and Dues; each of these is further split into fit and eval subsets (i.e. D and Deval ). fit and eval will correspond even and odd sessions, respectively. and then parameter fitting will be repeated and the scores averaged. In the top panel, generated programs are evaluated by first fitting params on D and then evaluating them on Deva After the evolutionary process has completed, the resulting program is tested on Des in an analogous manner. This latter score is what is used for our comparisons throughout the paper.
14

---

Discovering Symbolic Cognitive Models from Bchavior
B. Extended Dataset descriptions
Generally, we divided all of our datasets into training and test sets to validate programs, which were then further divided into even and odd datasets validate parameters, However, the Human, Rat, and Fruit Fly datasets differed in key structural ways which affected the data partitioning.
Dataset splits
Human
Rat
Fly
Bessoon 5 862 subjects
session 90
347 subjects
Do0 20 subjects
-600 trials
5000 trials
150 trials
Test
Train
Figure 9. Organizing data for train and test. In Human and Rat datasets we use half of the subjects for training, and half for testing; for each train subject, we use half of its sessions for parameter fitting and half for evaluation. For Fly dataset we use half the subjects for training and half for testing, and proceed similarly as for the other datasets.
As detailed in Fig. 9, the Human dataset contained a large number of subjects and trials, but only five sessions per subject. Despite the small number of sessions, we found we could fit the ten parameters of CogFunSearch programs (using cross-validation given the data splits denoted by the horizontal dashed line in Fig. 9), and programs were optimized using scores aggregated across subject (validating on the test subjects denoted with the dotted lines). The four parameters of the PFQ baseline were fit similarly.
However, the small number of sessions meant that attempting to fit the parameters of an RNN to each subject resulted in extremely low performance. Thus, as described in the main text, we followed the approach of Eckstein et al. (2024) and fit an across-subject RNN model. joining sessions from each subject to form one long concatenated session. While parameters therefore not adapted per-subject for this model, the RNN can adapt in its latent space reflect subject-specific activity, The data organization that we used for fitting RNNs ignored the per-subject division into even/odd splits (denoted by dashed line in Fig. 9), and only required the train/test split indicated at the across across-subject level (like the division illustrated for Fly). While is not an apples-to-appee comparison, it does compare each model in the regime where they are best set up to succeed.
The Rat dataset contained an intermediate number of subjects, but with far more sessions and trials per session, permitting per-subject parameter fitting for all models. An important property of the experiment design is that 20% of trials were randomly selected be "instructed",, meaning that the animal was instructed make a particular choice rather than allowed to choose freely. Failure to make the instructed choice resulted in adverse outcomes (loud noise and timeout). This was included to enforce exploration, and ensure that the animal is receiving informative experience to drive the learning process that is the targetted object of study. This is similar to "teacher forcing", which is commonly used for training sequential models.
The fly dataset contained a large number of subjects, but contained only one session per subject. This precluded any possibility of fitting parameters per-subject; rather, we followed prior work (Mohanta, 2022; Rajagopalan et al., 2023) and fit parameters of metasubject model for CogFunSearch, DFQ baseline, and the GRU. The way split the data is different for this data also differed (Fig. 9): rather than having train and test subjects that were divided into sessions, we had train and test subject populations and only the training subjects were further subdivided into even / odd splits. Model evaluation is then performed by fitting the parameters of the model to all subjects in the training set, and then reporting performance on the held-out test subjects.
nuance about the Human dataset that what we refer to as "sessions", for notational consistency, , would more accurately
15

---

Discovering Symbolic Cognitive Models from Behavior
be referred to blocks, because they are experienced by the subjects consecutively. Between blocks, subjects received summary feedback about the previous session and took short break. The appearance of arms changed after each block and the drifts reset. Because of the temporal proximity of sessions in this dataset, there are likely dependencies between these sessions that we are not explicitly capturing.
C. Extended Results
C.1. Normalized likelihoods
Table 1 presents the average absolute likelihoods for the Baseline, RNN, and CogFunSearch methods, with standard errors.
Baseline
RNN
CogFunSearch
Human
55.77 + 0.64 66.63 + 0.53
62.04 = 0.62
60.93 + 0.63
Rat
67.41 = 1.21
67.22 + 1.25
Fruit Fly
53.99 + 0.24
54.71 = 0.57
54.62 + 0.57
Table Normalized likelihood scores (as percentages) and standard error, as computed over held-out test subjects. evaluated methods each of the datasets.
The number of samples from the LLM required by CogFunSearch to produce these models was 1.2M (Human), 300K (Rat), and 200K (Fly).
C.2. Differences across Seed Program
Here, present extended results the performance differences induced by the different seed programs considered throughout training. Figure 10 shows how performance of the best discovered program changes over the first 150,000 LLM samples for each CogFunSearch (fruit fly and human datasets, corresponding plot for the rat dataset can be found in the main text as Figure 4). Figure 11 shows the performance of the best program discovered at the end of evolution for each CogFunSearch run. Performance is often better than in Figures and 10 as many CogFunSearch runs were allowed to proceed for more steps of evolution.
Normalized Likelihood
Fly Dataset: Evaluation
Normalized Likelihood
Human Dataset: Held out Subjects
54.5%
Lowinfo
60.0%
Lowinfo
Structured1
Structured1
50.0%
Structured2
53.5%
FullMdee
FullModel
CogModel Baseline
40.0%
CogMadel Baseline
52.5%
30.0%
0k
30k
60k
100k
Ok
10k
20k
30k
Number of Samples
Number of Samples
Figure 10, Impact of seed program on quality and discoverability of models. Evaluation results of (left) fly dataset and on (right) human dataset, computed at each timepoint in evolution for the best program far discovered for each CogFunSearch run. Individual runs are plotted as light lines, while dark lines represent the average across all runs.
16

---

Discovering Symbolic Cognitive Models from Behavior
Human
Rat
Fly
Difference Norm alcedd Litell hood 1hss Points)
0.5
0.4
0.4-
0.2
0.2-
6.1
Structued1
Structured2
Figure Impact of seed program on final model performance, from different seeds and independent runs, for each datasets. Each dot represents the average normalized likelihood improvement (over haseline) averaged over the held-out test subjects. Each dot represents an individual run, illustrating the variance incurred from independent runs.
C.3. Complexity Statisties
We evaluate the interpretability of the generated programs as well as the baselines, relative to their scores. To do so, we use of the Halstead measures (Halstead, 1977), which is standard metric for software complexity. Figure 12 plots, for each dataset, the generated programs from different program seeds, well the respective baseline.
Program Seed
Lowinfo
Structured
Structured22
MallM
PHQ Sase ins
-0
Dataset
RHG Base inc
64
64
64
64
75
DFO Base inc
Co
54
Figure 12. Halstead Difficulty for generated programs given different seed programs. Each row represents different dataset (Human, Rat, Fly), each column different seed program, each point different program.
We additionally explored novel measure of complexity by prompting large language model (LLM) rate the relative complexity of our programs (see Appendix C.3.1 for the prompt used). Figure 13 plots the complexity score against their likelihood, for each of the generated programs from the best seed prompts for each dataset.
Here, to reduce LLM calls, we limit the Human programs to those with values above 56% normalized likelihood, which why PFQ Baseline appears to be an outlier in Fig. 13 but not in 12. Far fewer programs were generated for Fruit Fly and Rat datasets so filtering was not required.
17

---

Discovering Symbolic Cognitive Models from Behavior
Human
Rat
Fly
LUM DOGOO pairant vatrantse compleity
1000
400
750
1000
300
500
250
500
01
Highest scoring program Exampple Program
Helttedd biffectt
Example Program
Example Peogram Basoline
20
53.75
54.25
Normatize lksithood (>
Norallzzd likeihoodd (>
Normalized
Figure 13. LLM-based and Halstead Complexity Program Scores for Different Datasets. Each column represents different dataset (Human, Rat, Fly), and point different program.
18

---

Discovering Symbolic Cognitive Models from Behavior
C.3.1. LLM-BASED PAIR-WISE COMPLEXITY SCORES
We the prompt below to query LLM decide which of pair of programs less complex. These "preferences" then used to compute the scores displayed in Figure 7 and Figure 13.
Below code anippeta. Which of them is essier to understand?.
Code enigpet 1:
(program >>
Code snippet 2:
<program 2>
Please explain your reasoning. Ead your answez with "Pinal Answert <ans>", where ans is 0 if the first, understand
snippet simpler and/or second an.ppet. Otherwise, is 1. Do not output anything after that.
This comparison operation is applied to all programs in the dataset with MergeSort, which returns a ranked list from least to most complex. This rank what shown as LL.M-base complexity score in Fig. 7 and Fig. 13.
C.4. Quantifying variability discovered models
(Human)
(Human)
20
froes
Params odd
40
-10
40
-10
10
20
Params fIL to even
Params fit to even
Figure 14. Relationship between parameter values fit to even and odd splits of Human dataset. Each color denotes different parameter, and each dot human subject dataset. a-coordinate denotes values parameters fit to even sessions, the y-coordinate the values of parameters fit to odd. (left) Parameters for best Cug program, (right) parameters for PFQ Baseline.
Best FunSearch (Rat)
RHG Baseline (Rat)
Bewent
Params to 000
5.0
Bust
2.5
2.5
Reaadd
8.0
Dueel
0 Params fit even
Params even
Figure 15. Relationship between parameter values fit to even and odd splits of Rat dataset. Each color denotes a different parameter, and each dot in the dataset. The z-coordinate denotes the values parameters fit even sessions, the y-coordinate the values of parameters fit to odd. (left) Parameters for best CogFunScarch program, (right) parameters for RHG Baseline.
19

---

Discovering Symbolic Cognitive Models from Bchavior
Best FunSearch (Fly)
DFQ Baseline (Fly)
Params fit oda
20
B:
10
Kiow
Bb
Dia:
5-
Ra
D:
Parame
Parame to
Figure Relationship between parameter values even and odd splits of Fruit Fly dataset. Each color/dot denotes a different parameter. Multiple dots are not shown because single set of parameters was fit for all subjects. The z-coordinate denotes the values of parameters even sessions, the y-coordinate the values of parameters fit to odd. (left) Parameters for hest Cog FogFurrearch program, (right) parameters for DFQ Baseline.
In order to assess the robustness of the fit parameters under the discovered program, we compare the values of different experimental parameters found across different data splits (e.g. even v. odd sessions). In Fig. 14, 15, and 16, we show the parameters recovered for the even/odd split for the best discovered program for each dataset. We see that parameters discovered across different splits of the data covary closely. We find these parameters to be even more consistent than those found under the cognitive model.
C.5. Optimizability of the discovered models
A risk of our optimization pipeline is that we may evolve programs that are overadapted for the specific initial parameters used for optimization, and show extreme sensitivity to the random seed used for parameter fitting. In order to determine whether this was the case, we performed parameter fitting 50 times with different random seeds for both the best Rat program and the RHG cognitive model baseline, for comparison.
We explore this for the best Rat programs found with FullModel seed program, and find that the discovered programs do show preference for the seed that was used during optimization (normalized likelihood training seed = 67.206%, mean normalized likelihood across other seeds = 67. 187 + 0.029, t - p = 5 10 5, only 15/49 seeds have larger values than seed=0) (Fig. 17). Regardless, the best Rat program continues to outperform RHG baseline across all seeds, and also appears to show less dependence on seed than the baseline (normalized likelihood 66.64 + 0.079). Thus, while some specialization for the training seed is evident, we conclude that does not seem to result in more brittle, hard-to-optimize programs.
Fitting Model Parameters with different seeds
Normalzedd likelihood
67.2
X Seed 0
All seeds
67.0
All seeds (RHG)
66.8
66.6
66.4
RHG Baseline
BestRatBandit
Figure 17. Evaluating whether evolved programs overadapted to the seed used for optimization, when doing parameter fitting. We observe less dependence on seed than the baseline, and all random seeds continue to outperform all fits to RHG haseline.
Another risk of our pipeline is that we might learn programs that are overparameterized, with many different parameter values capable of yielding similar solutions. this case, we would expect to be substantial variability in the parameters discovered when using different seeds. However, looking at the best CogFunSearch program, we find that the value of parameters discovered from different seeds tend to covary with that of the seed used for optimization. Shown Fig. 18 are
20

---

Discovering Symbolic Cognitive Models from Behavior
the parameters discovered from different values.
RHG Baseline
Params seed!=best
4
0
-4
Params (seed=best)
RatBanditBest
Params seed! =0
0
-4
-4
4
-4
-4
-4 0 4
-4 -4 Params (seed=0)
4 0
-4
4 -4
-4
Figure 18. Evaluating whether our models are overparameterized. We plot the parameter values found by seed 0 against those found by other seeds, find that they are generally covariant, suggesting our models are not overparameterized.
C.6. Rejection sampling for compute efficiency
We experimented with form of rejection sampling for reducing the computational expense. Specifically, we maintain set of scores Ω = (Ro, ...... } } for all programs evaluated thus far (see Section 3.2 for the details on how 12 is computed). We use the 90th quantile as a threshold r. When new program Φ; is generated, we first evaluate it on a small subset the subjects (in our experiments we used 20) to produce proxy score SY, If 2 2 т, we compute Q2 on the full set of subjects and add to 52; if not, we reject di.
This simple technique has the effect of only evolving promising programs. This myopic strategy naturally trades off exploration for computational savings, yet as Figure 19 and Table demonstrate, the resulting programs significantly improve over the baseline (Wilcoxon signed rank tests on score improvement by subject: W = 2203 (human), (rat), 1322 (flies); all with p < 0.02), and are competitive to the full runs.
Human
Rat
Fly
30
15
Difference in Normalizer alized
Likelihood (% Points)
10
X
-10-
-20
Subsampled
Full
Figure Rejection sampling can produce accurate models fraction of the compute. Despite trained with order of magnitude fewer samples, the discovered models still significantly outperform the baseline (p < 0.02 for all datascts).
21

---

Discovering Symbolic Cognitive Models from Behavior
Baseline
CogFunSearch (Subsampled)
CogFunSearch (Full)
Human Rat
55.77 = 0.64
59.38 0.62
60.93 + 0.63
66.63 0.53
67.03 1.27
67.22 + 1.25
53.99 = 0.24 Fly
54.58 0.56
54.62 + 0.57
Table Normalized likelihood scores (as percentages) and standard error, as computed over held-out test subjects, for the evaluated methods each of the datasets.
D. Model Parameters
In this section we present the numerical values and distributions for the parameters in both the baseline and discovered models, for each dataset.
D.1. Human parameters
Figure 20 displays the distribution of parameter values of the best model discovered by CogFunSearch for the human dataset, after fitting the parameters to the data. Table 3 displays the average and standard errors for these parameters.
B Bleam
Ofteret
persey
6
6:
9g
8a
S.ubjeetts
Subjects
200
250
200
even
ood
-3.5
21
-3.4
10
10
06
0
0g
0g
Subject
200
200
100
200
Figure 20. Distribution of parameter values of the best model discovered by CogFunSearch for the human dataset.
PFQ Baseline
BestHumanBandit
Parameter
Value
Parameter
Value
Assigned name
Transformed Value
Breward
5.9 + 1.7
Ou
15 = 4.5
beta.r
14 + 3.7
Oleam
0.19 + 1.6
01
3.1 + 2.6
lapse
0.83 + 0.2
Offoreet
-0.6 3.7
02
-4.1 1.6
prior
0.044 + 0.091
persev
0.0018 + 0.12
03
2
alpha exploration rate
0.2 + 0.18
01
-.4
decay-rate
0.026 + 0.05
05
0.65
attention.biast
-0.5 + 0.65
Os
5.7
attention.bias?
-5.7 + 1.8
07
1.1
perseveration strength
0.75 + 0.48
0g
-6.6
switch.strength.
-6.6 + 1.7
0y
4.6
lambda lambda.parrm param
1 0.036
gamma
4.6 + 0.92
temperature
4.6 0.92
beta.p
4.6 0.92
Table 3. Parameter values for PFQ Bascline (left) and HandiBand discovered FunSearch program. For the discovered program, also show values after they are transformed inside FunSearch program, as well as names they are assigned hy the programs. Shown the means and standard deviations across all subjects and across both and odd session splits. Note that additional named variables transform d they typically reflect attempts to index further into the parameter array than the length (which in Jax. defaults selecting the last element).
22

---

Discovering Symbolic Cognitive Models from Behavior
D.2. Rat parameters
Figure 21 displays the distribution of parameter values of the best model discovered by CogFunSearch for the dataset, after fitting the parameters the data. Table 4 displays the average and standard errors for these parameters.
Subborts
Preecand
Shatt:
Poomater
dreaard
Brastt
Subjects
Preward
Buath
Ppambtrr
Grewerd
Ghabt
muen
even
ood
bias
Ggomber
bias
Subjects
-0.00
5.6
-0.51
0.21
5.88
5.6
-0.5
0.21
Figure Distribution of parameter values the best model discovered by CogFunSearch the rat dataset.
RHG Baseline
Best Rat Program
Parameter
Value
Parameter
Value
Assigned name
Transformed Value
Breward
1.2 +: 0.54
On
-1.1 + 0.46
alpha.r
0.35 + 0.074
Phabit
0.63 + 0.48
0i
-0.41 + 1.4
alpha h
0.55 + 0.21
3 gaamler
0.64 + 0.66
02
2 0.2
alpha.g
0.98 = 0.025
Orwward
0.084
fg
-0.031 = 0.48
alpha.q
0.6 0.1
O'abit
1.5 + 1.2
0.72 0.49
beta r
1.1 0.32
Oggamller
2.3 2 2
0.76 0.87
beta.h
1.2 0.46
bias
0.0059 + 0.16
1.1 0.74
beta_g
1.5 = 0.56
07
-1.7 0.33
beta_q
0.18 = 0.056
0s
0.27 + 0.25
bias
0.27 + 0.25
0g
-0.26 + 0.12
beta.q.bias
-0.26 + 0.12
beta h bias
-0.26 0.12
exploration.n iniiseeennieeeee
0.57 0.051
lapse rate
0.43 = 0.028
gamma q
0.57 + 0.051
gamma_w alpha bias reward.bias
0.64 = 0.02
0.43 = 0.028
-0.26 +: 0.12
bias 1
-0.26 + 0.12
Table Parameter values for RHG Baseline (left) and Best Rat Program discovered by CogFunSearch. For the discovered program, we also show the values after they are transformed inside the CogFunScarch program, as well as the names they are assigned by the programs. Shown are the means and standard deviations across subjects and across both even and odd session splits. Note that 8, additional named variables transform 6g, and typically reflect attempts to index further into the parameter array than the length (which in Jax, defaults selecting the last element).
D.3. Fly
Table 5 displays the average and standard errors for the parameters of the best model discovered by CogFunSearch for the fly dataset, after fitting the parameters the data.
23

---

Discovering Symbolic Cognitive Models from Behavior
DFQ Baseline
BestFlyBandit
Parameter
Value
Parameter Oo
Value
Assigned name
Transformed Value
OTeam
-3.8, -4.1 -1.9
-3.8, -4.1
params [0]/10
-0.38,-0.41
O'oorgtt
01
-1.9,-1.9
1r temperature
0.13,0.14
Kreward
1.7, 1.5
02
1.7,
5.3,4.5
bias
-3.1
0g
-3.1,
lapse.rate
0.043,0.039
04
1,2
alpha
2.8,3.2
0:
0.083, -0.22
bias [0]
0.74,0.59
0s
0.14,00.071
bias
0.76,0.66
07
-1.3,-1.3
forget.rate
0.24,0.25
0g
-.1,
decay factor
0.98,0.98
0g
0.97
exploitatioo..aae
0.61,0.71
Table Parameter values for DFQ Baseline (left) and FlBannitBBest discovered FunSearch program. For the discovered program, we also show the values after they transformed inside the FunSearch program, well as the names they assigned the programs. The two values shown correspond to the parameter values for (even, odd) subjects. Since only one set of parameters is recovered per dataset split, error provided.
24

---

Discovering Symbolic Cognitive Models from Behavior
E. Baseline Programs
E.1. Human Cognitive Model Baseline
We used the following implementation the PFQ Baseline from Eckstein et al. (2024).
agent(
params: chex. Area,
choics: int,
reward: int, ageat_atate: Agcionsl[chexAcrra/11.
Tuple [eleex.errrar Array, chex. :
"""oggiccvv model desersbing human behavior on a multi-armed bandit task.
Assumes the agent is presented with four options on each trial.
Arge:
perame: list containing (beta_f, deera_ leern, sipha_ forget, and p]
cheice: choics made the agent the previous trial. 0, 1, 2, reward: Revord received by the agent on the pravious trial. scalar between
agent_atate: (oi, 02, 03, 04, previous_ [ If None, assumes this is the first triat of new sesaion and detaults to (0, 9, 6, 0, -1j.
Returns:
pecice_logitas The probabilities that the agent will choose option or 3 on the next trial, expreased as logits. agent_stater agent
agent_state is None:
ageat_state arrpy ((IO, 0, 0, 0, -11))
- [gent__tatte[4]]
pzev_choice jap. int.2 [aaent
beta_r params [0]
alpha_leara = (1 exp (-perams
alpha_farger 1 [1 jxp.exp (-params (2]1)
p params [3)
back perseveration pazem should $ Nele Indicate this using prev_choice of first triel session p (prev_choice !- -1)
# Update 0 for chosen action
choice "jp.. int (choice)
- at [choice) .ssoicceesstte set (slpha_learn) (zeward (sccoiice) qe [choice)]
Update 0 for unchosen actions using maak mask sap.ones_lae Like (gs, dtype-bool) nask.at [choice] .sell.setee (False)
= whare (maak, qa sipha_forgat, ga)
Velues for choice: @ plus boaus for qu_for_choize (pseetppeecchcccel..co
(ppee__chaieee
agent_state = jap aqp.a choice)
# Compute choice togits
cheice_iogils beta_r qs_for_choice
zeturn choise_logits, agent_state
25

---

Discovering Symbolic Cognitive Models from Behavior
E.2. Rat Cognitive Model Baseline
We used the following implementation of the RHG Baseline model from Miller al. (2021).
def a
params: Arex.
choice: int,
reward: int,
ageat_atate: (ocional[chex.Areee
(eple(chex.Arron Array, ohex, Array]:
"*Cognitvve sodei describing behavior binary two-armee bandit task.
Arge:
perams: Model params. niffarent perameters are ueed for different rats.
choice: Previous choice, Values: 0
reward:
Previous reward. Values: agent_atater Previous state of the agent
Returns:
choice_ logits: Vector of shape (2,)
probabilitics choose opt or 1 the next trial, expressed as logits.
rae
agent_ state: New state the agent, after observing the previous choice and reward.
agent_state
ageat_state jnp. zeroz ((3,))
beta_r params
beta_h pa:ams [1]
beta_g params [2]
alpha_: -
[1
jap.expl exp params (31
alpha_h
jap exp (-params [411)
alpho_g 1 [1
jap exp (-params (5]1)
blas parama (6]
# Convert choice and reward 0/2 to +-1 chaico_for_upánt 2-choice-1 seward_for_updata 2sreward-1
# Update 8
agent_ state (0) alyhe_r agent_state[0] (1 alpha_r) for + choice_For_update
# Update R agent_state- alpha_h agant_state(1] + alpha_h) choloe_for_update
Update C agent_ state
alpha_g agant_state(2)
f alpha_g) aitrice_for_upder
reward_for_updan choico_for_updates
choice_term
Update chaice choice_Leat i- beta_ [gent_atatee0]
$ Update choice term from choice_term 1- bete_h ageat_state(l]
f Update choice term from G. choice_term beta_g agent_state[2]
$ Update choice term from blas. choice_Lerm 1- blas
chaice_iogita anpay 1]) chaice_term choice_logils, agent_state
26

---

Discovering Symbolic Cognitive Models from Behavior
E.3. Fly Cognitive Model Baseline
We use the following implementation of the DFQ Baseline model from Ito & Doya (2009). See Ito & Doya (2009) and Mohanta (2022) for more details.
I
params: chex.
choice: int,
reward: int,
ageat_stater Optional [chex Array]
Tuplo [coelecce..Arar chex. :
""Cognitivv modei describing fly behavior binary two-armed bandit task.
Args:
params: Fit paraneters of the model.
chofce: The choice made by
this trial.
The reward received by fly on this trial. 0 or agent_atate: current state the cognitive model.
reward:
Returns:
logice the probabilaties that the fly will choose option or the next triai, expressed as logita.
agent_atater updated state of cognitive modei.
# Toieialize Q-vaiues co
if agent_state is None:
ageat_state jap.errey (10.5, 0.5))
# thpack and rename parameters for readability.
alppa__eern_logit- paranr [0]
alpho_ hipho_forgot_loott parama[1]
<appa_reward params[2]
sappa_omission: parama[3]
alpha_leara 1 alrga_forget
exp l-gnaalgh ttpp.oogr-iipph_ffgg___...o.
qs = agert_state
# Chosen aotion's O-value will be updated cowarda thís target va lue. learn_tazget kappa_reward zeward kappa_omission x reward)
f Update gevalue for chosen action.
qs.at (choica) ( alpha_learn (learn_target qe qs(chcice]] + 11 alpha_leazn) qa [choice)]
Update g-value for unchosen scrion: decay towards 0. qs q1 settcc).set(() Torge_) x Igs [1 choice])
logits q3 new_agent_ qs
return Logils, aewae
27

---

Discovering Symbolic Cognitive Models from Behavior
F. Seed programs
In this section we seed programs provided to CegFunSearch for the rat dataset. Note that the programs used for human and fly datasets are functionally identical, with the exception the FullModel seed program.
F.. LowInfo seed program
"**pits a modei to a dataaet.
def acore__ acore_funation (data_1d):
**aFica pazameters of model and returne the normalized likellhood.
Arge:
dota_1d: Data id,
Returne: Normai fxed 1ikei ihood thodd of the fit moda on held out data, in range 10, 11.
dataset load_dataset id)ts_id) hormalize__lkkelatee
twofold_.cooevvlldaettt (datasat, model) return norma 1edd kalllhon
medellledeeeecc.. (parame, X, Yr ***Implements predictive model.
perams: parameters of the model toput
second input
state:
curzent
model
Returns:
output: output of the model
state: updated state of the modai .**
if state is None:
stase jnp. zeros (10,))
Fill in code here
outpot Jnp.array (10, 8))
return output, state
F.2. Structured1 seed program
"nspinds fike hooo fuect hinery bandit task the rat dataset.
def acore_ funation (rat_id)
***oomputes normailaed likellhood
given model.
Retrieves saved datasets disk. Datasets are keya "choices" and "rewords", Each of these is list of lists of integers, representing the choices and rewards made by the ret on each sesston.
Json files
Computes twafold cross-validation, using avan-numbered and odd-numbered sessions.
Args: rat_id: Rat 1d.
Returns:
tikelincod of the fit model, in range 10, 3/. en*
dataset
irra____)_d
normal 11 sed__ 1hood twofold_erosswe idatioon (dataset, agent)
return normslized likelihood
28

---

Discovering Symbolic Cognitive Models from Behavior
def agent (params, choice, reward, agent_state):
""Cggiiiive model describing rac behavior on a binary two-armed bandit task.
Arge:
perams: Model parama. Dirferent paranetera are seed different rats.
choice: Previous choice. Values:
reward:
Previous reward. Values: or agent_atate: Previous state of the agent
Keturns:
choice_ logits: Vector of shape (2,) with the probabiiitles that the rat will choose option 2 on the next trial, expressed logits. stateestate: New state of the agent, after observing the previous choice and reward.
***
if staae is None:
agent_state jnp. 2aros((3))
param0 pazame [0] paraml (1 jap.exp perans (1)))
FITT cade here
cheice_iogits jap. (0,
return choite_logits, agent_state
F.3. Structured2 seed program
For conciseness, we use to indicate a section of the prompt that is the same as in Structuredl.
def agent (parame, choice, reward, agent_state)
not remove commence
progras.
# TCD0: Initialize agent state. agent_state is None: agent_state (ep.zerro(((2))
1000: Define all parameters.
porameter assignment
singie line.
Change the names of the parameters descriptive.
Include comments describing the parametez transformation, if applicable. Placeholdez perameter definition.
param) params
paraml (1 jnp.exp (-params # Placehnider poramoter definition with logistic function.
YCD0: update atant_
not update more than one elament at tise.
Sach component of the agent state update ahould be single line.
Include comments describing the computation in each line.
agent_state[1]
Compute choice_logits.
Bach component of the choice logit computation should be single line. Thotude comments desoribing the computation in each iine. choice_iogits jap.array 0]]
choise_logits, agent_state
E.4. FullModel seed program
For conciseness, to indicate section of the prompt that is the same in Structured1.
def agent (perams, choice, reward, agent_state):
29

---

Discovering Symbolic Cognitive Models from Behavior
agent_state is None:
agent_atate 3ap.zeros((3) (3,))
beta_r params (D)
beta_h ame
beta_g parame [2]
alpha_:
jnp.exp exp params [3]1]
alpha_h
jap
(411)
alpha_3
jap.exp (paaamss (511)
bias parama
choice rewazd 0/1
2-choice-1
foward_for_update. 2ereward-1
f Update 8
agent_state agent_atate..e (
algha_r , agent_ sgent_state[0]
=]gha_e) fowerd choice_for__pddts
Update
agent_state agent_ttate.s:()..ett
alpha_h + agant_state[[] + [1 alpha_h) chcico_for_updie
Update G
agent_state
alpha_g agant_state[[]
* plphagg)
lchoise_for_updan roward_for_upéate chcice_for_upáata,
cheice_term
Update choice term from choice_term +- beta_r agent_state(0]
# Update choice term from 8. choice_term
[gent_state(1]
Update
choi choice_term += bata_g agant_state[2]
Update choice from bias. choice_term +- blas
chaice_iogits array 11) chaice_term return choloe_logits, agent_state
30

---

Discovering Symbolic Cognitive Models from Behavior
G. Best Discovered Programs
Here we show the best discovered program for each dataset. Quantitative results from these programs are featured Fig. 1 (see Sec. 4.1). We note that different prompts elicited different best programs depending on the dataset: Best Human Program was produced with Structured2 prompt, whereas Best Rat Program and Best Fly Program were produced with FullModel prompt.
G.1. Human Structured2
This program was the overall best discovered program for Human. was generated using the Structured2 seed program and occurred after 1, 239, 338 samples.
def agent (params, choica, reward, stant_ None) num mum_params 13
params jap.clip (parame, -5, 5)
beta_r onp.ol wpp.clppppaee (wa..nn.sse.. softpl ua (parama [0]), 0.01, 201
lapse = cnp.clip nn..nn..... wn..nnn..anmm..... aigmoid (params [=]] 0.01, 6.991 prlor jnp.clippjax..i waw..n..ooftplai..
(params [2]), 6.31, 0.99)
Llpha_explozation_aad - jnp, clip (alp:jax..... aax..nn.. -n. aignoidd (parama 0,01, 0.99)
decsy_iate jap.clip(jex. wtx..nnn.iggmiiee.........o..... (perams (41), 0.01,
attencior_biarl params [5]
attencter_blas2 params
perseveration_strenne
fax.nn.affpllus (parama
switci_stzength parsms
lambda_param jnppccllpjaax..r. wwa..xn..ootppuueee.. softplus (parans [9]1, 3.0,
ganma - jax. iex.nn. (10)) Loss aversion pazameter
tempersture jap. clip. (jax. un. softplus (params [111] 1e-6, 1e-6, cemperature
beta_p jax.zn. softplus (parans (12])
agent_state is None:
g_values jap.ones(10)) ((1,) prior
old_choice
stiel_eince_las__ate 0
exploration_rate sipha__eplloaation.aa
cumchoice = 2ep.2eroa(14.)
else:
q_values agent_statel:4]]
choic agent_ststed
LilaLeeine_less_sv...o
stae:_
exploration_rate agent_state(6]
cuschaice agant_state[?111]
choice is not None and reward is not None:
delza reward gamma (1-reward) s_values[chhied g_values at
(a_vel_ee [caoice] [clovee
inp.whora (choice ald_choice, fribl_ainne_laat. lattoe_llst__ttttee exploratioon_ate exploration__att cuncheice cunchoice. at [choice] set teuncheice [choice]
1e-31 # decay exploration rate slowly
g_values = [1 apploration_ratal x g_values exploration_rate jap.mean (np.meannqq valuees q_values g_values decay_rate
chaice_probs (1 lapse) jax.nn bax.nn.softmar. [beta_r q_values terparsture beta_p log
1 cumchoicet) + lapse /
choice_iogits jap.log www.oogggcccccccc. (cobcc__robb))
choice is not None: peraeveration choice, num_elasses=4; switch bonus (choice old_choioe) switen rengtth jax. fax. one__hoo hot (chofce, run_elasaes=4) hot coldcchoce,, nus_olasses-()
(choice old_choice)
wttpvverrtioossstennnneee jax. www.na.one_eee
attention bonusl actention_bies
attention_ actention_bies? one (choice 2) 4, num_classes=1)
chaice_le logittt #
choice_logits perseveration
switch attent Lon_tiooobbbeeee attention_bonus2
jax.nn.one_hoo hot (shoica, + jnp. logitrial paatrraal_sirre sirro__lsst last_avited 13)
agent_state jap. concatenssel
Iq_values, arpayrray (Ichoice, trial since last switch, exploretion cumchoice]]
choice logits, ageat_etate

---

Discovering Symbolic Cognitive Models from Behavior
G.2. Rat FullModel
This program was the overall best discovered program for Rat. It generated using the FullModel seed program and occurred after 292, 690 samples.
def agent (parame, choice, zewazd, agent_state None)
agent_state is None:
ageat_atate jap. zeros (18)
alpha_= jax. wwg.mn..iiidiiiioo.mmmmm... (parems [01)
alpha_h jax.nn. w...iigmmoiie (params [1])
alpha_g jax. -n1 gmoodmdddgg (params [2]
alpha_q jax.na na.aa wox.na.aigmmen aigmoid [parama
beta_r nex softplus (parems
beta_h
softplus (params
beta_g fax.nn.softplue (parama
beta_g jax.nn. sofeplus (params
bias parama[8]
beta_g_bias params [9]
bela_h_bies params [10]
exploration_noise jax. softplus (pazams [111)
lapse_rate jox.nn. www.:iiiiigmiiiiioo.ommmm (porams [12])
genma_g w.k.na.zoftplus (params [13]]
ganma_N 1./ [[L_jxx..n.softtte..
(perams (1411)
alpha_a'as jax.nn. gmo (params| (15])
reward_bias params [155]
bias_l params [17]
alpha_" alpha_r 0.1
alpha_h alpha_h + 0.1
alpha_g alpha_3 0.1
alpha_q alpha_q 0.1
Dacice_for_update 2 choice
Saaard_for_updatt 2 reward reward
state - agent sgent_statees:1 { (eet(
alpha_r agant_state(())
alpha_r) reward_for_update choicee_oor for_dttt
agent_ state se (
alpha_h agent_state[[]
alpha_h) chcice_for_uppate
agent_state agent_state.s:12)..co
algha_g agent_state[2] alpha_g) @roloe_for_updcct
reward_for_update
state_g agent_staate3;$1) state_N agert_state
state_g state_g.at [chhice set setie state_g[chnice]
[1 alpha_g) tewara_for_update
alpha_bias gamma_g state_wll cheicell)
state_w state_x.st (choics), (elipha state_w(choice]
[1 alpha_g) raward_for_updas
alpha_blas x gamma_w (ttate_[[-chhotoe)
agent_state sgent_state.s:: (3t (state_g)
agent_state agent_state.a (state_k)
Q agent_state agent_stato([:3]
choice_term: -
agent_sta:ee0]
beta_h agent_state[]]
beta_g agent_s:ate[[]
beta_g (WIC] Q10] w[ll ((1))
bias
beta_g_blas (010) 0[11
beias ageat_s.ate[]] bias_]
chaice_logita {
lagse_rate l 0.5)))
(1 Iapse_rate) (rpaarra 11) choice_term)
zeturn choisa_logits, agent_state
32

---

Discovering Symbolic Cognitive Models from Behavior
G.3. Fruit Fly FullModel
This program was the overall best discovered program for FlyBandit. It was generated using the FullModel seed program and occurred after 195, 962 samples.
def agent (parame, choice, zeward, agent_state None) a_params_v2 = 10
if agent_state is None:
ageat_atate sepete [perans n_pazram_vv2
else:
agent_atate agent_aataa
1: jax nn. softplus (parems [1])
tempersture jap.exp (perams
lepse_iate = jax.nn. sigmeid (igmo1i(paassmss]111
alpha exp.exp (pazems [41) sipha
bias jax.nn. jaw.nn.softpliae softplus (porams [5:7]] # 8ias should be pasitive
torget_rate wax.rn.softppp.... [p])
decay_sate woltp.ussolluuaa..ooo.oma (paramss
exploitation_rete jaw.nn.eigooddd pprgmii
g_vals agent_state
rpe [_vess[choiee] P_choice = flv.mm.clgmotd(hiaaall [q_vala[1 - choica] __vall[(cooico)
decty_facter jnp.ed exp dasey_race)
deita_gval 1r xpe p_choice
v_alo_s___ooooootttt s_ (choicc)) (g valv [choice) delta_gval)
updated_g_ jax.lxx.selec:
_evatedd___vllee luea_chaioe, (q_avla) decay_factor)
reward, updated_s_va
choice).sets
updatedd___values__ cheice]
forget_rata jop.naximum() (0, rpe (1 P_choical fecty_
this iine reduce averwriting
podaate____valaee jap.clip (xpdaced q values, -alpha, alphs! updated___vallcee updated_g_ v_alues 11 lapse_rate)
choice_iogits
exploitation__rat rate jxp.exp (bias
(blass + temperature updatad_ q exploltation jnp.exp (blas)
1e
return choire_logits, u
H. Seed program comparison
Here we show example best CogFunSearch programs that are discovered for the Rat dataset under different seed programs (Sec. F). We note that these programs differ from the best found programs, as they use a later cutoff. This earlier cutoff was used here to permit comparison with programs from other prompts given similar number of samples.
H.1. Rat LowInfo
Best LowInfo program given a cutoff 165,000 samples.
def nodel (params, x, stata None):
hparams len (params)
jap.max
arzay (1 [Ljap.abs abs
0.5, jnp. 0.51))
jnp smpp.ggrt/ (jpp. (x))
jnp..agrs.. (jap. abs
x_clip jnp.cl (x -20, 20)
y_elip jag.clip(y, -10, 10)
*_olip (p -20, 20)
state is None:
state jnp. zeros (3,))
atate_mean jnp. aap.meen (atate) (params e_clip perans x_clip pazams y_eelpp
33

---

Discovering Symbolic Cognitive Models from Behavior
params [3] state_ream params (4))
(pacams [5]
X_clip pazama [6] + jap. (y_clip) params [7] z_clip
[8]
state[l] params [9])
output jap.
alpha ampp.arryylloo [xx,, x2]1
beta 0.7
ganmo C.9
atate inp.array
gamna (xl
(1 gamma) atate[0] params [10] x_elip x_elip
params [1l] y_clip y_slip,
beta x 1x2 + 1) + (1 beta) s1ate(1) parama [12] x z_clip +
pazams [13]
alpha
params
x_alip params [15] tp params state(2)) 1 inp. heppebss(_eiippp
[14] p14)
sipha + (state[2) (y -
pazams [17] x_elip y_elip:
return output, state
H.2. RatBandit Structured1
Best Structured program given a cutoff of 165,000 samples.
agent (params, choicce, reward, agent_sstate
None)
aum_params 16
alpha - jax.nn. softplus (pazams [0])
beta jax.nn. softplus (params [1]]
decay - ne...n.. slgmoid (pacame [2]]
initial_bies jax. softplus (pazams (31)
bias_dacap jax.nr. mox.nn..50 wow.nn......ttii.. stw.softtluu [params[4]] atr.
exploratiion_otte
jaw.nn..oot... wwt.mmooottpaaaaaommmmmmmoo (parmmm (5))
alpha_pe__os- jax.na.ao..e.ocoooooo..aooooo.oo www.nna.olluu (pa
alpha_pe_neg jax.nn. nn. www.nnn.sfftll. softpllst (params [7])
decay_pe jax.an. w...nnssoouusss (parama
tempersture = jax. An. [parems(]]] 0.1
lapse_rate jax.nn.e. www.nn....gmiiiie. (params [100 [10])
initial_g_vall
(params [11]]
choice_bias params [12]
exploration_tatio
www.nn....lluuuuppcoomm (pazams [13])
initiel_bias_directi params [144
beta_pe max.na.soflpluc (pacams [15])
agent__ctate None:
'nitial_0 = full (2, initial_g_vae ___value]
initial_bias jax. sigmoid (initial blas Inp.array
[initisl_biss_direes direction,,,
maitial_bias _ias__drrrettoo)))
initial_pe jep.zeros
ageat_etate
'V_values': initiel_g,
"bias': initial_bias,
(PE': initial_pe,
delta = reward agent_stato 'O_valuas'' [choice]
valaue agent_ state ('_aattl' Lat [choice] sett ictt
agent_ste atreaaa'GG..e va__vavlsss [chofce] alpha deltal
vavele decay) Q_valuea decay v_iliel
bias (1 bicaydecayl agent_statt("biaa"" blas += initial bias (np.array choice, choice])
po_pos jap. anp.moximem (Jelma, 0)
pe pe_neg -jnp. (deltà,
decay_pe) agent_stale(''E']
pe sipha_pe_pes pe_pes jnp. array [[1 choice, choice]) sipha_po_nag pe_neg orpay ([1 choice, choica])
exploration_noise exploza.ion_noizs jax. random jax, waandom.PRMcKey (0), iuploortioo_ratee shape= (2,)) exploration_noise exploracion_noiee jap.neaa texplorat ion_nlaaee
cheice_iogits beta (L-values bias
temporature exploration_noile pe choiss_bias (
[1, -1]) + beca_pe pe
choice_iogits jap. www.cclpphhiceeooooo logiiec 10)
34

---

Discovering Symbolic Cognitive Models from Behavior
fAdá Lapse lapse_prob. rape /2 cheice_iogits choice_logizs lspse_prob papse_prab
agent_state {
'0_valuea': Q_values,
'bias': bias, '2E": pe,
choice logits, agant_state
H.3. Rat Structured2
Best Structured2 program given a cutoff of 165,000 samples.
def agent (perame, choice, reward, agent_state
None)
alpha www.trr.oofpliooo.o (parama (0:2])
inv_temp = jax.an, an..an.sofplll softplus (pazana [2])
lapse_rate jex. an. softplus (pazams (3])
Changed sigmoid softplus
stickiness jax.nn ar. softplus (params [4])
decay_rate- sox.nn.softp (parama (5)) Changed sigmaid to softplus
omega jax.nr..afft aoftplus (params [6:8]) two osigga parama
bias parama [8]
init_q params [9:11]
sigms - jax.nr. softpll (pazams
explorstion_tate ja..en..e.................. soštplus (params (121)
momencum = sox.nn.coftplus (params (13]
nax.nn.softs (params [14]]
lapse_mode wtw.nn..ooolllli (params [15]]
lapse_a pha jax.nn, sontsooflllue us (parars(16))
agent_state I
agent_state
linit_q, jnp.2eroa 0.5,
(np.array((055)) jap.array|1011)
vavues agent_8Late(:2]
stickiness_state sgent_state[2:10
lspse_rate_state agsnt_atate(4)
trial_count agant [gent_state[S)
reward_noisy reward lapse_mode
jax. www.randm.PPRcKee.... (U)CKey()), shape- () aigna
reward_noisy jnp. eepeellpppeeee niiy,
deita = reward_naisy [choicee
@valuea_updatee [hp.wh choice 0, jap.array!
[g_values alpta [0] della,
[_]] cmega alpha[0] delta facto_)
anp.array(
(eevalues omega [1] alpha dalta omega_factor,
(_válues[1] []] alpha[1] delta]),
Updated stickiness calculation
stickiness_state
12 Soacay_rate) stickiness
stickineas reward_nolay) jax.nr.one__he (chofic,, 21)
choice_iogits inv_temp q_saluss_updated stickiness_etate
alas x ((1.))) -1.0]1)
explorstion_term
exploration_rap jnp.expt-trial_comm / 20) arp.array (lg_ [1),
[_vallee__pdteee(()))
choice_logits choice_logias exploration_tezm
updated_lepse_ratt
[ lapge_alphat lapso_rate_atate
lapse_alphs A jnp.clip/lapaeeratte 0.001, 0.9991
P
F updated_ Baaddd_lpa________e_aee tax.nn.soffren (chooot
updated jap.arzay (10.5, 0.51))
35

---

Discovering Symbolic Cognitive Models from Behavior
p (p 1e-8, 1e-8)
chaice_iogits ip / (1 pl)
v_vaue - uomentum L.vallue_uuddttee (1 momentuni q_values
g_values rate)_rate) g_valuas decay_rate irit_e
tsial
stant_state jmp imp.consatena:el
(g_values, stickiress_state,
@apse_rale_ataee [None], trial_count [None]]]
return choice_logits, ataet_state
H.4. Rat FullModel
Best FullModel program given cutoff of 163,000 samples.
def agent iparams, cholica, reward, agent__state None)
alphaa Jax www.siee... sigmold [pazams + 0.05
params [5:10]
bias_term params[10)
tan www.an.softplls...... (params [11]1
lapse_rate jax.nr. sigmoid (params (12]) 0.25 0.0
agent_state
ageat_etate jnp.array((0.,, 0., O., O., 0.5, 0.51)
for_update = 2 choice -
foward_for_uddate
joward_prediction.eees rewari agent_staaeef
agent_state
slphas[0] agent_state[[]
[1 alphas [0]) rewar__rradcction.errr
tordoa_tor_update (1 agent_state(s])
agent_state
[lp agent_atate[l]
[1 alphas[]]) choice_for_update
agent_state st
alphas [2) agent_atate(2] t f1 alphas +
ceward_for_update choice_for_update
chcice_for_updatss
agent_state seat__taate.::33...r
a3phas[3) agent_atata
[1 alphas[3]) reward_predlotiioneerro agent_atate[5])
agent_state sgent_state.s:(4..ce
alphas[4) agent_stata[4] s1phas[4)) reward_prediction..ess
agent_state(S))
agent_state- agant_state.s: .set
alghas[4] agent_atate[5]
(1 s1]h reword_prediction.eee. [1 egent_ssatte[]]
choice
betas(0]
agent_stato[(]
betas[]] agent_state(]]
betas[2] açeat_state(2)
betas[3] agent_state(3]
betas[4] agent_state[4]
a'am_ / tau
cerorce_erre (betas[0] agent_ [0] betaaff] =gger_sstar((1] #betas[2] agent _ggnt_asated state[2] betaal3] agent_ scate(3) #betas!4] jap. mwm.maxximmmmm (aaemummtggen__etttell]]
biss_term) tau
choice_'ogits lapse_rata jap.orrayt
[ohoice_term) -choice_teral) lapse_rate log (Jnp. ((U.5, 0.5113 choice_logits, agent_state
36

---

Discovering Symbolic Cognitive Models from Behavior
I. Example Programs that vary in Complexity
These are all programs sampled from the complexity/accuracy frontier, as discussed in Section 4.2 and illustrated in Figures 13.
L1. Human Example Programs
Here we share the Example Programs for the Human dataset, indicated in Fig. 13. These programs are ordered from least complex lowest scoring to most complex highest scoring, along the frontier. We provide brief description the code's function, for convenience. These programs were all generated with the Structured2 seed program, which also generated the best Human CogFunSearch programs.
Example Program 1. This program substantially overlaps with the PFQ baseline, exhibiting Q-learning and forgetting. Interestingly, it does not exhibit any perseveration (it records previous choice in its state, but never uses it). The key addition in this program is optimistic initialization in which all Q-values start at 1.
agent (pacame, choice, reward, agent_state):
agent_state is None:
ageat_atate jmp. ones ((5,)) & stare all Q-values at
beta_r jxp.exp (parama (0))
alpha_leara - jax. www.snnsoooplii softplus (parans [1]]
elpha_forget jax. jax.nn. nol.nn.soooe softplus (pazams [2])
q_values sgent_state(:
# Update 0 values
updated_qvallee Lqvalues q_values alpha_learn) alpha_learn reward podaate__qvalues qpvauee_ [chcice].set]
updated [choicel alpha_learn (reward a_vallus(chhiiee
Decay values becayingggva lees updated_quv qualuge jnp. lepbexxaapppaaaargget
# Compute choice probabilities using softmex
chaice_iogits beta_r dacayingg_valuus
Update agent state.
new_agant_atate concasenatt pvalus, amp.array
thoice_logits nev_agent_atate
Example Program 2. Example Program 2 is similar to except that initial Q-values optimistically initialized to 1, and forgetting decays toward 0.5 rather than O (as in Example Program 1) or the mean Q-value (as in the best CogFunSearch program). This program introduces an additional parameter compared to Example Program 1, lapse rate, which also found in the highest scoring CogFunSearch runs and evokes the classic psychology literature.
fittable parameter instead of
agent Iparame, choice, reward, agent_state):
# Deffine parametera
beta_r jax.nn. soltplus (parans [0])
alpha_leara jax. si signoid ( (params [1]]
alpha_farget lapse_rate jax.nn.a gmot (params [3]]
sigmodd (parens (2])
initial_g jax.nn. nn. sigoo [4])
roitialize 0 values close to 0.5
#initialize state
stent_state is None: agent_state jnp. repeat (initial_s,
vals agant_state
37

---

Discovering Symbolic Cognitive Models from Behavior
Update g-values
prediction_ oorror reward g_vals [choice]
g_vals at [choice].sed (Lvaice.
[1 alpha_laarn) g_vala[choien alpha_learm reword
Forgetiing process
g_vals g_vals 11 sipha_ forget) + 0.5 sipha_ forget Adda forgetting that favors 0.5
# Incorporate lapse rare chaice_probs
[1 [ppee_rate] sox.tn. (heta__ (_vals)
lapse_rate (1.0 4.01 v onp.ones (41)
choice__ogits jap. log (croice_probs)
choite_logits, g_vals
Example Program 3. is considerably higher scoring and more complex Figure 13.
Like Example Program 2, this program also includes fittable initial Q-values. Both programs also track previous state, but don't actually use to inform any computation. The practice of decaying Q-values of chosen actions toward mean Q-value, which is present all of the top CogFunSearch programs, also introduced in this program. This program uses recent reward to modulate exploration, with more exploitation if the immediately preceding trial was rewarded.
def agent (parame, choice, reward, agent_state):
param_names = ('beta_r', Dipha__earn' 2ipha__ooget! tomega", parama_dict diet fsip (param__ames, pazama))
"initug")
agent_state is None:
Default
ageat_state
0 initialize
agent_state agent_atata [at[[4]]]] o ctw.ccoooffpllaaaprrre.didddddd (iriiiiiiddctiiiiaa"""a ('init_a'])) Set volues
agent_state agent_atate.at s4).
#Dffine parameters
beta_r jax. softplus (params_dictte "bata_=']] 0.1
alpha_leara jax. an. sigmoid/params
(params _icttdddiitiie diairmmmddict"eelhaaleeeaas ['alpha_leara"]]
flrge_forget jax.nn. nn, si sigmaid pparans_dicttt ("lip''"aaaa_oorg="""
omaga jax.nr.saftplus parama_diee ('omaga']]
temerae jax. on. sofuplus (params_dict ('Liverse_temeeratuue 0.01
init_q jn. s.x.tnls (params_dee sicctms_dii____e "init_g"1)
f Update agent state v_valu [gent_state[:4] previous_choice [gen=_ataze[4]
q_values q_values.e [choice].seti
g_values [choices[choice] [choice] sipha_lesrn (seward q_values (choice))) Update Q value for chosen arm
alpha_ forget_mttttt 1.0 alpha_forgat
v_veue
q_vauues alpha_forget. www.ettrrriiiteo mult
alpha_forget inp mean valluues) Apply forgetting weighted by overali mean 0 value
#Compute choice logits
Apply inverse teeperature and reward modulation choice_ ogita beta_r v_valuea / inverse_temperatee 11 omega rewarda*2) choice_probabiillless jax. softmaxx cha axl 0) +Apply softmax for probabi iea
chaice_iogita lng.log
log htt.chhcc..rooabilttiie.. 1e-8)
state - jap. concatenccee
vq_vals arpay ([thoice]))) Update final
return logits, state_
I.2. Rat Example Programs
Example Programs for the Rat dataset, indicated in Figs. 7, 13. These programs are ordered from least complex lowest scoring to most complex highest scoring, along the frontier. We provide a brief description of the code's function, for
38

---

Discovering Symbolic Cognitive Models from Behavior
convenience. These programs all generated with the FullModel seed program, which also generated the best Rat CogFunSearch programs.
Example Program 1. This program preserves the terms R, H, and G, expanding them to reward, habit, and guess. It is notable that the LLM maps in the seed program to habit, as the word habit is not included anywhere in the seed program or initial prompt documentation. It maps G to guess, which is different from the RHG baseline where it refers to gambler's fallacy. This program has only parameters, compared the RHG's model's seven, and same learning rate used across all state updates. The code shows similar update RHG, except instead of tracking the recently rewarded choice, simply tracks recent reward. Then, all state variables are summed.
def agent (params, chol co, reward, agent_ statt_
agent_state is None:
agent_sttae
((2,))
v_reward - params [0]
v_habit params
v_gueas parama [2]
bias params [3]
learning_rate wwx.an.sigmoidpparamm
prev_choice
chol
prev_reward - 2 zeward
# Update reward prediction. agent_ state sttt:tssttteaal0l.s.....
[1 learning_rats) agent_statts[] learaing_rate prev_reward
Update choice value.
agent_ state agent_ a:lte.aal1).ae
[1 learning_rate; agent_atate(1] learaing_rate prev_cheice
update guassing value.
agent_state agent_state.a: [2]]
[1 lesrning_rate: agent_state[[] learning_rate (prev_choice prev_reward prev_choioe)
Compute choice logits, choice_term
v_reward agent_state[[] v_habit agent_state(1] V_guess agent_state[[]
bias)
chaice_iogits jnp.array ([-1, 1)) chaico_term
return choice_logita. agent_state
Example Program uses the same parameters as the RHG baseline, plus an additional temperature parameter controlling exploration. It also defines 7 additional state variables that separately track reward and omission rates, well a complex interaction of reward and choice.
def agent (params, choice, reward, agent_state):
a_params params. params.shape[U) state_ is None:
ageat_state zerozeros (eros(((1,))
alpha_r jax, max.cllmpp0... 001, 0.995, jax.nn. 1ggo (params [0])]
alpha_h jax. c0.
0.99,,
jax.nn..s sinmoi (params [1]
alpha_g jax. www.lax.caaa.. cww.ccampp0..ooommooooo 00001, 0,999,
jax.nn. gmo (pa rams [2]
beta_r cax soltplus (params
beta_h jax.nn. sex.rn.softpuu (params [4]]
beta_g jax.nn. softplus (params
blas_params = parama (6:8] blas now a vector
inverse_Lemperttun jax sof.plus (pacams
choice
fwarrd_oo__uddst reward
delta_r = reward for for_for__ppeee update proice_oo__upee
agent_state agent_state.. [0] set (alphe_r agent_stateef
alpha_r) della_r)
39

---

Discovering Symbolic Cognitive Models from Behavior
deita_h cheice_for_ priice_oor_uddst
stant_state agent_state... == .ee (alpha_h agent statee]
alpha_A) x delta_h)
della_g cerrocefor__uiiie ravare_ forrre_fo____pee cor_update carooce_________eeee agent_state - agent_atate,at [2] aet (alpha_g
(1 alpha_g) delta_g)
# Separate parameters for biases.
Exponential decay, with a dynamic decay zate depending on the zewazd expectation.
faay__ccte_for_ewweed jap.clip(0.2 (0.2 0.5.inp. tank decay_rate__fdt htt__raatt_frr__msssheeeee jnp.clip(0.5 D.5xinp, tanh state 0.1, 0.91
0.1, 0.9)
stane - [qent_state.s.(3]
rary__rat__aarraaaa_ !
agent_atate[3]] reward_for__pdatt choice_Cor_update bias_pacamm(01)
Bias for choice
agent_state agent_state.s:(4).ees
decoy
agent_atate[4]] reward choioe_ror_update blas [aa___arams]]
# Bias for choice
agent_state state 1
0,99 agent_state[6] 0.01 [eward_[oo__uddate)
agent_state state agent_state (t .set
0.99 agent__taata[]] 0.01 (l-roward_for__pdnee
staet_ agent_state({] Rchanged
chaice_term [bsta_r agent_
agent_staze[[]
bete_g agent_state[3] state_l sarams[-1]]
chgice_iogits jap. 11) chaice_term Laverse_cengersture
return atoice_logita, agent_state
Example Program 3. This program has very similar code to the RHG baseline. It includes a different parameterization of the parameters, which constrains all learning rates to be clamped at 0.01 and 0.99, lacks bias term, and includes step where tanh applied to the choice term. Charmingly, this tanh step is accompanied by a comment claiming: "This is different from the other version, potentially an important improvement". Verifying this claim would require future investigation.
def agent (perams, choice, reward, agent_state):
if agent_state is None: ageat_state jnp. zeros (3,))
alpha_: jax wal.la..llaaa clamp 20.00 01, 0.39, nn. seftplus (params[0]]
alpha_h jax. jax. lax. clamp 01, 0.39,
jax.nn. seftplus (params [1]])
alpha_g jax lax. clamp (0. 01, 0.39,
jax.nn. seftplus (params (2]))
beta_r jax. softplus (parans
beta_h jax.nn. scftplus (params (4])
beta_g jax.nn.softpiiln (params (5])
# Convert choice and reward 0/2 to
dacice_for_updat 2 reward
choice 1
caarrd_for_uppcttee
Update of values 8, H, 6.
ta_r reward for por_updata
caatee_oor_uudttte
agent_state agent_atate.a: (alphaa_ agent_state(0]
alpha_r) delta_r)
deita_h cheice_for_updee agent_state agent_state.ss (allpaa agent_ [gent_ssate(l)
alpha_h) delta_h)
agent_state agent_state.a; rawarrd_oo__ueeppeeeere for Sor_uppatt choice_far far jar_update
choice_for_updhte della_g choice for for_updata- raward
(alpha_g aqent_atate(2]
alpha_g) delta_g!
# This ia different from the other veraions, potentialiy Important Improvement.
choice_term (ap.Leah(
beta_r agent_stszee01
beta_h agent_ataatel]]
beta_g agent_atatte[1]
pa:ams[6]
params [7] Aloice_for_updete
params reward_for__pdae
choice_" logito array 1]) cheice_tern thia do better aoftmax
return ohoite_logite, agent_state
40

---

Discovering Symbolic Cognitive Models from Behavior
1.3. Fruit Fly Example Programs
Example Programs for the Fruit Fly dataset, indicated in Figure 13. These programs are ordered from least complex / lowest scoring most complex highest scoring, along the frontier. We provide brief description of the code's function, for convenience. These programs were all generated with the FullModel seed program, which also generated the best Rat CogFunSearch programs.
Example Program This program closely imitates the DFQ baseline, adding a bias toward which the Q-value of the unchosen action decays.
agent (params, ce, reward, agent_sttate)
agent_state state is None: agent_state jnp.array 0.0))
Unpack rename pazameters readability.
alpho_learn = jax.rn.softplin softplus (porans [0] Ensures alpha_learn is positive
alpha_torget jax.nn nax.nn.aoftplllooooooooooooo...n softplus (params [1]) Ensures aipha_forget forget 1a positive
kapps_ieward jax softplus (pazams [2]) ? Easuzes sewaad is positive
kappa_omiaator jax. on, www.nnn.ooooo.... oolsplus (parsms [3]] Ensures rappa_onisaton is positive
[4]ama[4]
- agget_sttate
# Chosen action's (-value wfll i be updated cowarda this target vo lue. learn_target kappa_reward reward kappa_omisaion (1 reward)
# Update O-value for chosen action.
aew_q qs [choice) aephe_learn (leern_target qa [chhoice)]
qs.at [choice] .set set (zew_g)
# update Q-value unchosen scrion: decay towards bias.
[s.at choice] set [q choice] alpha_forget (qa[1 chofce] bias))
for initial preference. logits jax..n. softplus t blas) aew_agent_state gs
# Add bias
zeturn Logits, naw_agent_state
Example Program 2. This program closely imitates the DFQ baseline, adding a bias toward which the Q-value of the unchosen action decays. Like other discovered programs, including the best CogFunSearch Fruit Fly program, it includes a lapse rate term, Like the best CogFunSearch program, this program includes somewhat unusually complex exploration policy.
agent (porams, choice, reward, agent_sstate)
stant_ is None:
agent_state (np,array 0.5))
Unpack parameters
alpha_learn_lon params(0)
beta_logit params(1]
kappa_reward_loyi_ params
keppo_omlestor_lagtt params [3]
tau_logit parama[4]
lepse_iate_logit params[5]
elpha_farpet_logi:: params
alpha_learn taw.rils paaaaaa
beta tbpteeptbeeta_ooooo
# enforce beta >
kappe_reward nox.nt..oftplie Iaa peoppp_rewar_____ (ogit] logit) kappa_onission jax sofiplus (keppa__misss:on_oggil zau wow..am.oftpllsstsuuuuooo
lapse_sate jex.on. sigmeid (lapse_rate_iogit
liphafforget jax. softp haa (alpha aapphaafrr_______1_1__ ogitt)
agent_state
update g value.
learn_target kappa_reward reward + kapps_omission (1 reward) (leart_targat -a (cchiice))
us as set (qs [choice) Leaha_learn eιzoz)
qs.at[1 choice] set [qs[1 choice] forge_ esror!

---

Discovering Symbolic Cognitive Models from Behavior
inp clipiqs, C., 1.)
q= q= 0.001
# aod lapse rate, ensuze that lapse_ rabe is not 1.0 logit (1. lapge_rase) tau [0) beta go beta) Lapse_rate jrp.log(0.3. (0-5)
aew_agent_state qs
choice_iogits imp.arrey(((ogg) -logit]] return logits, new_agent__ttate stata
Example Program This program again includes the functionality from the DFQ baseline, and in fact closely resembles Example Program also includes an unusual, dynamically updated exploration term which aggregates the difference between the reward and the lapse rate, and uses this to scale eploration.
def agent (porams, choica, reward, agent_statco
if agent__ctate is None:
agent_atate anp.array(10.5,, 0.5, 0., 0.)1
# Unpack parameters
alpha_lears_logit params
beta_logit params [1]
kappa_reward__oge params
lappo_omlosiar__lagtt porams [3]
tau_logit [4]ama[4]
lapse_cate_logit [S]
alpha_ parems [6]
alpha_lapso__loit params
alpha_leara jax. soltplus
beta (bgg kappa_:eward
enforce beta >
sof:plus (ksspa__cmisssioooooggg keppa_omiesien: (ogu__ ) zau jax.an.softplaa. wtw.:mmmmmlllpoo_rttttttto
lapse_rate jax.nn.aofle aatpplluaaapphh_foqeee_..
alpha_torget alpha_lapse iax.nn.soltpiu.. lelpha_
qs - agent_state((12 lapse agent_state(2] lspse_rate_ agant_state
# update q value.
learn_target kappa_reward
kapps_onission
reward)
exror = (learn_target q[(chiice)
(shoics) set (gsichoice) alpha_learn error)
= [1 choice] .set [1 choice] alpga_forget. + error)
inp enp.elip(qs, C., 1.)
qs qs 0.001
# update 1sose rate. lapse alpha_lapss (reward lapse)
Japse zale,
lapse_rate Is
logit = (1. lapse) tau [qs[0] va beta qa[1] av betal rate log (0.5)
gov_agent_statt arpay (0g qili,, lapse, lapse_rate]) choice_ choico_logita ogiit arrayy -lagit]] return logits, saw_agent_atate
42


Reviews:
[{"review_id": "g7St2rIuqZ", "reviewer": "Reviewer_4sbb", "summary": "This paper introduces CogFunSearch, an automated approach to discovering symbolic cognitive models that accurately describe human and animal behavior. The method builds on FunSearch, a program synthesis tool powered by Large Language Models (LLMs) and an evolutionary algorithm, to systematically explore and optimize symbolic cognitive models.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful review of our submission and their useful feedback. We are glad the reviewer found the submission “well-structured, comprehensive, and methodologically rigorous”. Below we provide responses to some of the main concerns raised.\n\n## “Can researchers leverage these models to develop new methodologies?”\n\nComputational cognitive models for tasks of this kind are widely used for interpreting neuroscience data (do the internal variables of the model have correlates in the brain?) as well as understanding the effects of causal experiments (does damage to particular brain regions affect behavior in ways that are similar to ablation of the models?) and naturally-occurring variability (do patients with particular psychiatric diagnoses differ systematically in their model parameters?)\n\nA widely-held belief is that models which better capture behavior are more likely to be useful for these kinds of neuroscience applications. Exploring whether this is true for the models we have discovered here remains a direction for future work, but it is one we are optimistic about. This is especially true for the models discovered for the human bandit dataset, which fit data dramatically better than existing models and which contain structure qualitatively unlike the ones in those models.\n\n## “Do these models generalize to novel datasets, or do they merely excel in fitting existing data?”\n\nIn our evaluation we do cross-validation over held-out subjects. This can be considered as a form of generalization to new datasets, where the datasets come from new (held-out) animals performing the same task. Whether these models can generalize to entirely new tasks or different organisms is an interesting question for future work, but was not within the scope of our work here.\n\n## “Do the variable names in the high-performing models genuinely reflect their functional roles as implied by their nomenclature?”\n\nThank you for raising this question. We have manually inspected the top performing programs and found that, for the majority of cases, variable names are semantically meaningful with respect to their functional roles. We find that the meaningfulness of the variable names tends to correlate with the strength of the seed programs (i.e.starting from FullModel tends to result with more meaningful variable names versus starting from LowInfo).\n\n## “Lack of Comparison with Other AI Discovery Methods”\n\nOur focus in this work was on whether LLM-based discovery methods could produce accurate, yet interpretable, models. We were particularly compelled to explore this due to the fact that LLMs are able to utilize the semantics of its inputs (such as variable names) to guide its generation. Some of the alternative approaches suggested would lack this semantic “understanding”.\nNonetheless, we do agree it’s an interesting alternative to explore and will add a discussion of these ideas to our conclusions.\n\n## “Do you believe that similar methods could be successfully applied to other domains, such as physics, chemistry, or economic modeling?”\n\nYes! At the core of our approach is using funsearch to discover parameterized programs that can be fit to a dataset, and this can be applied to other areas of data-driven scientific discovery. One recent work in this vein is [1]. We will add a brief discussion of this extension in the final version of our manuscript.\n\n## “what characteristics of Symbolic Cognitive Models made this field particularly suitable for your approach”\n\nIt can be argued that this  field has relied heavily on models fit to data of one type (behavior) to understand data of very different types (neural activity, psychiatry, etc). This field is in crisis: deep learning has made it possible to benchmark longstanding models and find out that they just do not fit very well. It's therefore both theory-poor but also well-equipped to make use of new theories should they appear. Increasingly, it is also data-rich. These considerations make it a good place for data-driven theory discovery to have outsized impact. Further, the semantic-grounding of LLMs (discussed above) and the human-interpretable nature of Python programs allow us to leverage, and go beyond, prior cognitive models.\n\n\n[1] [Grayeli et al., 2024, “Symbolic Regression with a Learned Concept Library”](https://arxiv.org/abs/2409.09359)", "claims_and_evidence": "The article makes three key claims:\n\n1. The models discovered by CogFunSearch outperform state-of-the-art human-designed cognitive models in behavioral prediction across humans, rats, and fruit flies.\n2. CogFunSearch can explore vast model spaces, identify solutions superior to human-designed models, and provide novel insights into cognitive mechanisms, with strong supporting evidence.\n3. The discovered models remain largely interpretable.\n\nIn the conclusion, the authors state: “We find that CogFunSearch can discover programs that outperform state-of-the-art baselines for predicting animal behavior, while remaining largely interpretable.”\n\nHowever, the claim regarding interpretability requires further scrutiny. The paper primarily justifies interpretability based on code readability and the intuitiveness of variable naming but does not address several critical questions:\n\nCan researchers leverage these models to develop new methodologies?\nDo these models generalize to novel datasets, or do they merely excel in fitting existing data?\nDo the variable names in the high-performing models genuinely reflect their functional roles as implied by their nomenclature?", "methods_and_evaluation_criteria": "In summary, the proposed methodology is well-justified, and the evaluation criteria are scientifically sound.\n\nThe ​CogFunSearch framework demonstrates ​innovative potential by enabling ​automated hypothesis discovery, thereby ​reducing reliance on manually designed cognitive models. It has achieved robust performance across three distinct behavioral datasets  while adhering to rigorous evaluation standards.", "theoretical_claims": "This paper does not primarily focus on formal theoretical proofs, as its main contribution lies in data-driven symbolic cognitive model discovery using LLM-powered evolutionary search. I also did not find any obviously problematic theoretical statements.", "experimental_designs_or_analyses": "The experimental design of the paper is well-structured and suitable for evaluating the discovery of symbolic cognitive models. The datasets, evaluation metrics, and statistical analysis methods provide strong empirical support for the paper's key conclusions.", "relation_to_broader_scientific_literature": "The key contributions of this paper lie at the intersection of cognitive modeling, AI-driven scientific discovery, and program synthesis. It extends the FunSearch framework, applying LLM-guided evolutionary search to cognitive science, enabling automated symbolic model discovery. Compared to traditional methods, this approach reduces reliance on human intuition and automatically discovers models that outperform human-designed counterparts in predictive performance.", "essential_references_not_discussed": "I did not find missing references; the citations in this paper appear to be appropriate and comprehensive.", "other_strengths_and_weaknesses": "This paper is well-structured, comprehensive, and methodologically rigorous. However, there are a few areas that could be improved to enhance its practicality and impact:  \n\n1. **Computational Efficiency Issues**: CogFunSearch requires **hundreds of thousands to millions of LLM queries** to find the optimal programs, resulting in extremely high computational costs. Although the paper mentions **rejection sampling** as an optimization technique, it lacks a detailed analysis of **computational cost vs. predictive performance gains**.  \n\n2. **Lack of Comparison with Other AI Discovery Methods**: The paper does not compare its approach with **more computationally efficient AI discovery frameworks**, such as **differentiable architecture search, Bayesian symbolic regression, or reinforcement learning-driven symbolic search**. Including such comparisons would help clarify the trade-offs between computational resource consumption and model quality.", "other_comments_or_suggestions": "See above.", "questions_for_authors": "**Q1**: See Claims And Evidence\n\n**Q2**: \nYour work demonstrates the effectiveness of LLM-guided evolutionary search in discovering symbolic cognitive models. Do you believe that similar methods could be successfully applied to other domains, such as physics, chemistry, or economic modeling? If so, what characteristics of Symbolic Cognitive Models made this field particularly suitable for your approach？", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "56YV8OPq8j", "reviewer": "Reviewer_nJ74", "summary": "This paper extends the FunSearch evolutionary algorithm to autonomously uncover symbolic cognitive models that effectively represent human and animal behavior. The authors compare the top-discovered program with an RNN trained on data from all subjects collectively, showcasing the efficacy of CogFunSearch. Additionally, they explore the trade-off between quality of fit and program complexity.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful review of our submission, their useful comments, & are glad they found our proposed methodology “particularly impressive”.\n\n## RNN trained on data from all subjects…rather than…separate RNNs…\n\nTo clarify: This concern applies only to the human bandit dataset. We find that per-subject RNNs highly overfit to the training data & evaluated poorly on validation data, likely due to the small number of sessions & trials per subject. **We have revised the manuscript to include fits of RNNs to individual subjects’ data**, with a cross-validation setup comparable to CogFunSearch. We sweep network hidden sizes, used early stopping to mitigate overfitting, & find that the avg normalized score is 40.27$\\pm$0.50, which is considerably worse than the cognitive model baseline (55.55$\\pm$0.64), the best discovered program (60.93 $\\pm$0.63), and the multisubject RNN (61.83$\\pm$0.81). For the rat dataset, RNNs were trained separately for each subject. For the fly dataset, both RNNs and CogFunSearch programs were fit to the entire dataset, since each fly performed only one session, so it was not possible to compute held-out per-animal scores.\n\n## does building computational process models contribute to…scientific understanding…?\n\nWe thank the reviewer for raising this important philosophical question, which we agree is highlighted by our work. Computational models developed by humans reflect human-developed theories about the mechanisms at play, and are traditionally seen as tools for exploring the implications of these theories. These models can be used to make quantitative predictions about behavior in different  situations as well as different data modalities like neural recordings. They are also commonly used to understand differences in behavior, both experimentally-induced or naturally-occuring.\n\nWe believe that automatically-discovered models can play a similar set of roles, so long as the theories that they express are interpretable to scientists. They play an additional role in surfacing new ideas that might be different from those that would occur to researchers, as we discuss in section 4, as well as in accelerating their discovery.\n\n## How valuable is it…to dedicate years to mastering a skill that LLMs can now acquire with ease?\n\nWe agree that as technology evolves, the set of skills that will be valuable to learn will also evolve. We see methods like ours as adding to the toolkit available, but not necessarily as reducing the skillset required of our students. Many of the skills now required in order to build and use computational models are still required to interpret & use automatically-discovered models. We agree with reviewer fWgP that CogFunSearch very much does not mean “LLMs can substitute cognitive scientists”. Instead, we hope that by facilitating discovery of interpretable quantitative models we can accelerate scientific research & enable scientists to tackle increasingly ambitious challenges.\n\n## \"Can a two-step task truly capture the richness of cognitive behaviors…?” and “the behaviors studied…are too low-dimensional…”\n\nThe bandit tasks we consider strike a nice balance of not being too simple (so there’s something to learn from them) & not being too complex (so we stand a chance of learning it). Indeed, cognitive scientists have struggled with these bandit tasks for decades, despite their apparent simplicity. \n\n## LLMs have already reached a level of intelligence sufficient to model the intelligent behaviors of other agents\n\nIt is worth clarifying that currently LLMs can’t directly model this data: we are using LLMs to create variants of Python programs which serve as the cognitive models within the framework of FunSearch’s evolutionary algorithm. It is also worth noting that these LLMs are leveraging the knowledge we have built over decades of study via the seed programs provided, & exemplified in the semantically-meaningful variable names used. Reviewer fWgp phrased our goal well: “the authors did not claim \"LLMs can substitute cognitive scientists\" but rather focused on the specific problem of symbolic cognitive modeling, & the experiments have shown convincing results that the discovered models are of high quality.”\n\n## “the findings suggest mainstream symbolic…models may be oversimplifying…in problematic ways.” … “If AI can so readily outperform…do these models…meaningfully advance our understanding…?”\n\nWe partly agree, since RNNs (believed to achieve ceiling predictability) easily outperform previous cognitive models. Our approach largely closes this gap and arguably addresses this shortcoming. We feel that the issue is not with constructing cognitive models themselves, but rather with how we were finding them. Our results demonstrate both that better models are possible, & a novel mechanism for searching for them.\n\n## “are models…exploiting shortcuts in datasets?”\n\nAll our models were validated on held-out subjects that were never seen during training.", "claims_and_evidence": "I think most of the claims in this work are modestly supported:\n* The best-discovered programs demonstrate a strong fit to the data, effectively capturing underlying behavioral patterns.\n* The discovered programs reveal novel strategies, showcasing unique approaches to discover models of behaviors.\n* The discovered programs can readily be interpreted as hypotheses about human and animal cognition, instantiating interpretable symbolic learning and decision-making algorithms.", "methods_and_evaluation_criteria": "The bandit datasets are taken from well-established work, which makes sense.", "theoretical_claims": "There is no theoretic claim in this work.", "experimental_designs_or_analyses": "I found the baseline design in this work problematic. The authors use an RNN trained on data from all subjects collectively, rather than training separate RNNs for individual subjects, noting that the collective model generalizes poorly to held-out data. A more informative comparison might involve evaluating the best LLM-searched program against RNNs overfitted to individual subjects, as this could better highlight the strengths and limitations of each approach in capturing personalized behavioral patterns. They should also consider other variants of baselines, e.g., fine-tuning a small LM agent.", "relation_to_broader_scientific_literature": "Computational process models have been widely utilized in computational cognitive neuroscience to study both human and animal behavior. PhD trainees in this field typically spend four to five years developing the expertise needed to construct effective models that accurately capture behavioral patterns.\n\nThe CogFunSearch workflow presented in this paper is particularly impressive, as it generates computational process model hypotheses that outperform those designed by human researchers. However, its success raises fundamental concerns about the approach itself: (1) To what extent does building computational process models contribute to our scientific understanding of cognitive behavior? (2) How valuable is it, from a scientific training perspective, for students to dedicate years to mastering a skill that large language models can now acquire with ease?\n\nWhile CogFunSearch serves as a highly specialized science tool, it remains relatively narrow in scope compared to more general-purpose AI agents. Its strong performance suggests that computational process modeling may be a fundamentally \"low-dimensional\" task. This prompts a critical question: Can a two-step task truly capture the richness of cognitive behaviors in human and animal intelligence, especially if an LLM can model such data effortlessly?\n\nThe success of CogFunSearch implies two possible interpretations: (1) LLMs have already reached a level of intelligence sufficient to model the intelligent behaviors of other agents, or (2) the behaviors studied—such as those in the two-step task—are too low-dimensional, raising concerns about their suitability as paradigms for investigating intelligence.\n\nI knew this paper follows an established computational neuroscience approach, and given its AI focus, I lean slightly toward acceptance. **However, the findings strongly suggest that mainstream symbolic computational neuroscience models may be oversimplifying human and animal intelligence in problematic ways.** The fact that CogFunSearch can generate superior models so easily raises fundamental concerns about whether these models genuinely capture the complexity of cognitive behavior or merely reflect an artificial, low-dimensional abstraction. If AI can so readily outperform human-designed models, it calls into question whether these models have been meaningfully advancing our understanding of intelligence—or if they have simply been reinforcing oversimplified frameworks.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Why not compare the best LLM-searched program against RNNs overfitted to individual subjects?\n\n2. How can the authors ensure that the models are not merely exploiting shortcuts in the dataset? Or are symbolic cognitive models inherently just sophisticated shortcuts anyways?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8h6henrqz8", "reviewer": "Reviewer_fWgP", "summary": "This paper proposes to extend FunSearch (Romera-Paredes et al., 2024) to symbolic cognitive modeling, namely CogFunSearch, an LLM-based evolutionary program synthesis framework.\nExperimental results have strongly supported the value of CogFunSearch in discovering high-quality symbolic programs on human, rat, and fruitfly bandit datasets, outperforming human-designed symbolic models proposed in very recent literature.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful review of our submission, their useful comments, and are glad they found it “exceptionally well-written and (surprisingly) easy to follow”, and are thrilled (and grateful) they “enjoyed reviewing the paper and learnt a lot”.\nBelow, we address the main issues raised.\n\n## W1. “The authors only evaluate Gemini 1.5 Flash… the paper could be made stronger if more LLMs are tested.”\n\nOur main intent with this work was to demonstrate that the use of LLMs with CogFunSearch was capable of producing state of the art results. An important consideration in methods that use LLMs for the mutation function in an evolutionary algorithm is the tradeoff between model quality (relating to expected performance improvement per sample) and model throughput (relating to the number of samples that can practically be drawn). We chose Gemini 1.5 Flash specifically because it was designed to provide a reasonable tradeoff between these concerns, rather than providing the absolute best quality samples. Although we did not explore the use of other LLMs systematically, we agree this would be a natural idea for pushing the performance of this approach even further.\n\n## W2. “For RNN results it's unclear (1) whether the authors trained the RNNs themselves, (2) what architecture they used, (3) what hyperparameters they tuned.”\n\nWe trained our own RNNs, using  a GRU model (Chung et al., 2014) with a single hidden layer, and trained with early stopping. We performed a sweep over a set of possible hidden units (`[1, 2, 4, 8, 16, 32, 64, 128]`) and picked the value that gave the best performance. All the variants were trained with the Adam optimizer [(Kingma & Ba, 2015)](https://arxiv.org/abs/1412.6980) with a learning rate of `1e-3`.\n\nThank you for raising this point, as we inadvertently missed including these details, and will be including them in the revised version of the paper.\n\n## W3. “recommend moving related work discussion to the main text.”\n\nWe agree, and will be following your suggestion to move the related work (or the most important parts of it) to the main paper.\n\n## C2. “connecting to the work that uses LLMs to propose hypotheses in other CS-related areas would be helpful to provide a comprehensive background to the general ICML audience”\n\nThank you for this suggestion and for the provided references. We agree it would be valuable to connect to this body of related work, and will be including it in our revised related work section.\n\n## C3 and C4\n\nWe will address these issues, thank you for pointing them out.\n\n## Q1. “Is Figure 2B necessarily a Bayesian program synthesis pipeline? It seems that the \"prior\" $\\theta$ is not necessarily a distribution, but a point estimate. Is there any obstacle to generalize it to a distribution?”\n\nWe would like to clarify that $\\theta$ is not a prior, but rather the parameters that are fit in the internal optimization process (Fig. 2). However, we do essentially maintain a \"distribution\" over programs, via our program database (Fig 2A). The prior distribution over these programs is governed by the score $\\Omega$ of each program, which is proportional to the likelihood of it being sampled for evolution. When a new program is generated, this \"prior\" distribution is updated as a new program (with its score $\\Omega$) is added to the database.\n\n## Q2. “L155 (right): Why is $c \\in 0, \\ldots ,n$ instead of $c \\in 1, \\ldots ,n$? Is it allowed to have a \"null\" choice?”\n\nThank you for pointing this out, it should be $c \\in 1, \\ldots ,n$, as the reviewer suspected. We will correct this in the revised version.\n\n## Q3. “If I understood correctly, CogFunSearch is a general framework and should work for general data science tasks that involve symbolic modeling. Are there anything specific things about cognitive modeling that make CogFunSearch particularly suitable for it?”\n\nOur bilevel optimization arose from a common modeling choice in cognitive science, which is that programs capture across-subject patterns and parameters capture individual variations. For example, a cognitive modeling setup might assume or hypothesize that all subjects are executing the same RL algorithms, but each subject has a different learning rate and exploration parameter. In our case, CogFunSearch is set up to fit unique per-subject parameters in the inner loop, and programs are evolved based on the average across-subject score in the outer loop. \n\nHowever, as the reviewer suggests, this can be applicable to other situations with similar setups. The core of our approach is using FunSearch to discover parameterized programs that can be fit to a dataset, which is applicable to other areas of data-driven scientific discovery.", "claims_and_evidence": "In what follows, I will use S, W, C, and Q to denote Strengths, Weaknesses, Comments, and Questions, respectively, for ease of reference in the discussion.\n\nThe authors propose CogFunSearch, an LLM-based symbolic cognitive modeling framework, and demonstrate the effectiveness of the proposed framework on three bandit datasets of human, rat, and fruitfly behaviors. Results have clearly supported that the proposed framework can discover high-quality symbolic cognitive models that outperform human-designed symbolic models.\n\nS1. The paper is exceptionally well-written and (surprisingly) easy to follow. Although I am not an expert in the specific cognitive modeling task addressed in this paper, the authors have done an excellent job in describing the data and formulating the problem, so I can understand without difficulty.\n\nS2. The authors have made a proper yet exciting claim---the authors did not claim \"LLMs can substitute cognitive scientists\" but rather focused on the specific problem of symbolic cognitive modeling, and the experiments have shown convincing results that the discovered models are of high quality.", "methods_and_evaluation_criteria": "The method used in this paper is basically an extension of FunSearch (Romera-Paredes et al., 2024) to symbolic cognitive modeling.\n- The main evaluation criterion to compare models is the likelihood on the held-out data, which is a common metric in machine learning tasks with similar settings and purposes.\n- Additionally, the authors consider the quality of synthesized data, in terms of reproducing scientifically important features, as a secondary metric to illustrate model quality---qualitatively, the discovered programs have shown a sufficiently similar pattern to collected animal data.\n- Finally, the authors have examined the strong robustness of their results with respect to random seeds.\n\nS3. I endorse that the above evaluation criteria form a comprehensive and convincing evaluation protocol for the proposed framework.", "theoretical_claims": "This paper does not contain much theoretical claim, but rather proposes a practical framework for symbolic cognitive modeling using LLMs. However, there could be some interesting theoretical implications of the proposed framework around the interpretability of LLMs and the benefits of continuous vs. discrete representations in cognitive modeling.", "experimental_designs_or_analyses": "As mentioned in the Methods and Evaluation Criteria section, the main experiments are around comparing model likelihood on held-out data, a common metric in similar machine-learning tasks. The authors have also compared RNNs and LLM-based program synthesis, in terms of the gap between each of them and the human-designed baseline models.\n\nW1. The authors only evaluates Gemini 1.5 Flash---which is a considerably strong model but not among the best ones according to my machine-learning knowledge (in my area(s), Gemini 1.5 Pro, OpenAI o1, Anthropic Claude models, and DeepSeek models are usually considered the best ones at the current stage)---as the backend LLM for program synthesis. From a machine-learning perspective, the paper could be made stronger if more LLMs are tested.\n\nW2. For RNN results (Figure 5), it's unclear (1) whether the authors trained the RNNs themselves or they used others' trained RNN models, (2) what the specific RNN architecture they used, (3) what hyperparameters they tuned to optimize the RNNs. (3) is particularly important as given the data size, both overfitting and underfitting are possible, and the results could be very sensitive to hyperparameters. I would recommend the authors pay special attention to this part.\n\nC1. Related to above, it would also be good to cite the original work of RNN (Elman, 1990), and, if you used a specific RNN architecture (such as LSTM), please also cite the corresponding paper(s).\n\nS4. This paper has a strong discussion on the qualitative results of the discovered symbolic models (Section 4), which highlight the importance of the proposed framework in scientific discovery.", "relation_to_broader_scientific_literature": "W3. It is notable that the authors included the related work discussion in supplementary material. I would strongly recommend moving it (or at least the most important ones) to the main text, as it is crucial for grounding the proposed method to existing literature.", "essential_references_not_discussed": "C2. This is not a mandatory request, but connecting to the work that uses LLMs to propose hypotheses in other CS-related areas would be helpful to provide a comprehensive background to the general ICML audience. For example, you may wish to check out https://arxiv.org/abs/2408.05086 and https://aclanthology.org/2024.nlp4science-1.10.pdf.", "other_strengths_and_weaknesses": "N/A. I've discussed all strengths and weaknesses I found in the above sections.", "other_comments_or_suggestions": "C3. The bibliography needs some cleaning.\n- There is an incomplete \"Palminteri. 2015.\" entry in the reference.\n\nC4. Some presentation suggestions and typos:\n- You may wish to consider making the average lines in Figures 1 and 5 thinner, or consider using a box plot for it. I can only realize what the \"blue rectangles\" in Figure 1 left & right are by generalizing from the middle.\n\n- L649: \"Fig B.1.\" has pointed to Section B.1. You may wish to fix this.", "questions_for_authors": "Q1. Is Figure 2B necessarily a Bayesian program synthesis pipeline? It seems that the \"prior\" $\\theta$ is not necessarily a distribution, but a point estimate. Is there any obstacle to generalize it to a distribution?\n\nQ2. L155 (right): Why is $c\\in \\{0, \\ldots, n\\}$ instead of $c\\in \\{1, \\ldots, n\\}$? Is it allowed to have a \"null\" choice? If the latter is the case, you may wish to clarify, as it's not a common bandit problem formulation. \n\nQ3. If I understood correctly, CogFunSearch is a general framework and should work for general data science tasks that involve symbolic modeling. Are there anything specific things about cognitive modeling that make CogFunSearch particularly suitable for it?\n\nC5. Finally, I would like to thank the authors for the strong submission. I enjoyed reviewing the paper and learnt a lot.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jMFW3zMUDt", "reviewer": "Reviewer_wzmy", "summary": "This paper presents a new method, CogFunSearch, that automatically discovers symbolic models for a given dataset. Their approach builds on FunSearch [1], an LLM-driven evolutionary algorithm that searches over the program’s structure, by adding an inner level of optimization that fits model parameters to the data. Experiments on behavioral datasets from humans, rats, and fruit flies demonstrates that CogFunSearch discovers novel symbolic models that outperform hand-engineered programs and are human-interpretable.\n\n[1] Romera-Paredes et al. “Mathematical discoveries from program search with large language models” (2023).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful review of our submission, their useful comments, and are glad they found our approach “novel” and yielding “new behavioral insights”. \n\n## “missing reference to [1].”\n\nThank you for pointing out this relevant work! We will cite it and include it in our discussion.\n\n## W1. “requiring programs to be differentiable may hinder the application to other settings.”\n\nThis is true. We considered it a reasonable limitation, given that the most popular models for these datasets are all differentiable. One could in principle use a gradient-free or hybrid  optimizer for the inner loop (e.g. Acerbi & Ma, 2017), and this would expand the space of models that could be considered. However, this would affect the computational cost required for scoring the models.  We will revise the manuscript to be clear about this limitation.\n\n## W2 & Q2a. “how do the relative complexity scores compare to a human expert judgement.” \n\nAutomatic complexity of programs is an open problem with multiple proposed solutions in the literature, which is why we use multiple methods. Halstead Difficulty is a classic and widely-used measure designed to capture the difficulty of constructing or parsing the program. Several studies do exist linking it to behavioral measures and subjective complexity (e.g. Curtis et al., 2006; De Silva et al., 2012; Gao et al., 2025), but to our knowledge it has not been thoroughly evaluated for Python.  LLMs, having been trained on human data and optimized for human tasks, can also be considered a proxy for human judgments, though this approach also lacks thorough validation. For our programs, we find that these scores largely align (as shown in Figs 7, 12, 13). \n\nThoroughly validating these measures ourselves was beyond the scope of this work, but we do include example programs in the supplement so that readers can assess how they compare to their own intuitive notions of complexity.\n\n## W3. “CogFunSearch requires a significant amount of compute”\n\nWhile this is true, it is not clear how to compare compute cost appropriately against alternative approaches. The baselines models considered represent the results of years of human effort, and were therefore clearly vastly more expensive to produce than the programs we generate automatically. We expect that advances in hardware, model efficiency, and alternative evolution techniques (our rejection sampling being a simple example of one) will result in reduced computational costs for applying approaches such as CogFunSearch.\n\n## W4. “symbolic models may contain errors or be overly complex, hindering readability.”\n\nWhile discovered programs showed sufficient interpretability/readability to afford insights about learning strategies, it is true that there is room for improvement. Exploring tools for minimizing extraneous complexity or improving the readability of the discovered programs is a compelling direction for future work.\n\nIt is also worth mentioning that while the best-fitting models were more complex than human-discovered baselines, there were many discovered programs that were less complex but still outperformed the baselines. \n\n## Q1. “Are the set of parameters the same for all the programs? If so, how does one decide which parameters to use?”\n\nAll programs are provided with the same number of free parameters (10), but each program will use these differently (and some may even use fewer than 10).\n\nAs a concrete example: param[0] may be `learning_rate` for one program, and `lapse_rate` in another. Furthermore, each subject in a dataset can have its own value for each parameter, so `learning_rate` for rat 1 may be different than learning rate for rat 13, but the parameter will be used in the same way.\n\n## Q2b. “ optimal tradeoff between complexity and performance?”\n\nThis would depend on what the model will be used for. In some use cases, greater interpretability (possibly at the expense of accuracy) may be more desirable, and in others, the opposite may hold. One of the advantages of our approach is that it produces multiple programs along this efficiency frontier, enabling researchers to identify the program(s) that are useful to them. We expect a major use case for the tool will be to surface useful ideas about how to model the dataset, & that the best approach will be to examine multiple programs at different points along the complexity/performance tradeoff. \n\n## “How data-efficient is CogFunSearch compared to the RNN? For example, would CogFunSearch outperform the RNN if trained on only 10% of the subjects?”\n\nWhile we have not run this experiment directly, our comparison on the human dataset, where RNN struggles with per-subject parameters, is an indication that CogFunSearch is able to find more generalizable solutions that are still accurate.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes, I checked the experiments in the main paper and the evaluation details in Appendix B.", "theoretical_claims": "N/A.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "There exists some prior work that combine LLM-drive search over program structure with an inner parameter fitting step [1, 2]. As such, the main novelty of the paper is not in the method but instead the application to automating the discovery of symbolic cognitive models, which are typically hand-engineered by domain experts. \n\n[1] Ma et al. “LLM and Simulation as Bilevel Optimizers: A New Paradigm to Advance Physical Scientific Discovery” (2024). \n\n[2] Li et al. “Automated Statistical Model Discovery with Language Models” (2024).", "essential_references_not_discussed": "The key contribution in terms of methods is augmenting FunSearch with bilevel optimization over program structure and parameters. I believe that the paper is missing a reference to [1], which proposes a differentiable inner-level optimization with a LLM-driven evolutionary search. \n\n[1] Ma et al. “LLM and Simulation as Bilevel Optimizers: A New Paradigm to Advance Physical Scientific Discovery” (2024).", "other_strengths_and_weaknesses": "Strengths\n1. The application of LLM-driven program synthesis to cognitive modeling is novel.\n2. Experiments demonstrate that the best discovered programs achieve better performance than a hand-engineered baseline.\n3. Analysis of the complexity and readability of the outputted programs demonstrates that the discovered programs yield new behavioral insights.\n\nWeaknesses\n1. As the inner optimization loop fits parameters using gradient descent, it requires the programs to be differentiable, which may hinder the application of the approach to other settings.\n2. I’m a bit unsure how to interpret the complexity scores, as in how do the relative complexity scores compare to a human expert judgement.\n3. CogFunSearch requires a significant amount of compute, although the paper explores rejection sampling as a way to improve the efficiency.\n4. While more interpretable than a black-box network, the symbolic models may contain errors or be overly complex, hindering readability.", "other_comments_or_suggestions": "1. Missing period at the end of Figure 6.\n2. I would move Figure 1 closer to the results.", "questions_for_authors": "1. Are the set of parameters the same for all the programs? If so, how does one decide which parameters to use? \n2. How do the complexity scores align with a human expert’s perception of the complexity? What is the optimal tradeoff between complexity and performance?\n3. How data-efficient is CogFunSearch compared to the RNN? For example, would CogFunSearch outperform the RNN if trained on only 10% of the subjects?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces CogFunSearch, a novel extension of the FunSearch framework that automates the discovery of symbolic cognitive models using large language models (LLMs) combined with evolutionary algorithms and bilevel optimization. The approach is tested on behavioral datasets from humans, rats, and fruit flies, and demonstrates superior predictive performance over hand-crafted models, while preserving a degree of interpretability.

Reviewers consistently praise the clarity of the writing, the thoroughness of the evaluation, and the novelty of applying LLM-based program synthesis to cognitive science. The inclusion of detailed supplementary material and a rigorous experimental protocol further strengthens the work. However, several concerns were raised, including the high computational cost of the method, limited diversity of baseline comparisons (e.g., against better-tuned RNNs or alternative model discovery methods), and the need for a deeper discussion of the interpretability and scientific utility of the discovered models. Some reviewers also noted that the methodological novelty lies more in its application than in algorithmic innovation, and there are open questions about the broader scientific implications of automating cognitive model discovery with LLMs. Nevertheless, the author's rebuttal addresses most of the concerns. Hence, the consensus among reviewers leans clearly toward acceptance and I agree with them.
Decision: Accept (spotlight poster)