Technical Summary: EditSR – Enhancing Neural Symbolic Regression via Edit-based Rectification

1. Problem Statement

Neural symbolic regression models, particularly those utilizing Transformer architectures and pretraining (e.g., NeSymReS), have improved inference efficiency by shifting the structural search burden to the pretraining phase. However, these models rely on one-pass autoregressive decoding, which is prone to error accumulation. In symbolic regression, where expressions are tree structures governed by strict arity constraints, an early error (e.g., generating a binary operator instead of a unary one) distorts the syntactic context for all subsequent tokens. This often leads to structurally incorrect expressions that cannot be self-corrected.

Existing post-hoc rectification strategies attempt to mitigate this but often rely on restarting global search (e.g., via Monte Carlo Tree Search) or performing latent-space optimization during inference. These approaches incur substantial computational overhead, effectively negating the efficiency gains originally achieved through pretraining. Furthermore, many existing methods remain susceptible to error accumulation because they often condition decisions on the entire history rather than the current state.

2. Methodology: The EditSR Framework

The authors propose EditSR, a two-layer framework designed to achieve efficient prediction and post-hoc rectification without restarting the global search.

2.1. Architecture

• Layer 1 (Neural Predictor): A standard neural symbolic regression model (instantiated as NeSymReS) that maps datasets directly to expression skeletons via one-pass autoregressive decoding.
• Layer 2 (The Rectifier): A module that performs multi-step, edit-based rectification on incorrect predictions from Layer 1. It consists of two sub-components:
  • Tagger: Determines the optimal edit position and the type of edit action.
  • Editor: Predicts the content for the selected edit action.

2.2. Edit Actions and Syntax Constraints

The Rectifier operates on a subtree level rather than a token level to ensure syntactic validity. The defined edit actions are:

• Keep: Leaves the token/subtree unchanged.
• Replace: Replaces a token while preserving its arity.
• Delete: Collapses an internal subtree into a leaf.
• Rewrite: Regenerates an internal subtree.
• Insert: Expands a leaf into a new internal subtree.

To maintain syntactic validity, the framework enforces strict constraints during decoding. The "deficit count" (the number of pending child nodes required to complete the expression) is tracked. An edit is only permitted if the resulting state allows for a valid completion within the remaining token budget.

2.3. Training and State-Transition

The core innovation is formulating rectification as a step-by-step state-transition chain rather than a global search.

• Data Construction: Initial states $f^{(0)}$ are generated by applying random corruptions to target expressions $f^*$.
• Deterministic Chain Construction: A state-transition algorithm maps $f^{(0)}$ back to $f^*$ by determining the optimal sequence of edit actions. This algorithm aligns the current state with the target, calculates a cost function based on subtree differences, and selects the lowest-cost edit (position and action) to transition to the next state $f^{(t+1)}$.
• Training Objective: The Rectifier is trained on these supervised one-step transitions ($f^{(t)} \to f^{(t+1)}$). The Tagger learns to predict the position and action, while the Editor learns to generate the content.
• Inference: During inference, the Rectifier applies a greedy selection of edits based on the current state $f^{(t)}$. Because decisions are conditioned on the current state rather than the full history, errors made in earlier steps can be corrected by subsequent edits, mitigating error accumulation.

3. Key Contributions

1. Revisiting Error Accumulation: The paper identifies that incorrect predictions in neural symbolic regression often retain reusable structural components (variables, constants, subtrees) and proposes that these can be rectified via local edits rather than restarting the search.
2. EditSR Framework: A novel two-layer framework combining a neural predictor with a syntax-constrained, edit-based Rectifier. It preserves the inference efficiency of neural models while correcting structural errors without reopening global search.
3. Iterative State-Transition: The rectification process is modeled as an iterative state-transition chain where each step is conditioned on the current state. This design suppresses error accumulation and allows for the correction of earlier mistakes.
4. Empirical Validation: Extensive experiments demonstrate that the Rectifier significantly improves symbolic structure recovery, particularly for complex expressions where one-pass decoding fails, with limited additional computational cost.

4. Experimental Results

The authors evaluated EditSR on standard benchmarks (Constant, Koza, Nguyen, etc.), SRBench 1.0 (Feynman, ODE-Strogatz), and SRBench 2.0 (Black-box, Phenomenological).

• Performance: EditSR consistently achieves high Accuracy solution rates ($R^2 \ge 0.999$) and Symbolic solution rates (structural correctness). It outperforms or matches baselines like uDSR, SR4MDL, and TPSR, particularly on complex expressions.
• Complexity vs. Accuracy: Unlike some baselines that trade accuracy for complexity, EditSR maintains a favorable balance, generating simpler expressions while achieving high $R^2$.
• Efficiency: EditSR demonstrates the lowest test time among compared models (17.03s vs. 174.24s for TPSR on standard benchmarks). This confirms that shifting rectification effort to pretraining avoids the high runtime cost of online search-based refinement.
• Robustness:
  • Noise: EditSR maintains high performance under Gaussian noise levels (up to 0.1), outperforming many baselines in the Feynman and ODE-Strogatz benchmarks.
  • Distractors: In the presence of irrelevant variables, EditSR maintains a high symbolic solution rate while using significantly fewer distractors than baselines, indicating better structural recovery rather than just numerical fitting.
• Ablation Studies:
  • Rectifier Necessity: Adding the Rectifier significantly boosts performance, especially on high-complexity expressions where the base model (NeSymReS) often fails.
  • Fine-tuning: While pretraining on artificial chains provides a strong baseline, fine-tuning the Rectifier on the specific error patterns of the first-layer model yields further gains.
  • Edit Steps: Successful rectifications typically require only 4–6 edit steps, supporting the hypothesis that initial predictions are often structurally close to the target.

5. Significance and Claims

The paper claims that EditSR offers a practical solution to the trade-off between search capability and computational efficiency in neural symbolic regression. By reframing post-hoc rectification as a learned, edit-based state-transition process, the method:

• Preserves Efficiency: It avoids the computational penalty of restarting global search or performing expensive latent-space optimization during inference.
• Mitigates Error Accumulation: By conditioning edits on the current state, it allows for the correction of earlier errors, a feature lacking in standard autoregressive decoding.
• Improves Structural Recovery: The objective is directed toward the target symbolic structure, leading to higher rates of exact structural recovery compared to methods that only optimize for numerical error.

The authors acknowledge limitations, noting that the Rectifier's robustness relies on the quality of the initial prediction from the first layer and that extremely complex problems where the initial prediction is structurally distant from the target may still remain unresolved. However, the framework is presented as a significant step toward making neural symbolic regression both efficient and reliable for complex scientific discovery tasks.