Equivalence of Dataflow Graphs via Rewrite Rules Using a Graph-to-Sequence Neural Model

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.

💡 Research Summary

The paper tackles the long‑standing problem of program equivalence by framing it as a graph‑rewriting task. Programs are represented as data‑flow graphs, where nodes correspond to operations (e.g., scalar, vector, matrix multiplication) and edges denote data dependencies. An “axiom of equivalence” is a semantics‑preserving rewrite rule that replaces a sub‑graph with an equivalent sub‑graph (for example, 1 × x = x, distributivity, commutativity). Two programs are equivalent if there exists a sequence of such rewrites that transforms one graph into the other; the sequence itself constitutes a constructive proof.

To automate the discovery of such rewrite sequences, the authors introduce pe‑graph2seq, the first graph‑to‑sequence neural architecture designed for program equivalence. The system consists of three main components:

1. Graph Encoder – A gated graph neural network (GNN) processes each input graph, producing node embeddings that capture operation type, data type (scalar, vector, matrix), and structural context. Global attention mechanisms allow the decoder to query any node during generation.

2. Sequence Decoder – A memory‑augmented LSTM generates a token stream describing the rewrite plan. Tokens encode the axiom identifier and the target node where the axiom should be applied. An end‑of‑sequence token signals completion.

3. Verification Engine – After inference, the generated rewrite sequence is applied step‑by‑step to the source graph. Each step checks whether the axiom’s pattern matches the current graph and whether the resulting graph equals the target. This verification runs in linear time with respect to graph size, guaranteeing that any reported equivalence is provably correct; false positives are impossible.

Training data are generated automatically. The authors define a rich linear‑algebra language with over 100 axioms covering scalars, vectors, and matrices. Random programs are created according to the grammar, then transformed by randomly applying axioms to produce a target program and the exact rewrite path used. Both equivalent (with a known path) and non‑equivalent (no path) pairs are collected, yielding a large supervised dataset without manual annotation.

Experiments evaluate pe‑graph2seq on 10,000 test pairs of 30‑operation programs. The model produces a correct rewrite sequence for 96 % of the equivalent pairs, with an average inference time of 16 ms per pair. Verification time is negligible (linear in graph size). Using a beam size of 10, the success rate remains above 95 %, whereas a naïve baseline that enumerates all learned sequences reaches only about 60 % accuracy, demonstrating the model’s ability to prioritize plausible paths.

Key contributions and insights include:

• Deterministic proof generation: By outputting an explicit rewrite sequence, the system eliminates stochastic uncertainty inherent in typical neural classifiers.
• Graph‑to‑sequence modeling: Combining GNN encoders with attention‑driven LSTM decoders enables the network to reason about structural changes and variable‑length rewrite paths.
• Scalable data synthesis: Automatic generation of massive labeled rewrite paths sidesteps the bottleneck of manual theorem‑proving data.
• Practical applicability: The approach can be used for automated grading of algebraic derivations, real‑time matching of high‑level expressions to optimized library calls (e.g., BLAS), and verification of compiler optimizations.

Limitations are acknowledged. The method depends on a predefined set of axioms; extending to new domains requires curating additional rewrite rules and retraining. False negatives can occur when a valid rewrite path exists but the model fails to discover it, leading only to a probabilistic “not equivalent” answer. Nonetheless, the authors argue that in many domain‑specific settings (linear algebra, domain‑specific languages) the axiom set is manageable and the benefits of provable equivalence outweigh occasional missed proofs.

In conclusion, the paper demonstrates that deep learning can be harnessed not merely for classification but for constructive reasoning tasks. By framing program equivalence as a path‑finding problem in a graph of rewriteable programs and training a graph‑to‑sequence model to propose such paths, the authors achieve high‑accuracy, real‑time, and formally verifiable equivalence checking. This opens avenues for integrating neural guidance into formal methods, automated theorem proving, and compiler technology.