用数学表达式自动提取神经网络学到的物理规律,让模型更可解释且更高效。
SymTorch: Symbolic Distillation of Neural Networks
- 通过符号蒸馏将神经网络转化为可读的数学公式,不依赖具体架构。
- 从少量数据中恢复出物理方程和常数,甚至能捕捉混沌系统的动态规律。
- 提升模型解释性,还能替换Transformer层以加速推理并节省显存。
神经网络组件究竟学到了何种数学函数?符号蒸馏通过可解释的闭合形式数学表达式揭示训练中学习到的功能结构。我们提出一种系统化、架构无关的符号蒸馏方法,并发布开源工具 SymTorch——一个基于 PySR、原生集成 PyTorch 生态的库。在多种架构上应用该方法,发现 SymTorch 能成功自动发现物理定律:(1) 从训练于 $n$-体观测数据的图神经网络中恢复成对相互作用力;(2) 从稀疏数据训练的物理信息神经网络中蒸馏出多个物理系统的精确闭式偏微分方程/常微分方程解,包括常数取值;(3) 从高维数据中揭示洛伦兹系统的混沌动力学,最终在下游预测任务上表现优于原始神经网络。进一步展示框架在模型可解释性方面的价值:优化实现的 SLIME(LIME 的符号扩展)在八个主流分类与回归基准上持续优于 LIME,同时提供可解释的局部符号模型。最后,探索用符号代理替代 Transformer MLP 层:替换 1-7 层可带来 2-19\/% 的吞吐量提升,最高减少 18.7\/% 显存占用,所得混合模型在同等规模开源大语言模型中位于吞吐量与困惑度的帕累托前沿。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。