Graph-based Symbolic Regression with Invariance and Constraint Encoding

We sincerely thank the reviewer for the constructive feedback and insightful questions. Below, we address each point in turn to clarify the framework and resolve any ambiguities.

1) Clarity in Mathematical Formulations

We appreciate the reviewer’s comments on notation clarity and will revise Section 3 for greater precision. Specifically:

• The prior probability matrix $\mathbf{P} \in \mathbb{R}^{|\mathcal{V}| \times |\mathcal{Q}|}$ (line 130, $|\mathcal{V}|$ is the number of nodes and $|\mathcal{Q}|$ is the size of operator dictionary) is defined as a matrix where each row corresponds to a node in the EG, and each column to an operator $\phi \in \mathcal{Q}$, giving the prior distribution for adding $\phi$ at that node. Line 132 clarifies that $\mathbf{P}$ and $r$ are predicted by the R-GCN-based policy $\pi_\theta$ (Eq. 1), with $\pi_\theta | \mathbf{P}$ and $\pi_\theta | r$ as separate MLPs on the R-GCN embeddings (Eq. 2, Appendix C.3). This is consistent: the initial mention introduces the concept, and the follow-up specifies the neural prediction.

• Actions $a_t$ are defined in the MDP settings (Sec. 3.2): at step $t$, $a_t$ is the operator $\phi \in \mathcal{Q}$ added to the current EG state $s_t = G_t$, growing the graph deterministically (e.g., adding a child node/edge). The policy $\pi_\theta | \mathbf{P}$ parameterizes the prior distribution over $a_t$ through $P(s, a)$ in PUCB (Eq. 3). We will add explicit definitions and parameterizations in revisions to eliminate ambiguity.

2) Novelty of hnMCTS

While hnMCTS builds upon naive and neural-guided MCTS [8, 13, 18], its novelty lies in its adaptive $\varepsilon$-policy (Eq. 3), which enables hybrid integration of exploration and guidance. This design allows early random exploration, followed by a smooth transition to neural guidance as R-GCN stabilizes. Expression-level constraints are directly encoded via priors $\mathbf{P}, r$, influencing the sampling policy during search. Furthermore, we provide convergence guarantees under this scheme (Appendix C.5, Theorems C.1–C.6).

This is not a simple combination of components: the $\varepsilon$-policy decay (Appendix C.6) is specifically introduced to mitigate sparse rewards, a known challenge in symbolic regression. This mechanism improves sample efficiency (Table 1) and enhances the balance between exploration and exploitation (Appendix C.5.2, D.3).

3) Limitation of EG Invariance

EG captures a broad class of equivalences beyond commutativity, including associative, distributive, identity, exponential, and trigonometric rules (Table 3, Appendix C.2). These are modeled through a convergent term rewriting system (TRS) that is both terminating and confluent, yielding unique normal forms. EG encodes these canonical forms as DAGs. We provide a detailed proof of this uniqueness property, quantify its impact on search efficiency, and discuss remaining limitations in our response to Reviewer RJ7g (Section 1), which we refer to for completeness.

4) Sequential EG Generation

We thank the reviewer for raising this insightful question. We would like to clarify that the sequential generation scheme remains efficient. While different traversal orders can construct the same EG, they often pass through distinct intermediate expressions, which are not mathematically equivalent. Therefore, these paths shouldn't be merged and do not result in redundant representations, as the same EG is only one of several possible descendants from such intermediates. The permutation invariance encoded by EG only applies to equivalent expressions with the same descendants, and different intermediate expressions are not redundant cases to inflate the state space.

5) Incorporating Expression-Level Constraints

Expression-level constraints are incorporated proactively via R-GCN-predicted priors $\mathbf{P} \in \mathbb{R}^{|\mathcal{V}| \times |\mathcal{Q}|}$ and scalar scores $r$, which leverage structural features (e.g., node/edge features) to estimate physical validity (e.g., favoring $r^{-m}$ near singularities; Eq. 5, Appendix C.3). These priors guide hnMCTS through $P(s, a)$ in PUCB (Eq. 2), biasing the search toward constraint-satisfying expressions early.

Unlike post-hoc filtering, this inductive bias reduces invalid rollouts and densifies the reward landscape, even though full evaluations are only performed at leaf nodes. We will clarify this mechanism in Sec. 3.1 and provide a concrete example in the final version.

We thank the reviewer again for the detailed and thoughtful feedback. We believe these clarifications address the raised concerns and will further enhance the presentation in the revised manuscript.

Best regards,
The Authors

We appreciate the reviewer bn25' detailed comments. To help address the raised concerns and questions, please allow us to provide point-by-point responses to clarify our method for better understanding.

1) Links to Compiler Theory for ET-to-EG Conversion

