Equality Saturation Guided by Large Language Models

TL;DR

This paper introduces LGuess, a novel method combining large language models and e-graphs to generate sound rewrite chains, improving correctness guarantees in formal rewrite systems.

Contribution

LGuess is the first approach to integrate LLMs with e-graphs for guided equality saturation, enabling more reliable rewrite chain generation.

Findings

LGuess outperforms straightforward equality saturation.

LGuess surpasses direct LLM querying in accuracy.

Prototype demonstrates effectiveness on polynomial factorization.

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.