CogFormer: Learn All Your Models Once

Jerry M. Huang; Lukas Schumacher; Niek Stevenson; Stefan T. Radev

stat.MLMar 20, 2026

CogFormer: Learn All Your Models Once

Authors: Jerry M. Huang, Lukas Schumacher, Niek Stevenson, Stefan T. Radev

Organizations: Center for Modeling, Simulation, and Imaging in Medicine (CeMSIM), Rensselaer Polytechnic Institute, Troy, New York, United States · Center for Economic Psychology, University of Basel, Basel, Switzerland · Amsterdam Mathematical Psychology Lab, University of Amsterdam , Amsterdam, Netherlands

Abstract

Simulation-based inference (SBI) with neural networks has accelerated and transformed cognitive modeling workflows. SBI enables modelers to fit complex models that were previously difficult or impossible to estimate, while also allowing rapid estimation across large numbers of datasets. However, the utility of SBI for iterating over varying modeling assumptions remains limited: changes to parameterizations, generative functions, priors, and design variables all necessitate model retraining, thereby diminishing the benefits of amortization. To address these issues, we pilot the CogFormer, a meta-amortized framework for cognitive modeling. Our framework trains a transformer-based architecture that remains valid across a combinatorial number of structurally similar models, allowing for changing data types, parameters, design matrices, and sample sizes. We present promising quantitative results across families of decision-making models for binary, multi-alternative, and continuous responses. Our evaluation suggests that CogFormer can accurately estimate parameters across model families with minimal amortization offset, making it a potentially powerful engine that catalyzes cognitive modeling workflows.

Figures & tables

Figure 1: Taxonomy of amortization scopes . In cognitive modeling, one fits concrete model instances

\mathcal{M}

nested within model families

\mathcal{F}

as part of a larger model class

\mathcal{C}

. Here, we train neural estimators that generalizes across all model families within a given model class (in this paper, Sequential Sampling Models).

Figure 2: Meta-simulator architecture. For a single model family ( middle ), a coefficient matrix

{\mathbf{B}}

is constructed under a flexible generalized linear model (GLM) configuration, and configured through a random binary parameter mask

{\mathbf{M}}

. Together with a sampled design matrix

{\mathbf{X}}

and element-wise link functions

g

, this construction induces the trial-level parameter matrix

\Theta

, which is passed to the observation model

\mathcal{M}

to generate observables

{\mathbf{Y}}

. The same generative process applies to every model in the class

\mathcal{C}

Figure 3: CogFormer architecture. CogFormer is an encoder-decoder architecture in which the encoder is a Set Transformer that converts a design matrix

{\mathbf{X}}

and model observables

{\mathbf{Y}}

into summary representations

{\mathbf{S}}

, and the decoder is a stack of mixed-attention layers that produce a learned velocity field

\hat{u}

. The latter parameterizes the learned velocity field

\hat{u}

. Training minimizes a conditional flow-matching objective, yielding a procedure for sampling from the resulting ensemble posterior.

Figure 4: Ensemble parameter recoveries for CogFormer- $\mathcal{F}$ applied to the five design configurations for the DDM benchmark cases overlaid in their corresponding mapping to the shared coefficient matrix. The pixel thumbnails indicate the active estimation targets for each design configuration, with lighter shades of color depicting higher Pearson correlations $r$ . The per-cell correlations are averaged across configurations.

Figure 5: Ensemble calibration ECDFs for CogFormer- $\mathcal{F}$ applied to five design configurations for the DDM benchmark cases overlaid in their corresponding mapping to the coefficient matrix. The pixel thumbnails (bottom row) indicate the active estimation targets for each design configuration in the shared coefficient matrix, with lighter shades depicting lower calibration errors (ECEs).

Figure 6: Ensemble parameter recoveries for CogFormer- $\mathcal{C}$ for 2,400 simulated datasets with 500 trials each over 12 randomly generated design configurations. The pixel thumbnails at the bottom indicate the active estimation targets within each configuration, with lighter shades of color indicate higher Pearson correlation $r$ . Results are overlaid on top of each other, and the per-cell values are averaged across configurations.

Figure 7: Ensemble calibration ECDFs for CogFormer-

\mathcal{C}

