SmoothE: Differentiable E-Graph Extraction

An e-graph packs every rewriting of a program into one compact structure, but pulling the best version back out is NP-hard — integer programming solvers take hours and sometimes never finish, greedy heuristics finish instantly and miss shared subexpressions, and both are stuck with cost models that are a plain weighted sum; SmoothE (Cornell and Maryland, and the ASPLOS 2025 best paper) turns each yes-or-no choice into a probability, rewrites the correctness constraints so they too are continuous — including acyclicity, expressed as the trace of a matrix exponential — and then runs the whole thing as gradient descent in PyTorch on a GPU, which both matches the solvers in seconds instead of minutes and admits any cost function you can differentiate, including a learned one.

Why this problem is worth solving

Every optimising compiler is really a search over equivalent programs, and the traditional way of conducting that search destroys information at every step.

But a structure holding every option is only useful if you can get the good one out.

What people did before, and where it broke

• Formulate extraction as integer linear programming. Most published e-graph work does this, encoding selection as binary variables, completeness as linear inequalities, and acyclicity through an auxiliary topological ordering. It is exact when the solver finishes. Where it broke. Scalability. The paper measures a commercial solver needing 211 to 678 seconds on realistic datasets under a fifteen-minute cap, with open-source solvers failing outright on several graphs, and notes that even the best solvers can take hours on medium-sized problems. All of it runs on CPUs, so none of the parallel hardware sitting in the machine helps.
• Run a greedy bottom-up heuristic. Initialise every class to infinite cost, walk up from the leaves with a worklist, and keep whichever node gives its class the lowest total. It is what the popular e-graph libraries ship, and it finishes in a fraction of a second. Where it broke. It adds up the cost of each operand independently, so when two branches would use the same subterm it counts that subterm twice and picks a worse program to avoid a cost it would not actually pay. The paper builds two adversarial datasets specifically rich in shared subexpressions, and the heuristic comes out 2 to 6 times worse than optimal on them.
• Encode the problem for a satisfiability solver instead. A published formulation maps extraction onto maximum satisfiability. Where it broke. It requires enumerating all cycles in the graph, and the number of cycles grows exponentially with the number of nodes.
• Prune the graph first so the hard constraint disappears. Some systems delete all cycles as a preprocessing step, which lets the acyclicity constraint be dropped and dramatically speeds up the solver. Where it broke. Pruning removes feasible solutions along with the awkward ones, so the answer you get is the best of a smaller space, and nothing tells you how much you gave up.
• Work around linear cost models from outside. Recent systems that need a non-linear objective generate a pool of candidate extractions and rank them with a learned model afterwards, or fall back on simulation to get the quantity they actually care about. Where it broke. Ranking a pool is not optimising — the extraction itself is still being driven by a linear proxy, so the good solutions have to happen to be in the pool.
• The gap. Nobody had asked whether the discreteness was essential. Every method took the binary in-or-out decision as given and then differed in how it searched over binary assignments, which meant the whole field inherited both the scalability wall and the restriction to objectives a solver can express.

How SmoothE turns extraction into a gradient descent problem

The idea

Stop searching over discrete assignments and start moving continuously through the space of them. The reframing is the entire contribution, and it is the sort that looks obvious afterwards. A binary variable saying whether a node is selected becomes a real number saying how likely it is to be selected; the cost becomes a function of those numbers; and a function of real numbers can be differentiated, which means the search direction is computed rather than guessed. Nothing about the underlying problem got easier — it is still NP-hard, and SmoothE offers no optimality guarantee — but the search now moves by following a gradient rather than by branching, and it moves in every coordinate simultaneously instead of one decision at a time. That is what makes it a GPU workload rather than a solver workload. Relaxing a discrete program-optimisation search into a smooth one so that gradient descent can replace enumeration is the same manoeuvre as the tensor-program work covered elsewhere on this site, which found the same leverage in a different corner of the compiler.

The constraints have to be relaxed too, and that is where the real work is. Making the objective differentiable is easy; making the rules of correctness differentiable is not. There are three, and SmoothE handles each differently. Exactly one node per class becomes a softmax, so it holds by construction and never needs enforcing. Choosing an operator implies choosing its operands becomes a probability propagated from the root outward — and this is where cycles bite, because with a cycle there is no order in which to do the propagation, so SmoothE abandons ordering altogether and iterates the update in parallel everywhere until the numbers stop moving, which happens to suit a GPU far better than a sequential sweep would have. The third constraint, no circular dependency in the result, is handled by importing a theorem from causal structure learning: the trace of the matrix exponential of an adjacency matrix counts closed walks of every length, so it hits a known value exactly when there are none. A hard combinatorial property becomes one smooth scalar in the loss. Where the region-based work on encrypted programs covered elsewhere on this site made a hard problem tractable by finding a decomposition where the pieces were exactly solvable, this takes the opposite route — keep the problem whole and change what kind of problem it is.

Once it is a tensor program, the objective can be anything you can differentiate. This is the consequence that matters most in practice. Every previous extraction method needed its cost to be a weighted sum, because that is what a linear solver can take, and real compiler costs are not weighted sums — how operations cluster during technology mapping, how switching activity determines circuit power, how a fused kernel performs. SmoothE never assumes linearity anywhere; the cost is just a node in the computation graph. So a small neural network can be the cost model, trained on whatever the real objective is, and gradient descent optimises through it exactly as it optimises through a weighted sum. The paper measures that switching from a linear cost to a learned one leaves the per-iteration time essentially unchanged, since the cost model is not what the workflow is built around. Making a learned model the objective of a compiler search rather than a filter applied afterwards is the same shift as the deep-learning cost model work for tensor program tuning covered elsewhere on this site.

Walking one e-graph through SmoothE

