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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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 M\mathcal{M} nested within model families F\mathcal{F} as part of a larger model class C\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 B{\mathbf{B}} is constructed under a flexible generalized linear model (GLM) configuration, and configured through a random binary parameter mask M{\mathbf{M}} . Together with a sampled design matrix X{\mathbf{X}} and element-wise link functions gg , this construction induces the trial-level parameter matrix Θ\Theta , which is passed to the observation model M\mathcal{M} to generate observables Y{\mathbf{Y}} . The same generative process applies to every model in the class C\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 X{\mathbf{X}} and model observables Y{\mathbf{Y}} into summary representations S{\mathbf{S}} , and the decoder is a stack of mixed-attention layers that produce a learned velocity field u^\hat{u} . The latter parameterizes the learned velocity field u^\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- F\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 rr . The per-cell correlations are averaged across configurations.
Figure 5: Ensemble calibration ECDFs for CogFormer- F\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- C\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 rr . Results are overlaid on top of each other, and the per-cell values are averaged across configurations.
Figure 7: Ensemble calibration ECDFs for CogFormer- C\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.
ModelMethodNRMSE ( ↓\downarrow )Cal. Error ( ↓\downarrow )Post. Contr. ( ↑\uparrow )
DDMBaseline0.078±0.0070.078\pm 0.0070.022±0.0020.022\pm 0.0020.902±0.0230.902\pm 0.023
 CogFormer - F\mathcal{F}0.087±0.0070.087\pm 0.0070.021±0.0020.021\pm 0.0020.881±0.0250.881\pm 0.025
 CogFormer - C\mathcal{C}0.091±0.0070.091\pm 0.0070.029±0.0030.029\pm 0.0030.862±0.0260.862\pm 0.026
RDMBaseline0.112±0.0080.112\pm 0.0080.061±0.0100.061\pm 0.0100.803±0.0320.803\pm 0.032
 CogFormer - F\mathcal{F}0.112±0.0090.112\pm 0.0090.026±0.0040.026\pm 0.0040.786±0.0330.786\pm 0.033
 CogFormer - C\mathcal{C}0.117±0.0090.117\pm 0.0090.029±0.0030.029\pm 0.0030.750±0.0360.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.
ComparisonDDMRDMCDM
CogFormer - F\mathcal{F} vs. Baseline0.693±0.0670.693\pm 0.0670.702±0.0760.702\pm 0.0760.668±0.0680.668\pm 0.068
CogFormer - C\mathcal{C} vs. CogFormer - F\mathcal{F}0.640±0.0490.640\pm 0.0490.600±0.0320.600\pm 0.0320.635±0.0410.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

CategorySymbolMeaning
Amortization scopesC\mathcal{C}Model class (root; e.g. sequential sampling models)
 F\mathcal{F}Model family (shared computational structure)
 M\mathcal{M}Model instance / observation model (fully specified)
 D\mathcal{D}Design configuration over the intrinsic parameters
Generative objectsθF\theta_{\mathcal{F}}Intrinsic parameters of family F\mathcal{F}
 B{\mathbf{B}}Coefficient matrix, B∈RR×DF{\mathbf{B}}\in\mathbb{R}^{R\times D_{\mathcal{F}}} (inference target)
Appendix
Table 3: Table of notations
 Θ=g ⁣(X (M⊙B))∈RN×DF\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 QKVQKV matrices are identical linear layer outputs of the input data. Its output HH is the encoder representation, which serves as the input KVKV matrices for pooling along with a learnable seed vector as the query QQ . The mixed attention later takes the encoder summary z as query, along with the constructed decoder embeddings listed in Figure 3 as the KVKV 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 MM before passing to the next layer. The resulting output xx serves as the decoder representation that is used to generate a velocity field vv .
HyperparameterCogFormer- F\mathcal{F}CogFormer- C\mathcal{C}
Layers  
Encoder layers88
Decoder layers88
Attention  
Heads (encoder & decoder)88
Projection dim dmodeld_{\text{model}}256256
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.
ConfigurationModelIntercept ( 1\mathbf{1} )Slope ( u1u_{1} )Slope ( u2u_{2} )Interaction ( u1×u2u_{1}{\times}u_{2} )
Intercept OnlyDDMϑ,a,z,τ,sϑ,sτ\vartheta,a,z,\tau,s_{\vartheta},s_{\tau}−-−-−-
 RDMϑ,ϑΔ,a,τ,sϑ,sτ\vartheta,\vartheta_{\Delta},a,\tau,s_{\vartheta},s_{\tau}−-−-−-
 CDMϑ,ϑθ,a,τ,sϑ,sτ\vartheta,\vartheta_{\theta},a,\tau,s_{\vartheta},s_{\tau}−-−-−-
Fixed VariabilityDDMϑ,a,z,τ\vartheta,a,z,\tau−-−-−-
 RDMϑ,ϑΔ,a,τ\vartheta,\vartheta_{\Delta},a,\tau−-−-−-
 CDMϑ,ϑθ,a,τ\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).
