Equality Saturation Guided by Large Language Models

TL;DR

This work tackles the correctness gaps of LLM-driven optimizations by introducing LGuess, which inserts $e$-graphs as an intermediate layer between LLMs and formal rewrite systems to produce sound rewrite chains via high-level checkpoints. The approach uses a phase-based workflow where each phase saturates an $e$-graph, then leverages an LLM to identify a checkpoint for the next phase, addressing the intractability of direct low-level guidance. A probabilistic bigram model derived from the LLM guides checkpoint extraction, enabling efficient search within the saturated graph and progressive refinement toward convergence. Empirically, LGuess on multivariable polynomial factorization outperforms direct equality saturation and direct LLM approaches, solving substantially more tasks and showing particular strength on harder instances, which suggests practical impact for boosting correctness-aware optimization in structured rewrite settings.

Abstract

One critical issue with large language models (LLMs) is their inability to guarantee correctness. Although this problem can be addressed by applying LLMs to formal rewrite systems, current LLMs are still far from adequate to generate sound rewrite chains. To bridge this gap, this paper proposes LLM-guided equality saturation, dubbed LGuess, by incorporating e-graphs as an intermediate layer between LLMs and rewrite systems. LGuess queries LLMs only for high-level rewrite checkpoints and uses e-graphs to supply low-level rewrite chains between these checkpoints. The key technical challenge in this procedure lies in effectively extracting a suitable checkpoint from a saturated e-graph, which LGuess addresses by learning a probabilistic model from the LLM. The model predicts probable checkpoints while remaining simple enough for effective extraction. We implement a prototype of LGuess and evaluate it on the problem of factorizing multivariable polynomials. The results demonstrate a significant advantage of LGuess compared to both straightforward equality saturation and the approach that queries the LLM directly for the rewrite chain.

Figures (6)

• Figure 1: The workflow of LGuess.
• Figure 2: Rewrite rules in a ring with char. 2.
• Figure 3: The outline of the simplification.
• Figure 4: The saturated e-graph in the first phase, where $\textit{sqr}$ denotes the square operator $(\cdot)^2$, and $*$ denotes the multiplication operator.
• Figure 5: A query to GPT-4o and the response.