Handling the one-per-class rule this way rather than as a penalty term matters more than it looks. A penalty can be violated during optimisation and has to be weighted against the objective; a softmax cannot be violated, so the optimiser never spends effort on a region of the space that was never feasible.

Why it works

• The hard part of extraction is interaction, and gradients see interaction. Whether a node is worth choosing depends on what else was chosen, since shared subterms are paid for once → a greedy method prices each branch in isolation and therefore double-counts → optimising all the probabilities jointly means every choice is evaluated in the presence of every other one.
• Constraints that hold by construction cost nothing to maintain. The one-node-per-class rule is expressed as a softmax rather than a penalty → the optimiser cannot leave the feasible region for that constraint and never spends gradient steps returning to it → only the genuinely awkward constraint, acyclicity, needs a term in the loss.
• Abandoning sequential order is what unlocks the hardware. Propagating probabilities in topological order is impossible on a cyclic graph anyway → so the update is applied everywhere in parallel and iterated to convergence → the necessity and the GPU-friendliness turn out to be the same thing.
• A combinatorial property becomes one smooth scalar. The trace of a matrix exponential counts closed walks of every length at once → acyclicity is the statement that this quantity takes a particular value → an exponentially large constraint enters the loss as a single differentiable number.
• Restarts are free when they are a batch dimension. Gradient descent on a non-convex landscape needs multiple starting points, and a sparse irregular graph cannot fill a GPU alone → running many seeds as one batched tensor program addresses both at once → quality and utilisation improve together rather than trading off.

What it costs (trade-offs and scope)

• On the adversarial datasets, the solvers win outright and it is not close. The paper builds two datasets from NP-hard problems converted into extraction, precisely to expose the heuristics — and the heuristics are indeed 2 to 6 times worse than optimal. But SmoothE lands 21.9% and 22.9% above optimal on average there, while the integer programming solvers find the optimum in between 0.2 and 33 seconds. The paper explains why: those graphs carry less graphical structure and are therefore more suitable for solvers. It is a real limit, reported without softening.
• The greedy heuristics are still far faster, and sometimes just as good. On the first dataset the heuristic takes 0.2 to 0.3 seconds and lands within 0.1% of optimal on average — the same quality SmoothE reaches in 8.4 seconds. The clear wins over the heuristic are on the two datasets where shared subexpressions matter, where it goes from 53.0% to 0.0% and from 2.9% to 0.2%.
• The non-linear cost model is synthetic and deliberately adversarial. The systems that motivate non-linear costs are not open-sourced, so the paper trains a four-layer perceptron on random valid solutions with random negative target values, and optimises against that. It demonstrates that a differentiable objective can be optimised, which is the claim. It does not demonstrate that optimising a real non-linear cost produces better programs.
• A modelling assumption is chosen per dataset by hand. Computing class probabilities requires assuming something about how sibling selections relate, and the paper uses the independent assumption on three datasets and a correlated one on two others. Which to use is a per-dataset decision made by the authors, not something the method determines.
• There is no optimality guarantee and no bound. The relaxation is not exact, gradient descent on a non-convex objective finds local minima, the acyclicity constraint is a penalty rather than a hard requirement, and the sampling step does not check for cycles. What the paper offers is a measured trade-off, not a proof.
• It needs a large GPU to be at its best. On an 11 GB card rather than an 80 GB one the seed batch shrinks by a factor of eight. Quality is reported as largely unchanged and sometimes better, but four graphs could not run at all because a single seed exceeded memory.
• The performance depends on an implementation detail more than it looks. Ablation shows GPU execution and the matrix exponential optimisation each contributing roughly a factor of ten, and without the latter SmoothE runs out of GPU memory on large graphs entirely. The algorithmic idea and the engineering that makes it viable are not separable here.
• The comparison against a genetic algorithm is acknowledged as unfair. The authors implemented one for the non-linear comparison, note it is pure Python, and exclude it from the main table for that reason as well as because it performed worse.

The evaluation

Check yourself

Sources

The paper is “SmoothE: Differentiable E-Graph Extraction” by Yaohui Cai, Kaixin Yang, Chenhui Deng, Cunxi Yu and Zhiru Zhang, published at ASPLOS 2025 Volume 1, pages 1020-1034, DOI 10.1145/3669940.3707262. The authors are at Cornell University and the University of Maryland, College Park. It received the ASPLOS 2025 Best Paper Award, and the implementation is open-sourced.

This explainer is written from the full text, read from the author-hosted copy at Cornell. Everything above comes from the paper’s own sections 2 through 5: the statement of the three extraction constraints and the integer programming formulation that encodes them; the description of the greedy worklist heuristic and the worked example showing it missing a shared subexpression; the survey of satisfiability formulations requiring cycle enumeration and of pruning approaches that shrink the feasible space; the relaxation of selection variables to probabilities and of the objective to a differentiable function; the softmax construction for conditional probabilities within a class and the propagation rules with their independent, correlated and hybrid variants; the parallel schedule for cyclic graphs; the acyclicity theorem based on the trace of a matrix exponential and its Pade implementation; the greedy sampling schedule and the note that it guarantees only the completeness constraints; the seed batching design; the dataset table covering DSP vectorisation, CGRA mapping, FPGA synthesis, datapath and tensor graphs plus the two converted NP-hard problems; the linear cost model comparison with its per-dataset times and percentage gaps against a ten-hour oracle; the adversarial results where the heuristics are 2.0 to 6.3 times off and SmoothE 21.9% and 22.9% off; the non-linear comparison using a four-layer perceptron trained on random targets; the ablation attributing roughly a factor of ten each to GPU execution and to the matrix exponential optimisation; the profiling showing sampling at 4.8% to 21.8% of runtime; and the portability study on an 11 GB card with its eight-fold batch reduction and four failures.