Table of Contents
Fetching ...

E-Gen: Leveraging E-Graphs to Improve Continuous Representations of Symbolic Expressions

Hongbo Zheng, Suyuan Wang, Neeraj Gangwar, Nickvash Kani

TL;DR

E-Gen introduces an e-graph–based corpus generation framework to produce large, diverse clusters of semantically equivalent mathematical expressions. By training embeddings with seq2seq and contrastive learning on this cluster-based data, the approach achieves strong performance on clustering, semantic understanding beyond syntax, mistake detection, and embedding algebra, often outperforming prior SymPy-based methods and GPT-4o. The results demonstrate that rich, transformation-aware representations can be learned from synthetic, structure-focused data, with significant implications for mathematical information retrieval and reasoning. The work highlights the importance of operator diversity and semantic clustering in embedding symbolic mathematics for robust downstream applications.

Abstract

Vector representations have been pivotal in advancing natural language processing (NLP), with prior research focusing on embedding techniques for mathematical expressions using mathematically equivalent formulations. While effective, these approaches are constrained by the size and diversity of training data. In this work, we address these limitations by introducing E-Gen, a novel e-graph-based dataset generation scheme that synthesizes large and diverse mathematical expression datasets, surpassing prior methods in size and operator variety. Leveraging this dataset, we train embedding models using two strategies: (1) generating mathematically equivalent expressions, and (2) contrastive learning to explicitly group equivalent expressions. We evaluate these embeddings on both in-distribution and out-of-distribution mathematical language processing tasks, comparing them against prior methods. Finally, we demonstrate that our embedding-based approach outperforms state-of-the-art large language models (LLMs) on several tasks, underscoring the necessity of optimizing embedding methods for the mathematical data modality. The source code and datasets are available at https://github.com/MLPgroup/E-Gen.

E-Gen: Leveraging E-Graphs to Improve Continuous Representations of Symbolic Expressions

TL;DR

E-Gen introduces an e-graph–based corpus generation framework to produce large, diverse clusters of semantically equivalent mathematical expressions. By training embeddings with seq2seq and contrastive learning on this cluster-based data, the approach achieves strong performance on clustering, semantic understanding beyond syntax, mistake detection, and embedding algebra, often outperforming prior SymPy-based methods and GPT-4o. The results demonstrate that rich, transformation-aware representations can be learned from synthetic, structure-focused data, with significant implications for mathematical information retrieval and reasoning. The work highlights the importance of operator diversity and semantic clustering in embedding symbolic mathematics for robust downstream applications.

Abstract

Vector representations have been pivotal in advancing natural language processing (NLP), with prior research focusing on embedding techniques for mathematical expressions using mathematically equivalent formulations. While effective, these approaches are constrained by the size and diversity of training data. In this work, we address these limitations by introducing E-Gen, a novel e-graph-based dataset generation scheme that synthesizes large and diverse mathematical expression datasets, surpassing prior methods in size and operator variety. Leveraging this dataset, we train embedding models using two strategies: (1) generating mathematically equivalent expressions, and (2) contrastive learning to explicitly group equivalent expressions. We evaluate these embeddings on both in-distribution and out-of-distribution mathematical language processing tasks, comparing them against prior methods. Finally, we demonstrate that our embedding-based approach outperforms state-of-the-art large language models (LLMs) on several tasks, underscoring the necessity of optimizing embedding methods for the mathematical data modality. The source code and datasets are available at https://github.com/MLPgroup/E-Gen.
Paper Structure (29 sections, 5 equations, 5 figures, 17 tables, 1 algorithm)

This paper contains 29 sections, 5 equations, 5 figures, 17 tables, 1 algorithm.

Figures (5)

  • Figure 1: Illustration of e-graph saturation (\ref{['fig:init_e-graph']} to \ref{['fig:e-graph-3']}) and grammar creation (\ref{['fig:grammar']}). An e-graph consists of e-classes (dashed boxes) containing equivalent e-nodes (solid boxes). Edges connect e-nodes to their child e-classes. Applying mathematical rules to an e-graph adds new e-nodes and edges (\ref{['fig:e-graph-1']} and \ref{['fig:e-graph-2']}), or merges e-classes (\ref{['fig:e-graph-3']}). Additions and modifications are emphasized in black. In \ref{['fig:grammar']}, the saturated e-graph is converted into a context-free grammar, where each grammar is defined using e-class id and the e-nodes with their child e-classes.
  • Figure 2: Overview of E-Gen and the cluster-based training framework. Equivalent expressions of each initial expression are generated using an e-graph-based approach, forming clusters of equivalent expressions. The seq2seq model is trained on equivalent expression pairs, while the contrastive learning model is trained on triplets, with each triplet containing a reference expression $\boldsymbol{x}$, an equivalent positive sample $\boldsymbol{x}^{+}$, and a non-equivalent negative sample $\boldsymbol{x}^{-}$. $\boldsymbol{z}$ is the latent space representation of the corresponding input $\boldsymbol{x}$.
  • Figure 3: Visualization of representation vectors for 17 different single-operator mathematical expressions and their equivalent forms using our method (left) and SemEmb (right). t-SNE van2008visualizing is applied to reduce the dimensionality of the embeddings from 512 to 2.
  • Figure 4: Example of the formula retrieval with seq2seq model. The query expression is $\cos(x)$ in darker blue, and the rest are candidates. Top-$4$ expressions are retrieved. t-SNE is applied to reduce the dimensionality of the embeddings from 512 to 2.
  • Figure 5: Two examples of mistake detection conversations in the GPT-4o prompt. Errors in the derivation are highlighted in red. In the left example, GPT-4o correctly identifies the erroneous step but also incorrectly flags a false positive. In the right example, GPT-4o fails to detect the mistake entirely.