[论文解读] Transformer-based Planning for Symbolic Regression

Parshin Shojaee, Kazem Meidani+2 more|arXiv (Cornell University)|Mar 13, 2023

Evolutionary Algorithms and Applications被引用 16

一句话总结

TPSR 将蒙特卡洛树搜索(MCTS)与预训练的变换器符号回归模型相结合，以优化方程生成以提高拟合精度和复杂度之间的权衡，优于基线方法。

ABSTRACT

Symbolic regression (SR) is a challenging task in machine learning that involves finding a mathematical expression for a function based on its values. Recent advancements in SR have demonstrated the effectiveness of pre-trained transformer-based models in generating equations as sequences, leveraging large-scale pre-training on synthetic datasets and offering notable advantages in terms of inference time over classical Genetic Programming (GP) methods. However, these models primarily rely on supervised pre-training goals borrowed from text generation and overlook equation discovery objectives like accuracy and complexity. To address this, we propose TPSR, a Transformer-based Planning strategy for Symbolic Regression that incorporates Monte Carlo Tree Search into the transformer decoding process. Unlike conventional decoding strategies, TPSR enables the integration of non-differentiable feedback, such as fitting accuracy and complexity, as external sources of knowledge into the transformer-based equation generation process. Extensive experiments on various datasets show that our approach outperforms state-of-the-art methods, enhancing the model's fitting-complexity trade-off, extrapolation abilities, and robustness to noise.

研究动机与目标

• 动机：将符号回归作为跨科学与工程的可解释控制方程的途径。
• 通过引入不可微分的反馈（拟合精度与复杂度），解决符号回归中令牌级预训练损失的局限性。
• 提出一种与模型无关的 TPSR 框架，通过将预训练的 SR Transformer 与 MCTS 规划结合，来优化方程发现。

提出的方法

• 将方程表示为前缀表示法，以便与变换器解码对齐。
• 使用预训练的 SR Transformer 主干来对数据进行编码并生成方程序列。
• 在推理过程中应用蒙特卡洛树搜索（MCTS）前瞻规划，基于平衡拟合度与复杂性的自定义奖励来引导标记选择。
• 结合前k个标记限制和基于束搜索的模拟来管理 MCTS 中的扩展与评估步骤。
• 定义一个奖励：将归一化均方误差（NMSE）为拟合度，与复杂度的指数惩罚相结合：r = 1/(1+NMSE) + λ exp(-l/L)。
• 引入缓存机制（前k缓存和序列缓存）以减少 MCTS 过程中的冗余计算。

Figure 1: Pareto plot comparing the rankings of all methods in terms of the $R^{2}$ performance and identified equation complexity for (a) SRBench Black-box datasets and (b) Feynman datasets . Our results with Transformer-based Planning (TPSR) applied on top of E2E transformer SR model improves its

实验结果

研究问题

• RQ1 在标准 SR 基准上，TPSR 是否优于束搜索和采样解码策略及其他基线？
• RQ2 TPSR 是否能在符号回归任务中提升外推能力和对噪声的鲁棒性？
• RQ3 TPSR 的缓存机制是否显著降低推理时间？
• RQ4 单独的 MCTS 组件对 TPSR 性能的影响是什么？

主要发现

• 与基线相比，TPSR 在 SRBench 和领域内合成数据集上在拟合精度更高且方程复杂度相当或更低。
• TPSR 在 Feynman 和 Black-box 数据集的拟合与复杂度权衡上达到第一帕累托前沿。
• 缓存机制总体降低推理时间约28%，其中序列缓存带来最显著的提升。
• 消融研究表明，增加 MCTS 的回合数、束宽和最大扩展通常会提升性能，其中回合数贡献最大。
• 与端到端 transformer 基线相比，TPSR 在尺度变化下的外推能力和对噪声的鲁棒性更佳。

Figure 2: An overview of our proposed method with MCTS-guided decoding at inference compared to the concurrent works with beam search/sampling decoding strategy.