RAG-SR: Retrieval-Augmented Generation for Neural Symbolic Regression

Q1.

Thank you for your suggestion. To improve clarity, we have replaced "pretraining-free supervised learning" with "supervised learning methods that do not rely on pretraining" to emphasize that RAG-SR does not require pretraining.

Q2./W.

Thank you for suggesting these relevant papers. We have included discussions on VQ-VAE for code generation, RAG for code generation, and LLMs for symbolic regression in Section N of the supplementary material, highlighting their relevance and differences from our work. Regarding the differences between VQ-VAE and RAG-NN, there are two key distinctions:

• Interpretable Retrieval Results: RAG-SR retrieves symbolic expressions from a library, offering human-interpretable solutions, whereas VQ-VAE learns high-level concepts through a discrete latent representation.

• Access to Latest Information Without Retraining: RAG-SR can utilize the latest information even without training the neural network on new data. Specifically, new symbolic expressions can be added to the retrieval library to provide up-to-date information. The neural network does not need to be retrained as long as validation performance does not degrade.

While RAG and VQ-VAE are different, they can be complementary and combined to enhance code generation from multiple perspectives, which could be explored as a potential direction for future research.

Q3.

The maximum depth of GP trees and the number of trees are both fixed at 10, consistent with common settings in the GP literature. The algorithm maintains these fixed numbers throughout the evolutionary process.

To investigate scalability, we have conducted an analysis of the impact of the number of GP trees, included in Section M.2 of the supplementary material. Results indicate that setting the number of trees to 10 strikes a good balance between accuracy, model size, and training time, making it a suitable default choice.

Q4.

The residual is a vector of real numbers. To minimize computational overhead, we pre-compute and store the semantics of all subtrees for each solution during the evaluation phase. This allows us to avoid re-evaluating the semantics of the GP trees during the replacement process.

Specifically, if the semantics of a solution is represented by $\hat{Y} = \beta_1 \phi_1(X) + \dots + \beta_m \phi_m(X)$, the updated semantics after removing tree $k$ can be computed as $\hat{Y} - \beta_k \phi_k(X)$, and the residual is given by $Y - (\hat{Y} - \beta_k \phi_k(X))$. When replacing a tree, the process uses pre-computed semantics $\phi_1(X), \dots, \phi_m(X)$, ensuring efficiency. For a tree $\sigma$ retrieved from the semantics library, the new error, $[Y - (\hat{Y} - \beta_k \phi_k(X)+\beta_\sigma \sigma(X))]^2$, can be efficiently calculated using the stored semantics of $\sigma(X)$ from the library. For a tree $\sigma$ generated by the neural network, we assume the neural network generates good trees and skip error evaluation to reduce computational overhead.

Q5.

Semantics are defined as the output of applying a function $\phi$ to the input data $X$. For a GP tree $\phi$, the semantics are represented as $\phi(X)$, denoting the result of executing the GP tree on $X$. For a solution $\Phi$, semantics are defined as the output of the entire model, which includes both constructed features and the linear model, formally given by $\Phi(X) = \beta_1 \phi_1(X) + \dots + \beta_m \phi_m(X)$. In both cases, semantics are represented as a vector where each element corresponds to the output for a specific training instance.

The semantics determine the solution's performance based on a metric such as mean squared error. We have revised the paper to make it clear that semantics refer to the output of a model $\Phi$, a tree $\phi$, or a subtree $\psi$ on the input data $X$.

Q6.

The probability $P_{\text{neural}}$, which controls the balance between generating new expressions using the neural network and retrieving symbolic expressions from the semantic library, is a hyperparameter.

In this work, $P_{\text{neural}}$ is set to 0.1, guided by the rationale that one tree in each solution is replaced by a neural network-generated tree, as there are ten trees in each solution. An analysis of the influence of $P_{\text{neural}}$ has been added to Section M.1 of the supplementary material. Empirical results confirm that $P_{\text{neural}} = 0.1$ achieves an effective balance between neural generation and retrieval, ensuring effectiveness while minimizing the computational burden of neural generation.

Q7.

We have added error bars to all figures related to runtime results. After analyzing these error bars, we confirm that our conclusions regarding runtime performance remain unchanged.

Q8.

Although neural networks are limited to training on a small number of function nodes, the height of GP trees is capped at 10 to allow for more complex ground-truth equations. In this scenario, the role of the neural network is to provide useful building blocks, allowing the evolutionary algorithm to refine them into the final expressions.

To investigate what the neural network generates when the ground truth exceeds the limits of neural generation (i.e., the failure modes), we have added an analysis in Section O of the supplementary material. Experimental results for neural generation in cases where the ground-truth tree height is 6 have been included, including the generated expressions and the edit distance to the ground truth. These results show that, even when the ground truth exceeds the limits of neural generation, the neural network contributes valuable subexpressions that aid the evolutionary search process.

