SymTorch: A Framework for Symbolic Distillation of Deep Neural Networks
This work addresses the problem of limited adoption of symbolic distillation for interpretability and efficiency in deep learning, particularly for researchers and practitioners, though it is incremental as it builds on existing symbolic regression methods.
The authors tackled the engineering barrier of integrating symbolic distillation into deep learning workflows by introducing SymTorch, a framework that automates the process and demonstrated it across various architectures, achieving an 8.3% throughput improvement in LLM inference with moderate performance degradation.
Symbolic distillation replaces neural networks, or components thereof, with interpretable, closed-form mathematical expressions. This approach has shown promise in discovering physical laws and mathematical relationships directly from trained deep learning models, yet adoption remains limited due to the engineering barrier of integrating symbolic regression into deep learning workflows. We introduce SymTorch, a library that automates this distillation by wrapping neural network components, collecting their input-output behavior, and approximating them with human-readable equations via PySR. SymTorch handles the engineering challenges that have hindered adoption: GPU-CPU data transfer, input-output caching, model serialization, and seamless switching between neural and symbolic forward passes. We demonstrate SymTorch across diverse architectures including GNNs, PINNs and transformer models. Finally, we present a proof-of-concept for accelerating LLM inference by replacing MLP layers with symbolic surrogates, achieving an 8.3\% throughput improvement with moderate performance degradation.