MOD-SR: Unifying Multimodal Learning and Direct Optimization with Gradient-Guided Diffusion Model for Symbolic Regression

Symbolic regression (SR) aims to discover interpretable mathematical expressions from observed data. While recent generative approaches have shown promise in treating SR as machine translation or multimodal learning tasks using NN methods, they suffer from uncontrollable generation process and training-evaluation misalignment. The training objectives (average cross-entropy loss on a token level across the distribution of historical data) differ from the evaluation metric (fitting error for every test data / complexity), necessitating extensive heuristic post-processing. On the other hand, direct optimization methods suffer from an exponential slowdown as the dimensionality increases, non-differentiability and local optima traps. We propose **MOD-SR**, unifying multimodal distribution learning during training with direct optimization at inference time. This is achieved by modeling the task as $p(x_0 \mid \mathcal{D}, y^*)$ and employing gradient-guided diffusion in embedding space, enhanced by contrastive learning and representation alignment. Furthermore, we introduce DFEX, a fixed-depth tree relaxation method that ensures differentiability for effective gradient guidance during inference. Experiments demonstrate that MOD-SR achieves superior performance on diverse benchmarks through a unified framework integrating distribution learning and optimization. Our code is available at [https://github.com/KROX777/MOD-SR](https://github.com/KROX777/MOD-SR).

Symbolic regression is the task of finding the most accurate and simple mathematical formula that explains a given dataset—much like a scientist discovering physical laws from experimental observations. Recently, artificial intelligence has been used to "translate" data directly into equations. However, these AI models often struggle because their training goals don't perfectly align with what we actually care about: discovering highly accurate and concise formulas for specific new problems. Conversely, traditional search methods that actively try to build and test formulas step-by-step become excessively slow and get stuck when dealing with complex data with many variables. In this paper, we propose a new AI system called MOD-SR that combines the best of both worlds. During its initial training stage, it broadly learns the general patterns and "language" of mathematical equations. Then, when faced with a new dataset, it uses a guided refinement process—similar to how AI generates high-quality images—to continuously tweak and directly optimize the formula until it fits the data perfectly. To make this possible, we treat the rigid structure of mathematical trees as a smooth, continuous space, allowing the AI to effectively "slide" towards the best equation instead of resorting to trial-and-error. Our experiments show that MOD-SR is significantly more successful at discovering the true mathematical laws across various challenging benchmarks.