Q1./W3.

Thank you for your comments. We have included a hyperparameter sensitivity analysis in Sections I and M of the supplementary material in the revised version.

• Analysis on the Probability of Neural Generation $P_{\text{neural}}$: In this work, $P_{\text{neural}}$ is set to 0.1, guided by the rationale that one tree in each solution is replaced by a neural network-generated tree, as there are ten trees in each solution. An analysis of the influence of $P_{\text{neural}}$ has been added to Section M.1 of the supplementary material. Empirical results confirm that $P_{\text{neural}} = 0.1$ is an effective choice, providing a good trade-off between neural generation and retrieval from the knowledge base. This setting ensures effective use of neural generation while minimizing its computational burden.

• Analysis on the Weight of the Contrastive Loss $\lambda$: As for the weight of the contrastive loss $\lambda$, five different values ranging from 0.01 to 0.2 are evaluated to study their effect on accuracy and runtime, as detailed in Section I of the supplementary material. The results indicate that runtime is not significantly affected by $\lambda$, while accuracy is the primary factor that should be considered when tuning $\lambda$.

Q2./W1.

To improve scalability on large datasets, RAG-SR employs instance selection every 10 generations, retaining the semantics of the top 50 most challenging instances based on loss rankings. This ensures that the time consumption for the retrieval and neural generation components of RAG-SR remains consistent regardless of dataset size, making the algorithm highly scalable. Specifically, this strategy provides two key benefits for large datasets:

• Reducing Query Complexity: This approach reduces the query complexity of the KD-Tree during retrieval, with a worst-case time complexity of $O\left(k \cdot n^{1 - \frac{1}{k}}\right)$ [2] for $k$-dimensional semantics (i.e., $k$ instances).

• Reducing Training Complexity of the Neural Network: Fixing the dimensionality of semantics via instance selection ensures that the neural network input size remains constant, thereby facilitating scalability to large datasets.

Further details are included in Section A of the supplementary material.

Q3./W2.

Data augmentation and double querying are not limited to SR tasks. For example, in RAG-based code generation, text queries can be augmented by paraphrasing using appropriate tools. A broader discussion on this applicability has been added to Section N.1 of the supplementary material.

Q4.

We agree that evaluating the robustness of language model-based SR methods against varied functional distributions and function overlap is important. In this work, RAG-SR has not been pre-trained on external datasets. Consequently, both the black-box experiments and experiments on the Feynman and Strogatz datasets represent unseen problems for RAG-SR, demonstrating its ability to generalize to novel tasks.

Q5./W4.

In RAG-SR, interpretability is measured by the number of nodes in the symbolic tree, a commonly used metric for estimating model complexity in domains such as physics [1]. RAG-SR incorporates two mechanisms to balance interpretability and accuracy:

• Complexity Control in Exact Retrieval: First, during exact retrieval and replacement, the top-10 nearest trees are retrieved, and the nearest tree that is smaller than or equal in size to the incumbent GP tree is selected for replacement, rather than directly using the nearest tree.

• Complexity Control in Neural Generation: Second, the neural network is constrained to generate trees within a limited size.

These measures ensure a good balance between simplicity and accuracy.

[1]. Tenachi, W., Ibata, R., & Diakogiannis, F. I. (2023). Deep symbolic regression for physics guided by units constraints: toward the automated discovery of physical laws. The Astrophysical Journal, 959(2), 99.

[2]. Lee, D. T., & Wong, C. K. (1977). Worst-case analysis for region and partial region searches in multidimensional binary search trees and balanced quad trees. Acta Informatica, 9(1), 23-29.

I have read the rebuttal and willing to keep my score.

Q3.

• High Efficiency in Retrieval: Regarding the runtime performance of RAG-SR compared to GP-only models, the baseline algorithm, SBP-GP [1], while not incorporating neural networks, can still exhibit high runtime because it enumerates a large set of GP trees in an external library to determine the best GP trees for replacement, resulting in a time complexity of $O(L)$, where $L$ is the number of GP trees in the library. In contrast, RAG-SR uses a KD-Tree for efficient retrieval, reducing the time complexity to $O(\log L)$.

• High Efficiency in Neural Network Training: Additionally, RAG-SR improves efficiency by monitoring performance on a validation set, skipping neural network training during certain generations, incrementally training the neural network, and employing early stopping. The KD-Tree is reconstructed at each generation, irrespective of whether training is skipped. Thanks to the RAG mechanism, the neural network accesses information from the latest trees in the evolution process via retrieval from the library, even without further training. These optimizations collectively enable RAG-SR to achieve better runtime efficiency.

