Bridging the Domain Gap in Equation Distillation with Reinforcement Feedback

TL;DR

Data2Eqn seeks interpretable equations $y=f(\mathbf{X})$ from data, but foundation models face domain shifts and semantics misalignment. The authors propose REEL, a reinforcement learning–based finetuning framework that uses data-instance subset sampling and numerical-semantic rewards (based on $R^2$) to adapt generationPolicy while preserving prior knowledge via KL regularization. They introduce a two-loop, data-centric training paradigm with an endorsement-guided reward and a clipped surrogate objective to stabilize policy updates. Across Feynman, Strogatz, and Black-box benchmarks, REEL yields higher $R^2$ proportions above $0.99$, improved average $R^2$, and stronger robustness to noise, while offering faster inference than many baselines, demonstrating practical domain-adaptive symbolic regression.

Abstract

The data-to-equation (Data2Eqn) task aims to discover interpretable mathematical equations that map observed values to labels, offering physical insights and broad applicability across academic and industrial domains. Genetic programming and traditional deep learning-based approaches suffer from search inefficiency and poor generalization on small task-specific datasets. Foundation models showed promise in this area, but existing approaches suffer from: 1) They are pretrained on general-purpose data distributions, making them less effective for domain-specific tasks; and 2) their training objectives focus on token-level alignment, overlooking mathematical semantics, which can lead to inaccurate equations. To address these issues, we aim to enhance the domain adaptability of foundation models for Data2Eqn tasks. In this work, we propose a reinforcement learning-based finetuning framework that directly optimizes the generation policy of a pretrained model through reward signals derived from downstream numerical fitness. Our method allows the model to adapt to specific and complex data distributions and generate mathematically meaningful equations. Extensive experiments demonstrate that our approach improves both the accuracy and robustness of equation generation under complex distributions.

Figures (6)

• Figure 1: Framework Overview. Given a domain-specific dataset, we first sample diverse instance subsets to form multiple trajectory training sets. For each subset, an agent (a pretrained Data2Eqn model) iteratively generates equation tokens, receives numerical rewards based on equation-level fitness, and updates its policy via reinforcement learning. This process explores the optimal policy distribution and aligns generation with domain-specific mathematical semantics.
• Figure 2: Reward assignment: we reinforce it when the finetuned model generates high-quality equations.
• Figure 3: Performance comparison of REEL and all SRBench algorithms for Feynman and Strogatz datasets. For each dataset, the performance is sorted by the proportion of $R^2 > 0.99$. Our approach presents strong accuracy-speed tradeoffs.
• Figure 4: Robustness to noise, where the performances are shown for various target noise levels.
• Figure 5: Complexity and Scalability Analysis of REEL. (a) Comparison of training time on the Strogatz dataset between REEL and the baselines; (b) Comparison of model size between REEL and the baselines; (c) Analysis of how training time per epoch and prediction time change as the number of variables (features) in the dataset increases.