[Paper Review] Equivalence of Dataflow Graphs via Rewrite Rules Using a Graph-to-Sequence Neural Model

Steve Kommrusch, Théo Barollet+1 more|arXiv (Cornell University)|Feb 17, 2020

Software Testing and Debugging Techniques50 references4 citations

TL;DR

This paper introduces pe-graph2seq, the first graph-to-sequence neural network for program equivalence, which generates provably correct rewrite sequences to transform one dataflow graph into another. Trained on a linear algebra language with 100+ equivalence axioms, it achieves 96% correctness on 10,000 test pairs with 16ms inference time, and all outputs are deterministically verifiable in negligible time.

ABSTRACT

In this work we target the problem of provably computing the equivalence between two programs represented as dataflow graphs. To this end, we formalize the problem of equivalence between two programs as finding a set of semantics-preserving rewrite rules from one into the other, such that after the rewrite the two programs are structurally identical, and therefore trivially equivalent. We then develop the first graph-to-sequence neural network system for program equivalence, trained to produce such rewrite sequences from a carefully crafted automatic example generation algorithm. We extensively evaluate our system on a rich multi-type linear algebra expression language, using arbitrary combinations of 100+ graph-rewriting axioms of equivalence. Our system outputs via inference a correct rewrite sequence for 96% of the 10,000 program pairs isolated for testing, using 30-term programs. And in all cases, the validity of the sequence produced and therefore the provable assertion of program equivalence is computable, in negligible time.

Motivation & Objective

• Address the challenge of provably checking program equivalence between two dataflow graphs using machine learning.
• Overcome the limitations of stochastic neural networks in program equivalence by generating deterministic, checkable proof sequences instead of binary predictions.
• Develop a neural system that outputs a sequence of semantics-preserving graph rewrite rules to make two equivalent programs structurally identical.
• Enable efficient and correct equivalence checking by ensuring all outputs are verifiable in negligible time, avoiding false positives.
• Demonstrate the approach on a complex, multi-type linear algebra expression language with rich algebraic rewrite axioms.

Proposed method

• Represent programs as directed, rooted dataflow graphs with nodes for operations and values, and edges for data dependencies.
• Formalize program equivalence as the existence of a sequence of semantics-preserving graph rewrite rules transforming one graph into the other.
• Define 100+ axioms of equivalence (e.g., 1*x → x, a*(b+c) → a*b + a*c) as rewrite rules applicable to subgraphs.
• Train a graph-to-sequence neural network (pe-graph2seq) to predict a sequence of such rewrite rules from input graphs.
• Use a custom automatic example generation pipeline to synthesize 10,000 training pairs with known correct rewrite sequences.
• Validate the output by applying the predicted rewrite sequence to one graph and checking structural identity via depth-first traversal.

Experimental results

Research questions

• RQ1 Can a neural network be trained to generate a sequence of graph rewrite rules that provably transform one dataflow graph into another equivalent graph?
• RQ2 Can such a system ensure that all non-equivalent inputs are correctly rejected, avoiding false positives?
• RQ3 Can the output of the neural network be deterministically verified in negligible time, enabling practical deployment?
• RQ4 How effective is the model in learning complex rewrite sequences involving subtree manipulation, node deletion, and reordering in a rich algebraic language?
• RQ5 To what extent can a graph-to-sequence model with attention mechanisms learn heuristics for navigating the space of equivalence-preserving transformations?

Key findings

• The pe-graph2seq model successfully generated a correct and verifiable rewrite sequence for 96% of 10,000 test program pairs in a multi-type linear algebra language.
• All outputs are deterministically checkable in negligible time by applying the rewrite sequence and comparing structural identity.
• The system never produces false positives: it correctly identifies non-equivalent programs as non-equivalent by design.
• The average inference time per program pair is 16ms, demonstrating high efficiency for practical use.
• The model generalizes across complex transformations, including subtree reordering, node deletion, and algebraic simplifications.
• The approach demonstrates that graph neural networks can effectively learn heuristics for navigating the combinatorial space of equivalence-preserving rewrites.