Q4.

Regarding the balance between accuracy and complexity achieved by RAG-SR, it arises from two size constraint mechanisms in our framework:

• Complexity Control in Exact Retrieval: First, during exact retrieval and replacement, as outlined in Algorithm 3, the KD-Tree selects trees based on their distance to the residual, and a size constraint ensures that only trees equal to or smaller than the current tree are considered, promoting simplicity.

• Complexity Control in Neural Generation: Second, the neural network is constrained to generate trees within a predefined size limit, further contributing to concise expressions.

[1]. Virgolin, M., Alderliesten, T., & Bosman, P. A. (2019, July). Linear scaling with and within semantic backpropagation-based genetic programming for symbolic regression. In Proceedings of the Genetic and Evolutionary Computation Conference (pp. 1084-1092).

Q1.

Thank you for your feedback. Figure 1 provides an overview of the key components of our approach, illustrating the interactions among them and emphasizing the role of the neural network in generating a new tree that is highly correlated with the residual. To improve clarity, we have reorganized the flow of Figure 1 and expanded it with more descriptive labels and details.

• Semantic Library Update Procedure: We have expanded the explanation in Section B of the supplementary material to describe the details of how the semantic library is populated, including the update procedure, the criteria for storing new entries/trees, and the data structure used for efficient retrieval. Specifically, the library is updated at the end of each generation during the evolutionary process to maintain a set of GP tree semantics (stored in a queue $S$) and their corresponding GP trees (stored in a queue $T$), using a First-In-First-Out queue mechanism to store the most recent entries. A hash set ($H_{\text{seen}}$) tracks previously seen semantics to prevent duplicate entries.

• Retrieval Process in Exact Replacement and Neural Generation: To provide further clarity on the retrieval process, we have included pseudocode and additional explanations in Section C of the supplementary material. These describe how residuals are matched to trees in the semantic library and how the most appropriate tree is identified in exact replacement. Specifically, the retrieval algorithm retrieves the top $\kappa = 10$ symbolic trees with the smallest distances to the residual $R$ in the retrieval library. If a tree that is smaller than or equal to the incumbent GP tree is found, it will be selected as the proposed tree ($\phi_{\text{new}}$). Otherwise, the current round of retrieval will be skipped.

  For the retrieval process in neural generation, it relies on a KD-Tree built based on all symbolic expressions with a maximum of $n_T$ function nodes. The KD-Tree enables efficient querying of the nearest semantics given the desired semantics $R$. The index of the nearest semantics will then be used to retrieve the corresponding symbolic expression from the library.

Q2.

For additional clarity regarding the training of the neural network, we have provided the following updates and clarifications to address these concerns.

• Use of Cross-Entropy Loss: We have provided an explanation in Section E of the supplementary material to clarify the definition of the ground truth labels used in the cross-entropy loss. Specifically, the ground truth labels include functions such as $+$, $*$, and terminals like $x_1$, $x_2$. A list of these functions and terminals corresponds to the GP tree $\phi$ that achieves the semantics $\phi(X)$ and is taken from the retrieval library. These ground truth labels are encoded as sequences in a pre-order traversal format of symbolic trees. For instance, a sequence like [mul, abs, add, $x_1$, $x_2$, $x_3$] corresponds to the mathematical expression mul(abs($x_1$), add($x_2$, $x_3$)). During training, the neural network predicts sequences similar to the pre-order traversal of the symbolic tree. The cross-entropy loss compares the predicted sequence with the ground truth sequence, token by token. For each position in the sequence, the predicted probability distribution over all possible tokens (output by the neural network) is compared to the one-hot encoding of the ground truth token for that position.

• Training Procedure of RAG-NN: For the neural network training procedure, a detailed explanation and pseudocode are provided in Section F of the supplementary material for further clarification. We explain how the neural network continuously updates its weights during the evolutionary process. Specifically, the following processes are performed sequentially:

  • Data Preparation: First, based on pairs of GP trees and semantics from the semantic library, the data preparation stage includes splitting the data into training and validation sets, augmenting the data, and using a KD-Tree for efficient retrieval of symbolic trees.

  • Semantic Library Update: Next, before training the neural network, a validation check is performed. If the validation loss does not degrade, only the retrieval library is updated to allow the neural network to access the latest information without further training.

  • Neural Network Update: Finally, if the validation loss degrades, additional training is performed. The neural network is trained iteratively using the Adam optimizer with cosine annealing and early stopping based on validation loss.

This procedure ensures that the neural network adapts dynamically to changes in the evolutionary process and efficiently learns from the evolving data.