Neural-Guided Equation Discovery

TL;DR

The paper introduces MGMT, a modular grammar-guided Monte-Carlo Tree Search framework for equation discovery that supports both supervised and reinforcement learning. It evaluates how neural guidance, grammar-based action spaces, and dataset embeddings affect discovery quality, finding that supervised learning often outperforms reinforcement learning and grammars outperform token-based action spaces. It also introduces AmEx-MCTS and Risk-seeking MCTS to improve search efficiency and contrasts various tabular-data embeddings via contrastive learning to inform dataset representations. The study articulates seven desirable properties for equation-discovery systems, demonstrates continual learning across tasks, and highlights practical insights for scaling, parallelization, and handling noisy data in future work.

Abstract

Deep learning approaches are becoming increasingly attractive for equation discovery. We show the advantages and disadvantages of using neural-guided equation discovery by giving an overview of recent papers and the results of experiments using our modular equation discovery system MGMT ($\textbf{M}$ulti-Task $\textbf{G}$rammar-Guided $\textbf{M}$onte-Carlo $\textbf{T}$ree Search for Equation Discovery). The system uses neural-guided Monte-Carlo Tree Search (MCTS) and supports both supervised and reinforcement learning, with a search space defined by a context-free grammar. We summarize seven desirable properties of equation discovery systems, emphasizing the importance of embedding tabular data sets for such learning approaches. Using the modular structure of MGMT, we compare seven architectures (among them, RNNs, CNNs, and Transformers) for embedding tabular datasets on the auxiliary task of contrastive learning for tabular data sets on an equation discovery task. For almost all combinations of modules, supervised learning outperforms reinforcement learning. Moreover, our experiments indicate an advantage of using grammar rules as action space instead of tokens. Two adaptations of MCTS -- risk-seeking MCTS and AmEx-MCTS -- can improve equation discovery with that kind of search.

Figures (9)

• Figure 1: Overview MGMT A grammar and a data set are the input in MGMT. The search is performed by a neural-guided MCTS, which receives guidance from the neural network for each state in the search. Once a complete equation has been generated, all constants in the equation are fitted. The results of the MCTS are stored in a replay buffer and used to train the neural network. The aim is for the neural network and the MCTS to reinforce each other. The output of the system are the $k$-best equations found during the search.
• Figure 2: Interaction between components in MGMT. For each state in the search tree (see the right-hand side of the figure), the neural network gives guidance on which action to choose next ($P$) and how good the current state ($V$) is. Each state has a number of visits ($SSA$), which indicates how often its child nodes are visited. The results of the MCTS are denoted by a tilde and are stored in the replay buffer. When the neural network is trained, a batch is sampled from the replay buffer and the difference between the results of the MCTS and the current prediction of the NN is used to update the weights of the model.
• Figure 3: Sequential generation of an equation using the neural network from MGMT. Based on the visit counts of the MCTS, a probability distribution is calculated to sample which production rule from the grammar should be used to expand the syntax tree. The neural network is trained with the visit counts from the MCTS as well as the reward of the finished equation.
• Figure 4: Improving MCTS by ignoring already explored subtrees and leaves by focusing on the unknown. Updating the search strategy within MCTS by separating "incrementing visit counts" (displayed in blue) from the selected child nodes (displayed in green) to explore more while keeping the number of iterations $n_\mathit{sims}$ the same.
• Figure 5: Visualization of the contrastive loss between data sets for different equations.