for 2,400 simulated datasets with 500 trials each over 12 model configurations. The pixel thumbnails at the bottom indicate the active estimation targets within each configuration, with lighter shades indicating lower calibration errors (ECEs). Results are plotted corresponding to these thumbnails at their corresponding locations within the shared coefficient matrix.

Model	Method	NRMSE ( $\downarrow$ )	Cal. Error ( $\downarrow$ )	Post. Contr. ( $\uparrow$ )
DDM	Baseline	$0.078\pm 0.007$	$0.022\pm 0.002$	$0.902\pm 0.023$
	CogFormer - $\mathcal{F}$	$0.087\pm 0.007$	$0.021\pm 0.002$	$0.881\pm 0.025$
	CogFormer - $\mathcal{C}$	$0.091\pm 0.007$	$0.029\pm 0.003$	$0.862\pm 0.026$
RDM	Baseline	$0.112\pm 0.008$	$0.061\pm 0.010$	$0.803\pm 0.032$
	CogFormer - $\mathcal{F}$	$0.112\pm 0.009$	$0.026\pm 0.004$	$0.786\pm 0.033$
	CogFormer - $\mathcal{C}$	$0.117\pm 0.009$	$0.029\pm 0.003$	$0.750\pm 0.036$

Table 1: Absolute metric values (median

\pm

SEM) aggregated across the five benchmark design configurations and all parameters in each configuration.

Comparison	DDM	RDM	CDM
CogFormer - $\mathcal{F}$ vs. Baseline	$0.693\pm 0.067$	$0.702\pm 0.076$	$0.668\pm 0.068$
CogFormer - $\mathcal{C}$ vs. CogFormer - $\mathcal{F}$	$0.640\pm 0.049$	$0.600\pm 0.032$	$0.635\pm 0.041$

Table 2: Joint C2ST (median

\pm

SEM) aggregated across 10 posteriors for each configuration.

Appendix figures & tables62 assets

Supplementary material from the paper’s appendix.

Appendix

Category	Symbol	Meaning
Amortization scopes	$\mathcal{C}$	Model class (root; e.g. sequential sampling models)
	$\mathcal{F}$	Model family (shared computational structure)
	$\mathcal{M}$	Model instance / observation model (fully specified)
	$\mathcal{D}$	Design configuration over the intrinsic parameters
Generative objects	$\theta_{\mathcal{F}}$	Intrinsic parameters of family $\mathcal{F}$
	${\mathbf{B}}$	Coefficient matrix, ${\mathbf{B}}\in\mathbb{R}^{R\times D_{\mathcal{F}}}$ (inference target)

Appendix

Table 3: Table of notations

\Theta=g\!\left({\mathbf{X}}\,({\mathbf{M}}\odot{\mathbf{B}})\right)\in\mathbb{R}^{N\times D_{\mathcal{F}}}

Appendix

Algorithm 1 Model Family Simulator

Figure 8: Breakdown of the attention blocks in the CogFormer architecture. For set attention, the input

QKV

matrices are identical linear layer outputs of the input data. Its output

H

is the encoder representation, which serves as the input

KV

matrices for pooling along with a learnable seed vector as the query

Q

. The mixed attention later takes the encoder summary z as query, along with the constructed decoder embeddings listed in Figure 3 as the

KV

matrices. The mixed-attention architecture places the cross-attention layer first, followed by the self-attention layer, and further alternates between the two. The layer-wise output is element-wise multiplied by the parameter mask

M

before passing to the next layer. The resulting output

x

serves as the decoder representation that is used to generate a velocity field

v

Hyperparameter	CogFormer- $\mathcal{F}$	CogFormer- $\mathcal{C}$
Layers
Encoder layers	8	8
Decoder layers	8	8
Attention
Heads (encoder & decoder)	8	8
Projection dim $d_{\text{model}}$	256	256

Appendix

Table 4: CogFormer hyperparameters for the family-level models (DDM, RDM, CDM) and the model-class model. QKV dim per head = projection dim

\div

heads.

Configuration	Model	Intercept ( $\mathbf{1}$ )	Slope ( $u_{1}$ )	Slope ( $u_{2}$ )	Interaction ( $u_{1}{\times}u_{2}$ )
Intercept Only	DDM	$\vartheta,a,z,\tau,s_{\vartheta},s_{\tau}$	$-$	$-$	$-$
	RDM	$\vartheta,\vartheta_{\Delta},a,\tau,s_{\vartheta},s_{\tau}$	$-$	$-$	$-$
	CDM	$\vartheta,\vartheta_{\theta},a,\tau,s_{\vartheta},s_{\tau}$	$-$	$-$	$-$
