Symbolic Governing Equation Discovery Using Neural Arithmetic Modules

TMLR Paper7906 Authors

Abstract: Neural architectures with arithmetic inductive biases, such as Neural Arithmetic Logic Units (NALU) and Neural Power Units (NPU), are designed to model arithmetic relationships for improved out-of-distribution extrapolation and interpretability. However, in practice, these models frequently exhibit unstable optimization behaviours such as gradient starvation and convergence to dense and numerically fragile parameterizations that obscure the underlying data structure. We show that arithmetic inductive bias alone is insufficient to guarantee the recovery of sparse symbolic equations. Instead, interpretability should be explicitly enforced through strict architectural constraints. We propose MSRNet, a structured neural framework for extracting sparse symbolic expressions from high-dimensional data. The model has two variants: MSRNet (Multiplicative Symbolic Regression Network), which allows multiplicative and discrete exponential arithmetic interactions via differentiable softmax relaxations, and ExMSRNet (Extended MSRNet), which further allows for logarithmic and exponential pathways. We use a composite training objective that utilizes description-length regularization via entropy-based measures to bias the model towards confident discrete operator selection. Our experiments suggest that MSRNet variants significantly reduce gradient starvation. This could be attributed to explicit constraining of the hypothesis space. We benchmark MSRNet variants on synthetic datasets, SRBench 2025, and AI Feynman I/II/III, where it achieves strong performance with significantly lower computational cost than other symbolic regression methods. Source code for MSRNet is available at: https://anonymous.4open.science/r/MSRNet-6B05

Submission Type: Long submission (more than 12 pages of main content)

Assigned Action Editor: ~Shandian_Zhe1

Submission Number: 7906