Recovery-Directed Symbolic Distillation of Neural Likelihoods
Authors: Kianté Fernandez, Xinwei Li
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
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 ( v , a , z , t ). Each panel title gives the Pearson r across subjects. (a) True parameter value ( x -axis) against the posterior mean estimated with NUTS ( y -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 .
λ
Cf
MSE
v
a
z
t
held-out
KS
beats null
1
3
44
0.32
0.96
0.91
0.89
0.98
0.87
0.134
✓
2
3
25
1.05
0.94
0.91
0.88
0.97
0.90
0.130
✓
3
1
42
0.27
0.88
0.90
0.88
0.98
0.81
0.140
✓
4
1
45
0.27
0.88
0.92
0.88
0.98
0.84
0.193
–
5
3
42
0.51
0.95
0.88
0.88
0.98
0.86
0.160
✓
Table 1: The top five expressions’ performance on the test bank. λ is the recovery weight in Equation 5 , Cf is the expression complexity, MSE is its squared error against the neural log-likelihood on held-out rows (pool median 0.52 ), and the v , a , z and t columns give the recovery correlation r for each parameter. Held-out gives the weakest of the four correlations on 90 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 .287 for the MSE expression in Figure 2 and about .04 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 r=.5 . (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
1000
0.010
0.015
0.27
0.59
0.68
1.3
1144
100,000
0.057
0.28
29
76
13
35
38,080
1,000,000
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 20 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 96 converged expressions. Black points and bars show the analytical likelihood’s posterior mean and 89% 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 z=0.5 . (b) Wall time for one fit of a simpler model, in which only the drift rate varies by participant, under each of 267 expressions ( 213 converged). Dashed lines mark the median expression ( 14 s) and the analytical likelihood ( 386 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
v
a
z
t
div.
pooled
per search
0
36
8
968
12 (1.2%)
1.5 [0.5, 2.9]
22.9%
51.4%
6.9%
97.8%
1
0.3
20
10
560
55 (9.8%)
9.5 [3.7, 16.3]
33.9%
78.0%
18.6%
99.5%
3
1
20
15
486
95 (19.5%)
18.5 [10.9, 26.7]
52.7%
82.3%
25.1%
99.4%
3
3
24
12
634
99 (15.6%)
14.6 [7.0, 22.9]
42.9%
61.2%
30.6%
89.4%
7
10
24
7
723
15 (2.1%)
2.1 [0.7, 4.0]
29.2%
28.5%
11.5%
80.9%
4
Appendix
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 95% bootstrap interval. The λ=0 row pools two sets of MSE searches (see text). The last column is the median number of divergent transitions per expression. Between 14% and 41% 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 RT and a choice classifier ( Equation 3 ), trained on 15,000 parameter draws ×20 simulated trials. It is admitted only after passing a diagnostic gate, with density calibration ECE .041 (gate <0.10 ), posterior-predictive KS median .092 (gate <0.15 ), choice calibration MAE .017 , choice-rate agreement r=.996 , and a cross-check against an independently pre-trained LAN, r=.982 .
v
a
z
t
weakest
ranking
R>1.01
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
Appendix
Table 4: Recovery on 90 subjects never used for ranking, with 2000 warm-up and 2000 sampling iterations in 4 chains. Weakest is the lowest of the four correlations, and ranking gives the same quantity on the 60 subjects used to rank the expressions ( Table 1 ). R>1.01 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 z sit at a bound for 71% 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 95% simultaneous band ( Säilynoja et al., 2022 ) , and a calibrated posterior stays inside the band. Each row uses 200 simulations of 1000 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 R≤1.05 and bulk-ESS ≥100 ( 167 – 194 of 200 ), and verdicts are unchanged on the full set. Nineteen of the twenty parameter–expression combinations lie inside the band, the exception being z for Expression 2 (dispersion 1.03 , χ2p=.50 ). The analytical likelihood run through the same procedure is itself outside the band on a , 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 ).
What mathematical functions do neural network components learn? Symbolic distillation addresses this question by expressing neural network components with interpretable, closed-form mathematical expressions that expose the functional structure learned during training. We develop symbolic distillation as a systematic, architecture-agnostic methodology, and release our approach as the open-source SymTorch package - a PySR-powered library built natively for the PyTorch ecosystem. Applying this methodology across diverse architectures, we find that SymTorch is successful in the automated discovery of physical laws. Specifically, our approach (1) recovers pairwise interaction forces from graph neural networks trained on empirical n-body observations, (2) distills the exact closed-form PDE/ODE solutions of multiple physical systems, including the value of constants, from physics-informed neural networks trained on sparse data, and (3) uncovers the chaotic dynamics of the Lorenz system from high-dimensional data, ultimately outperforming the base neural network on downstream prediction tasks. We further demonstrate the utility of our framework for model interpretability by providing an optimized implementation of SLIME - a symbolic extension to the LIME explainability method. SLIME consistently outperforms LIME across predictive metrics across eight popular classification and regression benchmarks, while still providing an interpretable local symbolic model. Lastly, we investigate replacing transformer MLP layers with symbolic surrogates: replacing 1-7 layers with symbolic approximations yields 2-19% throughput improvements and up to 18.7% VRAM reduction, with the resulting hybrid models lying on the Pareto front of throughput versus perplexity among open-source LLMs of comparable scale.
Elizabeth S. Z. Tan, Adil Soubki, Miles Cranmer
Department of Applied Mathematics and Theoretical Physics University of Cambridge
Symbolic regression discovers explicit, interpretable equations without assuming a functional form in advance. A Bayesian approach strengthens this through probability distributions over candidate expressions, thus quantifying uncertainty in the presence of noisy and limited data. Deep Symbolic Regression (DSR) uses a neural network to generate symbolic expressions, but it is designed to identify a single best-fitting expression rather than infer a posterior distribution over models. We introduce Deep Variational Inference Symbolic Regression (DVISR), a variational Bayesian extension of DSR. DVISR replaces the original reward with the integrand of the evidence lower bound. It also extends the network architecture to output distributions over constants within expressions, enabling posterior inference over both expression trees and their associated constants. We show that DVISR can recover the true posterior in simple settings, both with and without constant tokens, and we examine how its performance changes as the size of the expression space increases. These results position DVISR as a step toward scalable Bayesian symbolic regression with uncertainty over full symbolic models.
James Butterworth, Gevik Grigorian, Alejandro DiazDelaO
Clinical Operational Research Unit University College London London, United Kingdom · Department of Computer Science University College London London, United Kingdom
Neurosymbolic (NeSy) models integrate neural networks and symbolic reasoning for robust and interpretable AI. State-of-the-art NeSy models require that the symbolic component is expressed in a differentiable way, often complicating the use of approximate inference. We propose EM-NeSy which casts probabilistic NeSy learning as an instance of the Expectation-Maximization (EM) algorithm. In the expectation step, we compute the posterior over the neurally predicted symbols conditioned on the label via probabilistic inference. In the maximization step, we update the neural parameters based on this posterior using gradient descent only through the neural component. This formulation unlocks the full potential of the EM algorithm for NeSy learning. It allows NeSy to extend naturally to approximate reasoning without any additional modifications or differentiability requirements of the symbolic component. Furthermore, it recovers the standard end-to-end gradient-based NeSy setting under exact inference. Our experimental results demonstrate the scalability and computational efficiency of EM-NeSy.