Fixed Variability	DDM	$\vartheta,a,z,\tau$	$-$	$-$	$-$
	RDM	$\vartheta,\vartheta_{\Delta},a,\tau$	$-$	$-$	$-$
	CDM	$\vartheta,\vartheta_{\theta},a,\tau$	$-$	$-$	$-$

Appendix

Table 5: Test design configurations. Each cell lists the parameters assigned to the corresponding effect type. A dash (

-

) indicates no parameters are assigned.

\vartheta_{\Delta}

abbreviates

\vartheta_{\Delta}

(RDM).

Variant	Component Role
No SAB	Self-attention across parameter tokens; allows parameters to coordinate with one another during the denoising pass.
No MAB	Cross-attention from parameter tokens to the encoder summary; the primary pathway through which the decoder reads the observed data.
Baseline	Feature-wise linear modulation (FiLM) of the output projection, conditioned on the current diffusion timestep; scales and shifts the final representation before predicting velocity.
No FE	Learnable Fourier embedding (FE) of the diffusion timestep; projects the scalar $t$ into a high-dimensional frequency space before concatenation.
With FiLM	Feature-wise linear modulation (FiLM) of the output projection, conditioned on the current diffusion timestep; scales and shifts the final representation before predicting velocity.

Appendix

Table 6: Summary of component ablation variants and the architectural role of each ablated component.

Condition	Variant	NRMSE $\downarrow$	Cal. Error $\downarrow$	Contraction $\uparrow$
Fixed ( $n=4$ )	Baseline	0.055	0.021	0.966
	No SAB	0.058	0.039	0.944
	No MAB	0.255	0.017	0.037
	No FE	0.053	0.023	0.965
	With FiLM	0.058	0.017	0.966
Intercept Only ( $n=6$ )	Baseline	0.087	0.022	0.845

Appendix

Table 7: Component ablation metrics across all design conditions, ordered by number of active parameters. NRMSE and Calibration Error: lower is better (

\downarrow

). Posterior Contraction: higher is better (

\uparrow

). Bold indicates best per metric within each condition group. † No MAB collapses Posterior Contraction across all conditions; its Cal. Error optimum is not meaningful, since the prior itself is perfectly calibrated.

Design	Condition	NRMSE $\downarrow$	Cal. Error $\downarrow$	Contraction $\uparrow$
Intercept Only	Without Embedding	0.104 (+0.005)	0.033 (-0.009)	0.794 (-0.033)
Intercept Only	With Embedding	0.099	0.042	0.827
Fixed	Without Embedding	0.068 (+0.008)	0.044 (-0.019)	0.947 (-0.006)
Fixed	With Embedding	0.060	0.063	0.954
Regressed	Without Embedding	0.082 (-0.002)	0.031 (+0.000)	0.884 (-0.002)
Regressed	With Embedding	0.084	0.031	0.886

Appendix

Table 8: Model embedding ablation for the DDM across all five design conditions. Mean metrics across all active parameters.

\Delta

(in parentheses) is the change relative to the with-embedding baseline. NRMSE and Calibration Error: lower is better (

\downarrow

). Posterior Contraction: higher is better (

\uparrow

). Best result per condition per metric is bolded .

	S (2L, 2H)	M (4L, 4H)	L (8L, 8H; baseline)	XL (16L, 16H)
NRMSE $\downarrow$	0.179	0.170	0.124	0.127
Cal. Error $\downarrow$	0.024	0.032	0.022	0.045
Post. Contr. $\uparrow$	0.472	0.566	0.746	0.701

Appendix

Table 9: Mean posterior approximation metrics across all parameters and conditions for each CogFormer-

\mathcal{F}

size variant over the number of attention layers (L) and the number of attention heads (H). NRMSE and Calibration Error: lower is better (

\downarrow

). Posterior Contraction: higher is better (

\uparrow

). Best result per column is bolded .

	Coupling Flow	Diffusion Model	Flow Matching	Consistency Model