The conversion from Expression Trees (ETs) to EGs shares similarities with compiler optimizations on Abstract Syntax Trees (ASTs). In compilers, ASTs are often transformed to Directed Acyclic Graphs (DAGs) for expressions to enable common subexpression elimination (CSE) and algebraic simplifications (e.g., flattening associative operations, constant folding). Our EG construction (Appendix C.2) mirrors this: node sharing for common feature nodes (DAG property), unification of binaries (e.g., +/− to $\sum$ with edge features), and rewriting for equivalences (Table 3, like distributive expansion). Standard compiler techniques include bottom-up DAG construction during parsing (e.g., in GCC/LLVM for intermediate representations) and term rewriting for optimizations. However, EGs extend this with R-GCN-compatible features (e.g., operand relations $\alpha$) tailored for SR invariance. While no exact "standard" exists for SR-specific equivalences, our rules (TRS, defined in our response to Reviewer RJ7g Section 1, to yield unique EG representations) align with algebraic rewriting in computer algebra systems (e.g., SymPy). We will add these connections to Appendix B in revisions.

2) Literature Review and Appendix Usage

We thank the reviewer bn25 for the suggestion and will discuss Simon's BACON in the literature review for completeness. We agree that the main paper should be self-contained, and we aimed for that by including core descriptions and results in the main paper and providing optional expansions in the appendix. To improve clarity, we will include the key literature review and reproducibility details (e.g., experimental settings) in the main paper in the camera-ready version given the extended page limit.

3) Clarification on Neural-Guided Proposals and the Data Role

Thanks for raising questions in the GNN proposal mechanism, and we would like to provide more clarifications to assist better understanding. The Expression Graph as the DAG $G_t$ is represented as the state of an MDP, which the Relational GCN (R-GCN) encodes to predict a prior probability matrix $\mathbf{P}$ (operator selection distribution per node) and reward $r$ (estimated utility; Eq. 1-2). These guide hnMCTS simulations (Algorithm 2), balancing exploration via PUCB (Eq. 3).

The network does not directly encode the data $(X, y)$; instead, data influences proposals indirectly via RL rewards $R(\tau)$ from expression evaluations (fit to data) for current $G_t$, and are backpropagated to update $\pi_\theta$ (Eq. 4). This indirect data influence (via rewards) allows learning general structural priors (e.g., favoring valid constraints) without overfitting to specific datasets. Encoding data directly (e.g., via embeddings) could bias toward numerical patterns over symbolic ones, complicating generalization, and leading to computational overhead. Future work could explore more carefully designed hybrid data-graph inputs for data-aware GNNs to balance the computational cost and guidance enhancement.

4) Learning Approach

The model learns from scratch for each dataset, as SR is typically per-problem (exploratory discovery). This avoids foundation-model biases from unrelated data, but pretraining on synthetic expressions (as in Sec. 6) is a promising extension.

5) AIFeynman NaN Score

In SRBench, AIFeynman's NaN in R2 (Fig. 4a) arises from precision artifacts: exact recoveries yield zero error, causing NaN in some metrics (e.g., division by zero in normalized scores). This is a known benchmarking issue, and we report raw SRBench results for fairness.

6) Limited Baselines on Physics Dataset

The materials science dataset (Appendix F) requires physics-specific constraints (e.g., the essential scalar output constraint for valid evaluation; Sec. 5.2), for which many SOTA SR methods (e.g., AIFeynman, Operon) lack immediate support. We compared with fitting baselines: CGCNN (graph NN surrogate), GPs (from [25]), and MCTS (constraint-aware variant). GSR's superiority (lowest MAE, satisfies all constraints; Table 2) highlights its value. We will discuss this rationale in our revision.

7) No Test-Set Overfitting

We confirm that all developments were on separate datasets (e.g., Nguyen for initial tuning, holdouts from SRBench) without access to test sets during algorithm design. Hyperparameters (Appendix C.6) were fixed before final runs, following standard practices to ensure fair comparisons.

Overall, we thank the reviewer again for the detailed comments and hope that we have addressed all the reviewer's concerns and questions.

Best Regards,

The Authors

We thank the reviewer for their thoughtful comments and valuable suggestions. Below, we address each point in detail to clarify the contributions and scope of our proposed GSR framework.

1) Novelty of Constraint Handling

We appreciate the opportunity to clarify the novelty of our constraint incorporation strategy. The key distinction lies in the proactive constraint encoding via hnMCTS, which integrates expression-level constraints—such as $f(r \to 0) = \infty$—through R-GCN-predicted priors $\mathbf{P} \in \mathbb{R}^{|\mathcal{V}| \times |\mathcal{Q}|}$ (line 130) and scalar validity scores $r$, both computed from global EG features (e.g., node/edge types $\alpha$) (Appendix C.3).

