cs.LGSep 26, 2026
SaveRecovery-Directed Symbolic Distillation of Neural Likelihoods
Organizations: Department of Psychology University of California, Los Angeles Los Angeles, CA 90095, USA · Department of Civil and Environmental Engineering National University of Singapore Singapore
Abstract
Amortized neural likelihoods enable computationally expensive inference for models with analytically intractable or unspecified likelihoods, but their black-box nature limits interpretability. We introduce a symbolic distillation pipeline that converts trained neural likelihoods into explicit, interpretable expressions optimized for efficient parameter estimation. Our approach uses a recovery-directed objective to guide symbolic regression toward expressions that preserve parameter-recovery accuracy rather than merely approximating the likelihood function. Candidate expressions are evaluated on held-out datasets and selected using a criterion that jointly accounts for expression complexity, parameter-recovery performance, and distributional distance from the learned likelihood. We evaluate the pipeline on the diffusion decision model, a classical cognitive model, whose analytically tractable likelihood provides ground truth for controlled evaluation. The proposed recovery-directed objective improves parameter recovery over standard symbolic-regression objectives. The resulting symbolic likelihoods enable over 100 times faster parameter evaluation than both neural likelihoods and, when available, the exact likelihood, while maintaining a manageable loss in precision. We further demonstrate these computational benefits in Bayesian hierarchical inference on empirical data. Our pipeline provides a lightweight interface for integrating symbolic distillation with existing neural-likelihood estimation methods and can be adapted to a range of simulation-based inference settings.
Figures & tables
Figure 1: Pipeline for recovery-directed symbolic distillation of the likelihood of a generative model. Tr, V and T denote the training, validation and test data.
Figure 2: Parameter estimates on the test dataset. Each point is one subject, and columns are the four DDM parameters ( , , , ). Each panel title gives the Pearson across subjects. (a) True parameter value ( -axis) against the posterior mean estimated with NUTS ( -axis) under four likelihoods, one per row. From top to bottom, these are the analytical likelihood, the neural likelihood, an expression found under Equation 4 , and an expression found under Equation 5 (third in Table 1 , fit in a separate NUTS run). (b) Pairwise comparison of posterior-mode estimates under the analytical likelihood, the neural likelihood, and the expression in the bottom row of a .
| MSE | held-out | KS | beats null | |||||||
| 1 | ✓ | |||||||||
| 2 | ✓ | |||||||||
| 3 | ✓ | |||||||||
| 4 | – | |||||||||
| 5 | ✓ |
Table 1: The top five expressions’ performance on the test bank. is the recovery weight in Equation 5 , is the expression complexity, MSE is its squared error against the neural log-likelihood on held-out rows (pool median ), and the , , and columns give the recovery correlation for each parameter. Held-out gives the weakest of the four correlations on further subjects never used for ranking ( Section H.2 ). KS is the median Kolmogorov–Smirnov distance between observed and posterior-predicted RTs (signed by choice), where lower is better. For reference, it is for the MSE expression in Figure 2 and about for data simulated from the true parameters. A check mark means the expression’s KS is lower than that of an unbiased estimator with the same recovery ( Section H.2 ).
Figure 3: Effect of the recovery weight . (a) Recovery of each parameter against expression size. Each dot is one expression, colored by , and each line traces, for one value of , the best recovery reached at that size or smaller. The dotted line marks the analytical likelihood, and the dashed line the criterion of . (b) Proportion of expressions that recovered all four parameters. Each translucent dot is one independently seeded search, and the solid line gives the pooled rate. The dashed line gives the faster posterior-mode (MAP) screen. (c) Proportion of expressions that recovered each parameter.
| Expression 2 | Navarro–Fuss | LAN | neural | ||||
| trials | value | gradient | value | gradient | value | gradient | likelihood |
| 0.010 | 0.015 | 0.27 | 0.59 | 0.68 | 1.3 | 1144 | |
| 0.057 | 0.28 | 29 | 76 | 13 | 35 | 38,080 | |
| 0.51 | 2.7 | 315 | 905 | 161 | 279 | 383,500 | |
Table 2: Time to evaluate one summed log-likelihood (ms; median of repeated calls after warm-up). The expression example in Equation 7 was compiled through PyTensor, the analytical Navarro–Fuss series and the LAN are HSSM’s compiled implementations (PyTensor and JAX), and the neural likelihood was timed at the ODE steps used in estimation. Gradient is the value together with its gradient, which is what NUTS evaluates at each step.
Figure 4: Evaluation cost across dataset sizes , on log–log axes. Green curves are Expression 2, black curves the analytical series and blue curves HSSM’s LAN. Solid curves are NumPy, dashed curves are compiled, and dotted curves include the gradient. The remaining curves are neural likelihoods trained here, including a coupling flow that needs a single network pass per evaluation. Table 2 gives the values at three sizes.
Figure 5: Hierarchical DDM fit to Cavanagh’s data. (a) Participant-level posterior means in the full model, in which every parameter varies by participant. Each translucent line is one of the converged expressions. Black points and bars show the analytical likelihood’s posterior mean and interval, and squares show the neural likelihood’s posterior mean. Within each panel, participants are ordered by that parameter’s analytical posterior mean, so rows do not correspond across panels. The dashed line marks . (b) Wall time for one fit of a simpler model, in which only the drift rate varies by participant, under each of expressions ( converged). Dashed lines mark the median expression ( s) and the analytical likelihood ( s). All fits in b used four chains on one core of the same machine.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
| runs | any | pool | recover all 4 | div. | ||||||
| pooled | per search | |||||||||
| 36 | 8 | 968 | 12 (1.2%) | 1.5 [0.5, 2.9] | 22.9% | 51.4% | 6.9% | 97.8% | 1 | |
| 20 | 10 | 560 | 55 (9.8%) | 9.5 [3.7, 16.3] | 33.9% | 78.0% | 18.6% | 99.5% | 3 | |
| 20 | 15 | 486 | 95 (19.5%) | 18.5 [10.9, 26.7] | 52.7% | 82.3% | 25.1% | 99.4% | 3 | |
| 24 | 12 | 634 | 99 (15.6%) | 14.6 [7.0, 22.9] | 42.9% | 61.2% | 30.6% | 89.4% | 7 | |
| 24 | 7 | 723 | 15 (2.1%) | 2.1 [0.7, 4.0] | 29.2% | 28.5% | 11.5% | 80.9% | 4 | |
Table 3: Searches differing only in the recovery weight of Equation 5 , with identical operators, no-constants rule, neural likelihood, test subjects and estimator (NUTS posterior means). Per-parameter columns give the proportion of expressions that pass the criterion of Section 2.3 for that parameter. Runs is the number of independently seeded searches pooled at each value, and pool sizes differ because each search ran for a fixed time. Any is the number of those searches that found at least one expression recovering all four parameters. Pooled gives that count and percentage over all expressions, and per search gives the mean percentage over searches with a bootstrap interval. The row pools two sets of MSE searches (see text). The last column is the median number of divergent transitions per expression. Between and of expressions had no divergent transitions, more so at the two ends of the sweep.
Figure 6: The neural likelihood. The joint trial log-likelihood factorizes into a flow over and a choice classifier ( Equation 3 ), trained on parameter draws simulated trials. It is admitted only after passing a diagnostic gate, with density calibration ECE (gate ), posterior-predictive KS median (gate ), choice calibration MAE , choice-rate agreement , and a cross-check against an independently pre-trained LAN, .
| weakest | ranking | div. | ||||||
| Expression 1 | 0.91 | 0.90 | 0.87 | 0.99 | 0.87 | 0.89 | 25 | 28 |
| Expression 2 | 0.90 | 0.92 | 0.90 | 0.98 | 0.90 | 0.88 | 6 | 3303 |
| Expression 3 | 0.83 | 0.89 | 0.81 | 0.99 | 0.81 | 0.88 | 36 | 377 |
| Expression 4 | 0.84 | 0.90 | 0.86 | 0.99 | 0.84 | 0.88 | 33 | 12 |
| Expression 5 | 0.95 | 0.86 | 0.88 | 0.99 | 0.86 | 0.88 | 0 | 5 |
| Known leading-term form | 0.97 | 0.86 | 0.95 | 0.99 | 0.86 | 0.88 | 3 | 553 |
Table 4: Recovery on subjects never used for ranking, with warm-up and sampling iterations in chains. Weakest is the lowest of the four correlations, and ranking gives the same quantity on the subjects used to rank the expressions ( Table 1 ). counts subjects whose chains did not converge, and div. is the total number of divergent transitions. The MSE expression fails the criterion because its estimates of sit at a bound for of subjects.
Figure 7: Simulation-based calibration for all five expressions of Table 1 . Each panel shows the difference between the empirical CDF of the posterior ranks and the uniform CDF, with a simultaneous band ( Säilynoja et al., 2022 ) , and a calibrated posterior stays inside the band. Each row uses simulations of trials, with drawn from the same uniform range as the recovery banks and data generated from each expression’s own grid-normalized density. No run produced improper draws. Curves show the simulations with and bulk-ESS ( – of ), and verdicts are unchanged on the full set. Nineteen of the twenty parameter–expression combinations lie inside the band, the exception being for Expression 2 (dispersion , ). The analytical likelihood run through the same procedure is itself outside the band on , which sets the scale for reading a single excursion.
Figure 8: Per-subject posterior predictive distributions for three subjects spanning the observed choice balance. Each panel shows the observed choice-signed RT histogram against predictive bands from one fitted likelihood. Rows are the analytical likelihood, the neural likelihood, the five expressions of Table 1 , and the MSE expression of Figure 2 . The recovery-objective expressions track the RT distributions at both boundaries, whereas the MSE expression visibly misallocates mass (median KS values in Table 1 ).