Bridging the Domain Gap in Equation Distillation with Reinforcement Feedback

TL;DR

This paper introduces a reinforcement learning-based finetuning method to improve foundation models for data-to-equation tasks, enabling better domain adaptation and more accurate, meaningful equation generation from complex data.

Contribution

It proposes a novel reinforcement learning framework that directly optimizes equation generation models using numerical fitness rewards, addressing domain adaptation and semantic accuracy issues.

Findings

Enhanced equation accuracy and robustness on complex datasets

Improved domain adaptability of foundation models

Outperforms existing methods in equation generation tasks

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.