NRMSE $\downarrow$	0.128	0.131	0.117	0.129
Cal. Error $\downarrow$	0.020	0.025	0.017	0.067
Post. Contr. $\uparrow$	0.793	0.746	0.820	0.706

Appendix

Table 10: Mean posterior approximation metrics across all DDM parameters for each inference network. NRMSE and Calibration Error: lower is better (

\downarrow

). Posterior Contraction: higher is better (

\uparrow

). Best result per column is bolded .

Figure 9: Joint C2ST aggregated across 10 posteriors for each benchmark design configuration per increasing simulation budget for CogFormer-

\mathcal{F}

( top ) and CogFormer-

\mathcal{C}

( bottom ). The training epochs are normalized to between 0 and 1, with 0.2 indicating Epoch 1000 and 1 indicating Epoch 5000.

Panel of Figure 10 — Figure 10: Random ensemble parameter recovery ( top ) and calibration ECDF ( bottom ) for meta-amortized model family for DDM. The pixel legends at the bottom indicates the parameter mask for each benchmark design configurations.

Panel of Figure 11 — Figure 11: Ensemble parameter recoveries for CogFormer- $\mathcal{F}$ applied to the five design configurations for the RDM benchmark cases overlayed in their corresponding mapping to the shared coefficient matrix. The pixel thumbnails indicates the active estimation targets for each design configuration, with lighter shades of color depicting higher correlations $r$ . The per-cell $r$ -s are averaged across cases.

Panel of Figure 12 — Figure 12: Random ensemble parameter recovery ( top ) and calibration ECDF ( bottom ) for meta-amortized model family for RDM. The pixel legends at the bottom indicates the parameter mask for each benchmark design configurations.

Panel of Figure 13 — Figure 13: Ensemble parameter recoveries for CogFormer- $\mathcal{F}$ applied to the five design configurations for the CDM benchmark cases overlayed in their corresponding mapping to the shared coefficient matrix. The pixel thumbnails indicates the active estimation targets for each design configuration, with lighter shades of color depicting higher correlations $r$ . The per-cell $r$ -s are averaged across cases.

Panel of Figure 14 — Figure 14: Random ensemble parameter recovery ( top ) and calibration ECDF ( bottom ) for meta-amortized model family for CDM. The pixel legends at the bottom indicates the parameter mask for each benchmark design configurations.

Frameworks	Figures	Color Scheme
Baseline	15-19 (DDM) / 30-34 (RDM) / 45-49 (CDM)	Green
CogFormer- $\mathcal{F}$	20-24 (RDM) / 35-39 (CDM) / 50-54 (DDM)	Red
CogFormer- $\mathcal{C}$	25-29 (RDM) / 40-44 (CDM) / 55-59 (DDM)	Blue

Appendix

Table 11: Corresponding frameworks, figures, and colors for the additional results for benchmarking.

Figure 15: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for DDM baseline (Case intercept_only ).

Figure 16: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for DDM baseline (Case fixed ).

Figure 17: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for DDM baseline (Case regressed ).

Figure 18: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for DDM baseline (Case fixed_regressed ).

Figure 19: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for DDM baseline (Case interaction ).

Figure 20: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for DDM model family (Case intercept_only ).

Figure 21: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for DDM model family (Case fixed ).

Figure 22: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for DDM model family (Case regressed ).

Figure 23: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for DDM model family (Case fixed_regressed ).

Figure 24: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for DDM model family (Case interaction ).

Figure 25: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for DDM model class (Case intercept_only ).

Figure 26: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for DDM model class (Case fixed ).

Figure 27: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for DDM model class (Case regressed ).

Figure 28: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for DDM model class (Case fixed_regressed ).

Figure 29: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for DDM model class (Case interaction ).

Figure 30: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for RDM baseline (Case intercept_only ).

Figure 31: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for RDM baseline (Case fixed ).

Figure 32: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for RDM baseline (Case regressed ).

Figure 33: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for RDM baseline (Case fixed_regressed ).

Figure 34: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for RDM baseline (Case interaction ).

Figure 35: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for RDM model family (Case intercept_only ).

Figure 36: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for RDM model family (Case fixed ).

Figure 37: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for RDM model family (Case regressed ).

Figure 38: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for RDM model family (Case fixed_regressed ).

Figure 39: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for RDM model family (Case interaction ).

