[ICLR 2025] SymTorch：将深度学习神经网络“翻译”成人类可读的数学公式

TL;DR

SymTorch 是一个创新的开源工具库，它打破了神经网络（NN）与符号回归（Symbolic Regression）之间的技术隔阂。它能自动捕捉 PyTorch 模型层级的输入输出行为，并将其蒸馏为简洁、直观的数学方程。无论是从 GNN 中找回物理定律，还是通过方程替换 MLP 层来给 LLM 推理“提速”，SymTorch 都展现了符号化解释带来的独特价值。

背景：当神经网络遇到物理直觉

在科学 AI（AI for Science）领域，黑盒模型（如 Transformer 或 GNN）虽然预测精湛，但无法告诉科学家它“学到了什么定律”。传统的机制可解释性关注神经元激活，而 SymTorch 关注的是：这个组件到底计算了什么函数？

作者认为，物理学中的可解释性意味着简洁（Parsimony）。牛顿第二定律 $F=ma$ 并不追求 100% 的数值模拟精度，但它抓住了系统的核心动力学。SymTorch 的核心动机就是赋予深度模型这种“公式化”的表达能力。

技术核心：自动化符号蒸馏流程

SymTorch 的工作流可以概括为：封装 -> 采样 -> 蒸馏 -> 替换。

图 1：SymTorch 捕获 NN 组件 I/O 并使用 PySR 生成 Pareto 前沿方程的流程

1. 自动化的工程处理

此前，将 PySR 应用于深度学习需要研究者手动处理 GPU 到 CPU 的张量转换。SymTorch 通过 SymbolicModel 包装器，自动管理 Forward Hooks、数据缓存和内存同步，极大地降低了准入门槛。

2. 突破维数灾难：PCA + 符号回归

符号回归（SR）的计算量随输入变量数呈指数增长。为了处理 LLM 中上千维的隐藏层，SymTorch 引入了 PCA 降维策略：先将 MLP 输入投影到低维子空间（如 32 维），进行符号回归，再将输出投影回高维。

核心实验展示

实验 A：LLM 推理加速

作者尝试将 Qwen2.5-1.5B 中的部分 MLP 层（SwiGLU）替换为符号方程。

• 效果：替换 3 层后，推理吞吐量提升了 8.3%。
• 洞察：虽然 Perplexity 有所上升，但符号化后的模型在“精度-吞吐量”Pareto 前沿上依然极具竞争力（见下图 5）。

图 2：不同 LLM 的吞吐量 vs Perplexity 比较，SymTorch 改进版保持在 Pareto 前沿

实验 B：从 PINN 中提取偏微分方程（PDE）解

在物理信息神经网络（PINN）训练后，SymTorch 成功提取了一维热传导方程的解析解 $u(x,t)=e^{-{\pi }^2\alpha t}\sin (\pi x)$，精度达到小数点后两位。

图 3：PINN 在极少数据点下通过符号蒸馏展现了极高的拟合一致性

深度洞察：LLM 到底学会了算术吗？

一个有趣的案例是分析 Llama-3.2-1B 的算术能力。通过符号蒸馏，研究者发现 LLM 进行“3位数加法”时，内部其实在使用极其复杂的非线性模拟，其中包含大量的正弦项（sin）和倒数项（inv），这解释了为什么 LLM 在处理数值边界时常会出现系统性误差。

局限与未来

1. 计算代价：符号回归仍然是暴力搜索过程，难以扩展到极宽的层。
2. 分布偏移：目前的符号代理模型在训练分布外（Out-of-Distribution）的泛化能力仍需验证。

总结

SymTorch 将“不可解释”的权重矩阵转化为了“可读”的代数语言。它不仅为 AI for Science 研究者提供了一把手术刀，也为未来开发“神经-符号”混合架构的轻量化模型打开了新大门。

主编点评：这是一项非常务实的工程创新。它没有去发明新的回归算法，而是通过解决 PyTorch 内部复杂的 Hook 管理和加速设备通信模型，让“符号蒸馏”这一学术概念真正走向了工业界的可操作化。