VariantComponent Role
No SABSelf-attention across parameter tokens; allows parameters to coordinate with one another during the denoising pass.
No MABCross-attention from parameter tokens to the encoder summary; the primary pathway through which the decoder reads the observed data.
BaselineFeature-wise linear modulation (FiLM) of the output projection, conditioned on the current diffusion timestep; scales and shifts the final representation before predicting velocity.
No FELearnable Fourier embedding (FE) of the diffusion timestep; projects the scalar tt into a high-dimensional frequency space before concatenation.
With FiLMFeature-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.
ConditionVariantNRMSE ↓\downarrowCal. Error ↓\downarrowContraction ↑\uparrow
Fixed ( n=4n=4 )Baseline0.0550.0210.966
 No SAB0.0580.0390.944
 No MAB0.2550.0170.037
 No FE0.0530.0230.965
 With FiLM0.0580.0170.966
Intercept Only ( n=6n=6 )Baseline0.0870.0220.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.
DesignConditionNRMSE ↓\downarrowCal. Error ↓\downarrowContraction ↑\uparrow
Intercept OnlyWithout Embedding0.104 (+0.005)0.033 (-0.009)0.794 (-0.033)
With Embedding0.0990.0420.827
FixedWithout Embedding0.068 (+0.008)0.044 (-0.019)0.947 (-0.006)
With Embedding0.0600.0630.954
RegressedWithout Embedding0.082 (-0.002)0.031 (+0.000)0.884 (-0.002)
With Embedding0.0840.0310.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 ↓\downarrow0.1790.1700.1240.127
Cal. Error ↓\downarrow0.0240.0320.0220.045
Post. Contr. ↑\uparrow0.4720.5660.7460.701
Appendix
Table 9: Mean posterior approximation metrics across all parameters and conditions for each CogFormer- F\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 FlowDiffusion ModelFlow MatchingConsistency Model
NRMSE ↓\downarrow0.1280.1310.1170.129
Cal. Error ↓\downarrow0.0200.0250.0170.067
Post. Contr. ↑\uparrow0.7930.7460.8200.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- F\mathcal{F} ( top ) and CogFormer- C\mathcal{C} ( bottom ). The training epochs are normalized to between 0 and 1, with 0.2 indicating Epoch 1000 and 1 indicating Epoch 5000.
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.
Figure 11: Ensemble parameter recoveries for CogFormer- F\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 rr . The per-cell rr -s are averaged across cases.
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.
Figure 13: Ensemble parameter recoveries for CogFormer- F\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 rr . The per-cell rr -s are averaged across cases.
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.
FrameworksFiguresColor Scheme
Baseline15-19 (DDM) / 30-34 (RDM) / 45-49 (CDM)Green
CogFormer- F\mathcal{F}20-24 (RDM) / 35-39 (CDM) / 50-54 (DDM)Red
CogFormer- C\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 ).
More in this topic

Amortized Inference

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

See all 85 papers

Explore similar work

CardsList
Oct 1, 2026

Hypothesis-guided discovery of cognitive algorithms via program refinement

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.

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

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

Computational Cognitive ModelingAutomated Algorithm DiscoveryAlgorithmic Reasoning
AskPDF
May 6, 2026

Think-Aloud Reshapes Automated Cognitive Model Discovery Beyond Behavior

Computational cognitive models discovered using large language models have so far relied solely on behavioral data. However, it is well-known that models produced from the behavioral trajectory alone are typically under-determined. In this work, we explore the use of Think Aloud traces as an additional form of data constraint during automated model discovery. When applied to the domain of risky decision-making, we find that the models discovered with think-aloud achieve significantly improved predictive performance on held-out data. Additionally, we find that the discovered models belong to different structural classes than those discovered from behavior alone for the majority of participants (69.4%), specifically, it shifts from Explicit comparator towards Integrated utility. These results suggest that process-level language data not only improve model fit, but also systematically reshape the structure of the discovered cognitive models, enabling the identification of mechanisms that are not recoverable from behavior alone.

Hanbo Xie, Akshay K. Jagadish, Lan Pan +1

School of Psychology, Georgia Institute of Technology · Princeton AI Lab, Princeton University · Center of Excellence for Computational Cognition, Georgia Institute of Techonology

Computational Cognitive ModelingHuman Behavior Modeling
AskPDF
Jan 12, 2026

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

Agent-based models (ABMs) are widely used in economics and social sciences. They represent heterogeneous personal behaviors and their interactions by using simulations, whereas their likelihoods are often intractable for parameter estimations. Therefore, most traditional statistical inference techniques are hindered. Driven by recent advances in probabilistic machine learning, recent works have applied neural networks to achieve likelihood-free inference in large-scale ABMs. Specifically, Neural Posterior Estimation using normalizing flows provide amortized Bayesian inference from simulations, yet a pre-trained estimator may become inaccurate when the target task changes. In this work, we investigate the practicality of adapting a pre-trained conditional normalizing flow through a unified sequential neural posterior estimation test-time training (SNPE-TTT) framework. The framework includes four adaptation strategies: full-parameter adaptation, low-rank adaptation, and two update spaces identified from target-task gradients. We evaluate all four methods across economic simulation models with both tractable and intractable likelihoods. These benchmarks demonstrate that the sequential neural posterior estimation test-time training framework can adapt pre-trained posterior estimators to related inference tasks using a limited budget of target simulations. Posterior recovery and computational benefits depend on the task and adaptation strategy, with lower target-stage cost in some settings but no consistently superior posterior accuracy.

Junyao Zhang, Jinglai Li, Junqi Tang

School of Mathematics, University of Birmingham

Simulation-Based InferenceNeural Posterior EstimationTest-Time TrainingAgent-Based Modeling
AskPDF