These priors guide hnMCTS via $P(s, a)$ in the PUCB formula (Eq. 3), steering the search towards constraint-satisfying expressions early. This approach contrasts with post-hoc filtering, as it reduces invalid rollouts before evaluation—leading to a ~60% reduction in reward sparsity (Fig. 6a). While tree-based methods and EG-based methods differ primarily in invariance encoding, EGs enable richer structural (node/edge) features that strengthen neural guidance during search. As shown in the ablation study (Sec. 4.2), EG improves constraint learning even under the same hnMCTS setup. We will clarify this distinction and provide a motivating example in Sec. 3.2 of the revised version.

2) Feynman Metrics and Baselines

Thank you for your suggestions regarding benchmarking completeness. For the Feynman dataset (Fig. 4d), we adopt solution rate (exact recovery) as the primary metric, following prior SRBench protocols. For recovered expressions, $R^2 = 1$; for approximate solutions, the average $R^2 \approx 0.92$, based on best-fit expressions. Our discovered expressions have a median complexity of ~20 operations, comparable to DGSR.

As SRBench’s official release lacks comprehensive reporting of $R^2$, expression complexity, and runtime, we will extend our benchmarking with these additional metrics via re-implementation. We also plan to include recent symbolic regression baselines (e.g., PySR [Cranmer], uDSR [Landajuela et al.], TPSR [Shojaee et al.]) in our evaluation framework. If accepted, we will include these comparisons in the camera-ready version.

3) Application Scope and Generalizability

The 32-atom copper dataset from first-principles molecular dynamics [25] serves as a representative benchmark for materials science, consistent with typical system sizes in datasets like MPtraj (~31 atoms/structure [40]). While training occurs at this scale, the discovered analytical potential $f(r)$ (Table 2) generalizes to much larger systems. As shown in Appendix F.2, it enables efficient simulations involving millions of atoms, such as modeling copper crystallization [41], which are beyond the reach of DFT.

The incorporated constraints—unit consistency and electrostatic repulsion—are foundational physics principles rather than "basic" in a dismissive sense, which ensure valid behavior at short ranges and prevent simulation crashes (Fig. 5). These constraints are either absent or improperly handled in current machine learning models of interatomic potentials, leading to divergences in force calculations and causing real-world simulations using these models to crash. Without these fundamental constraints (though look simple), most, if not all, machine learning models cannot work properly when the system moves away from the equilibrium position, e.g., under temperature or strain fluctuation during molecular dynamics simulations.

4) Computational Expense and Scalability

EGs are DAGs with node sharing, resulting in fewer nodes than equivalent expression trees (Appendix C.2), which helps reduce memory and computational overhead. Moreover, our hnMCTS design does not rely solely on neural-guided search; the UCB component ensures efficiency during the exploration. Together, these design choices contribute to the method’s practical efficiency.

As shown in the ablation study (Sec. 4.2), GSR achieves strong performance while requiring the least training time among competitive baselines. Runtime configurations and sampling throughput are provided in Appendix A.

Regarding scalability, GSR performs robustly across datasets with varying input dimensionality. For example, in the black-box benchmark (Fig. 4c), the average training time remains manageable across increasing dimensions. We will include wall-clock timing and scaling results in Appendix D of the camera-ready version.

We again thank the reviewer for their detailed and constructive feedback. We will incorporate the clarifications above, expand on limitations, and include additional evaluations and discussion in the final submission.

Best regards,

The Authors

We thank reviewer RJ7g for the insightful questions on the uniqueness of Expression Graphs (EGs) and application scope. We clarify that, for the equivalences in Table 3, EGs yield unique representations (up to graph isomorphism) for equivalent expressions, modeled as a convergent term rewriting system (TRS). Below, we provide a formal proof addressing nested operators, quantify improvements, and discuss residuals/future work.

1.1) Theoretical Proof of EG Uniqueness

Definition: Let $\mathcal{Q}$ be the operator dictionary (Sec. 2.1), and $T(\mathcal{Q})$ the set of expressions over $\mathcal{Q}$. Table 3 defines a term rewriting system (TRS) $R$ with rewrite rules for equivalences. Terms $t_1, t_2 \in T(\mathcal{Q})$ are equivalent ($t_1 \sim_R t_2$) if they rewrite to a common term under $\leftrightarrow_R^*$ (reflexive-transitive-symmetric closure of $R$). EG is a function $EG: T(\mathcal{Q}) \to \mathcal{G}$, where $\mathcal{G}$ is DAGs with nodes labeled by operators/features (e.g., arity) and edges by relations $\alpha$ (e.g., base/exponent for $\wedge$; Fig. 2a). EG normalizes via bottom-up $R$ application (Appendix C.2). Uniqueness: $t_1 \sim_R t_2 \implies EG(t_1) \cong EG(t_2)$ (isomorphism).

