SymTorch: Symbolic Distillation of Neural Networks
Organizations: Department of Applied Mathematics and Theoretical Physics University of Cambridge
Abstract
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 -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.
Figures & tables
| Operation | Expected Equation | LLM-Learned Operation |
|---|---|---|
| Addition | ||
| Multiplication | ||
| Counting | ||
| Temperature Conversion |
Appendix figures & tables23 assets
Supplementary material from the paper’s appendix.
Appendix
| SR Parameter | Configuration |
|---|---|
| Binary operators | +, * |
| Unary operators | inv(x)=1/x, sin, exp |
| Extra sympy mappings | "inv": lambda x: 1/x |
| Number of iterations | 400 |
| Complexity of operators | sin: 3, exp: 3 |
| Simulation | Standard | Bottleneck | L1 | KL | Pruning | Variance |
|---|---|---|---|---|---|---|
| Charge | 18.20 | 19.12 | 18.06 | 39.80 | 19.40 | 56417.66 |
| 0.26 | 0.25 | 0.32 | 15.40 | 0.28 | 87.69 | |
| 24.07 | 25.25 | 21.67 | 57.80 | 23.77 | 98641.90 | |
| Spring | 0.24 | 0.16 | 0.20 | 7.29 | 0.23 | 55.84 |
| Simulation | Message | Standard | Bottleneck | L1 | KL | Pruning |
|---|---|---|---|---|---|---|
| Charge | 1 | ✗ | ✓ ‡ | ✗ | ✗ | ✓ |
| 2 | ✗ | ✓ ‡ | ✓ § | ✗ | ✓ § | |
| 1 | ✗ | ✓ | ✓ | ✗ | ✓ | |
| 2 | ✓ | ✓ | ✓ | ✗ | ✓ ∗ | |
| 1 | ✓ § | ✓ | ✓ § | ✓ ‡ | ✓ | |
| 2 | ✗ | ✓ | ✗ | ✓ ‡ | ✓ |
| Complexity | Equation | Loss |
|---|---|---|
| 1 | 0.0888 | |
| 5 | 0.0885 | |
| 7 | 0.0882 | |
| 8 | 0.0880 | |
| 9 | 0.0877 | |
| 10 | 0.0843 |
| Model | Train MSE | Test MSE | Train variance | Test variance |
|---|---|---|---|---|
| NN | 8.22 | 8.18 | 8646.09 | 8631.27 |
| Symbolic | 15.52 | 15.65 | ||
| Finetune | 12.51 | 12.66 |
| SINDy autoencoder | Ours (SymTorch) | |
|---|---|---|
| SR method | Sparse regression over fixed library | PySR (genetic algorithm, library-free) |
| Function library | Must be specified in advance | Not required |
| Encoder architecture | Unconstrained MLP | Single linear layer |
| Latent dynamics | Linear combination of library terms | Arbitrary analytic expressions |
| System | True solution | Distilled expression | True constant | SR estimate |
|---|---|---|---|---|
| Heat equation | ||||
| Wave equation | ||||
| Oscillator |
| Dataset | Task | |||||
|---|---|---|---|---|---|---|
| Iris | Classification | 120 | 30 | 4 | 10 | 500 |
| Wine | Classification | 142 | 36 | 13 | 10 | 500 |
| Breast Cancer | Classification | 455 | 114 | 30 | 22 | 500 |
| Pima Diabetes | Classification | 614 | 154 | 8 | 30 | 500 |
| German Credit | Classification | 800 | 200 | 61 | 40 | 500 |
| Adult | Classification | 39,073 | 9,769 | 108 | 100 | 500 |
| LIME | SLIME | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Point index | KL div | WCE | WP-MSE | Agree. | KL div | WCE | WP-MSE | Agree. | Example equation for prob. of class 0 |
| 4116 | 0.73 | 0.61 | 0.01 | 1.00 | 0.23 | 0.61 | 0.00 | 1.00 | |
| 11317 | 0.95 | 1.17 | 0.89 | 0.05 | 0.00 | 1.69 | 0.08 | 0.95 | |
| 11038 | 0.31 | 0.48 | 0.29 | 1.00 | 0.00 | 0.10 | 0.02 | 1.00 | |
| 184 | -0.00 | 0.00 | 0.00 | 1.00 | 0.68 | 0.27 | 0.11 | 1.00 | |
| 30797 | 0.38 | 1.17 | 0.70 | 0.32 | 0.04 | 15.94 | 0.98 | 0.32 | |
| LIME | SLIME | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Point index | KL div | WCE | WP-MSE | Agree. | KL div | WCE | WP-MSE | Agree. | Example equation for prob. of class 0 |
| 47 | 0.66 | 0.77 | 0.56 | 0.09 | 0.00 | 0.04 | 0.00 | 1.00 | |
| 130 | 1.03 | 1.22 | 0.89 | 0.00 | 0.01 | 0.31 | 0.05 | 1.00 | |
| 128 | 0.00 | 0.01 | 0.00 | 1.00 | 0.00 | 0.01 | 0.00 | 1.00 | |
| 2 | 0.04 | 0.07 | 0.00 | 1.00 | 0.00 | 0.04 | 0.00 | 1.00 | |
| 357 | 3.42 | 4.18 | 0.06 | 1.00 | 0.00 | 3.75 | 0.05 | 1.00 | |
| LIME | SLIME | ||||||
|---|---|---|---|---|---|---|---|
| Point index | WR 2 | WMAE | WMSE | WR 2 | WMAE | WMSE | Example equation |
| 1739 | -39.59 | 4.90 | 24.30 | -0.17 | 0.67 | 0.70 | |
| 4782 | -74.08 | 4.21 | 18.47 | 0.21 | 0.35 | 0.19 | |
| 4664 | -6.64 | 2.95 | 9.02 | 0.81 | 0.36 | 0.22 | |
| 77 | -30.12 | 2.96 | 9.40 | -0.07 | 0.41 | 0.32 | |
| 13014 | -31.57 | 3.28 | 11.10 | -0.18 | 0.48 | 0.40 | |
| LIME | SLIME | ||||||
|---|---|---|---|---|---|---|---|
| Point index | WR 2 | WMAE | WMSE | WR 2 | WMAE | WMSE | Example equation |
| 86 | -22.94 | 1.28 | 1.68 | 0.21 | 0.19 | 0.06 | |
| 237 | -13.86 | 1.85 | 3.49 | 0.62 | 0.26 | 0.09 | |
| 232 | -54.81 | 1.38 | 2.00 | -0.74 | 0.20 | 0.06 | |
| 3 | 0.69 | 0.18 | 0.05 | 0.81 | 0.13 | 0.03 | |
| 648 | -9.38 | 1.02 | 1.15 | 0.22 | 0.23 | 0.09 | |
| LIME | SLIME | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Point index | KL div | WCE | WP-MSE | Agree. | KL div | WCE | WP-MSE | Agree. | Example equation for prob. of class 0 |
| 83 | 1.48 | 1.65 | 0.01 | 1.00 | 0.01 | 0.24 | 0.00 | 1.00 | |
| 230 | 0.01 | 0.39 | 0.01 | 1.00 | 0.00 | 0.76 | 0.01 | 1.00 | |
| 225 | 0.06 | 0.85 | 0.15 | 0.05 | 0.00 | 0.74 | 0.05 | 0.05 | |
| 3 | 3.42 | 5.14 | 0.07 | 1.00 | 0.00 | 0.48 | 0.00 | 1.00 | |
| 629 | 1.63 | 2.52 | 0.02 | 1.00 | 0.00 | 0.31 | 0.00 | 1.00 | |
| LIME | SLIME | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Point index | KL div | WCE | WP-MSE | Agree. | KL div | WCE | WP-MSE | Agree. | Example equation for prob. of class 0 |
| 12 | 0.00 | 0.04 | 0.00 | 1.00 | 0.00 | 0.02 | 0.00 | 1.00 | |
| 33 | 1.40 | 1.42 | 1.14 | 0.00 | 0.00 | 0.02 | 0.00 | 1.00 | |
| 119 | 0.66 | 0.69 | 0.48 | 0.48 | 0.00 | 0.06 | 0.00 | 1.00 | |
| 0 | 0.41 | 0.41 | 0.23 | 1.00 | -0.00 | 0.00 | 0.00 | 1.00 | |
| 93 | 0.05 | 0.36 | 0.02 | 1.00 | 0.00 | 0.30 | 0.00 | 1.00 | |
| LIME | SLIME | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Point index | KL div | WCE | WP-MSE | Agree. | KL div | WCE | WP-MSE | Agree. | Example equation for prob. of class 0 |
| 64 | 0.17 | 0.70 | 0.15 | 0.34 | 0.01 | 0.55 | 0.00 | 0.99 | |
| 176 | 0.51 | 0.77 | 0.43 | 0.04 | 0.00 | 0.30 | 0.00 | 1.00 | |
| 173 | 0.02 | 0.68 | 0.03 | 0.91 | 0.00 | 0.66 | 0.02 | 0.98 | |
| 2 | 0.19 | 0.73 | 0.19 | 0.34 | 0.01 | 0.56 | 0.03 | 0.97 | |
| 483 | 0.90 | 1.20 | 0.76 | 0.00 | 0.00 | 0.38 | 0.01 | 1.00 | |
| LIME | SLIME | ||||||||
|---|---|---|---|---|---|---|---|---|---|
| Point index | KL div | WCE | WP-MSE | Agree. | KL div | WCE | WP-MSE | Agree. | Example equation for prob. of class 0 |
| 14 | 0.38 | 0.79 | 0.25 | 1.00 | 0.00 | 0.34 | 0.00 | 1.00 | |
| 39 | 1.04 | 1.52 | 0.06 | 1.00 | 0.00 | 0.25 | 0.00 | 1.00 | |
| 40 | 0.28 | 0.52 | 0.13 | 1.00 | 0.00 | 0.23 | 0.00 | 1.00 | |
| 0 | 0.24 | 0.45 | 0.03 | 1.00 | 0.00 | 0.24 | 0.00 | 1.00 | |
| 111 | 2.65 | 3.54 | 0.07 | 1.00 | 0.00 | 0.57 | 0.00 | 1.00 | |
| Operation | Example Prompt |
|---|---|
| Addition of two 3-digit numbers | Return only the numeric answer in the format . What is ? |
| Multiplication of two 3-digit numbers | Return only the numeric answer in the format . What is ? |
| Counting 1s in a 6-digit string of 1s and 0s | Return only the numeric answer in the format . How many 1s are there in the string 000101? |
| Converting from Celsius to Fahrenheit | Return only the numeric answer in the format . What is degrees Celsius in Fahrenheit? |
| Layer | PCA | PCA+SR |
|---|---|---|
| 0 | ||
| 1 | ||
| 2 | ||
| 3 | ||
| 4 | ||
| 5 |
| PPL | Throughput (%) | Memory Saved | |||||
|---|---|---|---|---|---|---|---|
| Layers | SR | PCA | Skip | SR | Skip | (MB) | (%) |
| 1 | |||||||
| 2 | |||||||
| 3 | |||||||
| 4 | |||||||
| 5 | |||||||
| Model | Perplexity on Validation Set | Avg. Latency (ms) | p95 Latency (ms) | Throughput (tokens/s) |
|---|---|---|---|---|
| Qwen2.5-1.5B (Baseline) | 10.79 | 216.99 | 217.13 | 4719.07 |
| Qwen2.5-1.5B SR (15) | 11.24 | 211.87 | 212.07 | 4833.25 |
| Qwen2.5-1.5B SR (15, 17) | 11.77 | 207.14 | 207.58 | 4943.49 |
| Qwen2.5-1.5B SR (15, 17, 10) | 12.44 | 202.16 | 202.41 | 5065.23 |
| Qwen2.5-1.5B SR (15, 17, 10, 13) | 13.34 | 197.64 | 197.86 | 5181.18 |
| Qwen2.5-1.5B SR (15, 17, 10, 13, 14) | 14.21 | 193.12 | 193.32 | 5302.48 |