Identity Testing for Circuits with Exponentiation Gates

TL;DR

This paper introduces a new randomized identity testing algorithm for circuits with exponentiation gates, motivated by neural network compiler optimization, and demonstrates its effectiveness both theoretically and empirically.

Contribution

It formalizes a black-box model for circuits with exponential gates and provides the first efficient randomized identity testing algorithm for this class.

Findings

Algorithm achieves perfect completeness and non-trivial soundness.

Implemented in the Mirage compiler with promising empirical results.

Proposes a conjecture ensuring high-probability soundness.

Abstract

Motivated by practical applications in the design of optimization compilers for neural networks, we initiated the study of identity testing problems for arithmetic circuits augmented with \emph{exponentiation gates} that compute the real function $x\mapsto e^x$. These circuits compute real functions of form $P(\vec x)/P'(\vec x)$, where both $P(\vec x)$ and $P'(\vec x)$ are exponential polynomials \[ \sum_{i=1}^k f_i(\vec x)\cdot \exp\left(\frac{g_i(\vec x)}{h_i(\vec x)}\right), \] for polynomials $f_i(\vec x),g_i(\vec x)$, and $h_i(\vec x)$. We formalize a black-box query model over finite fields for this class of circuits, which is mathematical simple and reflects constraints faced by real-world neural network compilers. We proved that a simple and efficient randomized identity testing algorithm achieves perfect completeness and non-trivial soundness. Concurrent with our work, the algorithm has been implemented in the optimization compiler Mirage by Wu et al.~(OSDI 2025), demonstrating promising empirical performance in both efficiency and soundness error. Finally, we propose a number-theoretic conjecture under which our algorithm is sound with high probability.