Proof: Uniqueness requires $R$ convergent (confluent + terminating), ensuring unique normal forms $nf(t)$.

• Termination: Finite expressions yield finite steps; $R$ satisfies Termination as rules are non-recursive (no loops recreating left-hand sides).

• Confluence: If $t \to^* u$ and $t \to^* v$, then there exists $w$ such that $u \to^* w$ and $v \to^* w$. $R$ is confluent, as rules are orthogonal (left-linear, non-overlapping). Arithmetic rules (commutative/associative/distributive/identity) mirror polynomial rewriting to sum-of-products form, confluent with multisets for commutative operands [1]. Trigonometric/exponential rules do not overlap with arithmetic rules, preserving confluence via modularity [2]. By Newman’s Lemma, termination and local confluence imply global confluence, so $R$ has a unique $nf(t)$.

• EG Encodes Uniquely: Builds minimal DAG from $nf(t)$ (Appendix C.2): flattens chains (multi-operand nodes with multisets), expands distributivity, simplifies identities/trigonometric. Fixed features ensure uniqueness up to isomorphism; node sharing for minimality, multisets for order-independence.

Thus, $t_1 \sim_R t_2 \implies nf(t_1) = nf(t_2) \implies EG(t_1) \cong EG(t_2)$. Holds for nested cases (e.g., $(a \times (b + c)) + d$) via recursive normalization (Fig. 1).

We will further clarify these points with the above analyses in our Appendix.

[1] Buchberger, Bruno. An algorithm for finding a basis for the residue class ring of a zero-dimensional polynomial ideal. Diss. Ph. D. thesis, University of Innsbruck, Austria, 1965.

[2] Huet, Gérard. "Confluent reductions: Abstract properties and applications to term rewriting systems: Abstract properties and applications to term rewriting systems." Journal of the ACM (JACM) 27.4 (1980): 797-821.

1.2) Quantification

While full algebraic uniqueness is intractable, EG captures dominant redundancies. In the ablation study (Table 1), MCTS‑EG (invariance only) reduces the average search‐space size by 33.9% relative to naive MCTS, directly measuring merged states due to EG's improved uniqueness.

We also notice that ~95% of merges in ablations are operation-level (commutative/associative/distributive), as expression-level equivalences require specific operands/constants and are hard to hold, confirming Table 3 covers bulk redundancy.

1.3) Residuals and Future Work

EG captures equivalences in Table 3, but some rarer expression-level ones (e.g., advanced trig identities) remain. Future work includes: expanding Table 3 with semantic rules; integrating pre-trained GNNs for equivalence detection or hybrid algebraic rewriters. We will discuss residuals in camera-ready.

2) Empirical Limitation

We agree that broader validation would strengthen generalizability. The 32-atom copper dataset [25] was chosen as a representative materials science benchmark, aligning with MPtraj (~31 atoms [40]). The analytical potential $f(r)$ (Table 2) generalizes to large systems, enabling million-atom simulations (Appendix F.2). Besides, the Feynman benchmark (Fig. 4d) already spans diverse physics domains (e.g., mechanics, electromagnetism, quantum), with GSR recovering more ground-truth expressions than baselines, implicitly validating constraint handling across domains. For future work, we plan cross-domain extensions like fluid dynamics (e.g., Navier-Stokes discovery) or biology (e.g., reaction kinetics), leveraging GSR's flexible constraint encoding. We will add this to the discussion in revisions.

We believe these clarifications address your concerns and reinforce the paper's contributions. Thank you again for your feedback—we will incorporate your suggestions to improve the final version.

Best regards,

The Authors

We sincerely thank the Reviewer RJ7g's again for the thoughtful review and insightful feedback. As the discussion period nears its conclusion, we would like to briefly follow up on your comment regarding the uniqueness of Expression Graphs (EGs), in case further clarification might be helpful.

EG normalization is modeled using a convergent term rewriting system (TRS) based on the equivalence rules listed in Table 3 (Appendix C.2). The TRS is both terminating and confluent, ensuring that mathematically equivalent expressions yield unique DAG representations up to isomorphism, including nested cases via recursive normalization.

This improved canonicalization is quantified by a 33.9% reduction in search space size (MCTS-EG vs. MCTS; Table 1). We acknowledge that certain rare or domain-specific equivalences remain outside current coverage, and we plan to extend the TRS and investigate GNN-based semantic normalization as part of future work (Sec. 6). We will detail this TRS proof and other clarifications in Appendix C.2 revisions

If you have any further questions or if clarification would be useful, we’d be happy to elaborate. Otherwise, we sincerely appreciate your insights, which have helped us identify areas for improvement in future revisions.

Best regards,

The Authors