Figure 40: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for RDM model class (Case intercept_only ).

Figure 41: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for RDM model class (Case fixed ).

Figure 42: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for RDM model class (Case regressed ).

Figure 43: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for RDM model class (Case fixed_regressed ).

Figure 44: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for RDM model class (Case interaction ).

Figure 45: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for CDM baseline (Case intercept_only ).

Figure 46: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for CDM baseline (Case fixed ).

Figure 47: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for CDM baseline (Case regressed ).

Figure 48: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for CDM baseline (Case fixed_regressed ).

Figure 49: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for CDM baseline (Case interaction ).

Figure 50: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for CDM model family (Case intercept_only ).

Figure 51: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for CDM model family (Case fixed ).

Figure 52: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for CDM model family (Case regressed ).

Figure 53: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for CDM model family (Case fixed_regressed ).

Figure 54: Parameter recovery ( top ), calibration ECDF ( middle ), and validation metrics (NRMSE, calibration error, and posterior contraction) for CDM model family (Case interaction ).

Figure 55: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for CDM model class (Case intercept_only ).

Figure 56: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for CDM model class (Case fixed ).

Figure 57: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for CDM model class (Case regressed ).

Figure 58: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for CDM model class (Case fixed_regressed ).

Figure 59: Parameter recovery ( top ), calibration ECDF ( middle ), and parameter-wise metrics (NRMSE, calibration error, and posterior contraction) for CDM model class (Case interaction ).

Amortized Inference

Amortized inference uses trained models to turn repeated probabilistic inference problems into fast predictions or posterior samples.

See all 85 papers

Hypothesis-guided discovery of cognitive algorithms via program refinement
Oct 1, 2026Huiwen Alex Yang, Mark K. Ho, Bill D. ThompsonComputational Cognitive Modeling Automated Algorithm Discovery
Department of Psychology University of California, Berkeley · Department of Psychology New York University
Think-Aloud Reshapes Automated Cognitive Model Discovery Beyond Behavior
May 6, 2026Hanbo Xie, Akshay K. Jagadish, Lan Pan +1Computational Cognitive Modeling Human Behavior Modeling
School of Psychology, Georgia Institute of Technology · Princeton AI Lab, Princeton University · Center of Excellence for Computational Cognition, Georgia Institute of Techonology
Adapting Neural Posterior Estimators for Simulation-Based Inference in Agent-Based Models
Jan 12, 2026Junyao Zhang, Jinglai Li, Junqi TangSimulation-Based Inference Neural Posterior Estimation
School of Mathematics, University of Birmingham

Oct 1, 2026J/K move · Enter open · S save

Hypothesis-guided discovery of cognitive algorithms via program refinement

Huiwen Alex Yang, Mark K. Ho, Bill D. Thompson

Department of Psychology University of California, Berkeley · Department of Psychology New York University

Computational Cognitive Modeling Automated Algorithm Discovery Algorithmic Reasoning

Open paper AskPDF arXiv

Developing cognitive models of algorithmic reasoning from behavioral data is a central problem in cognitive science that challenges current methods. Traditional approaches to cognitive modeling are interpretable and benefit from human expertise, but lack flexibility and scalability. Emerging techniques using large language models (LLMs) for de novo generation of cognitive models are scalable and flexible, but lack a role for human expertise and have mostly been applied to simpler tasks than algorithm recovery. We propose a hybrid system that treats discovery of cognitive algorithms as a program refinement problem. Human-created cognitive models are expressed as probabilistic programs and provided to a system of LLM agents with a mandate to: identify mismatches between model and behavior; propose code-level modifications within researcher-specified constraints; and verify structural fidelity. Revisions propagate to a probabilistic inference module that performs inference for latent variables and data likelihood computations. We evaluate the pipeline on human behavior in a problem-solving paradigm that exposes a variety of cognitive algorithms. Revised models consistently improve model fit relative to ancestral models and reveal a small set of recurring innovations that capture meaningful behavioral variability in this task.

Abstract

Amortized Inference

Explore similar work

Hypothesis-guided discovery of cognitive algorithms via program refinement

Think-Aloud Reshapes Automated Cognitive Model Discovery Beyond Behavior

Adapting Neural Posterior Estimators for Simulation-Based Inference in Agent-Based Models