# Optimizing Tensor Computation Graphs with Equality Saturation and Monte Carlo Tree Search

arXiv: https://arxiv.org/abs/2410.05534 · HTML: https://mlsanity.com/paper/2410.05534/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 37 references that are in the map.
- First posted on arXiv 7 Oct 2024.

## Builds on (most cited first, 16 of 16)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [Aggregated Residual Transformations for Deep Neural Networks](https://mlsanity.com/paper/1611.05431/) (2016)
- [Learning Transferable Architectures for Scalable Image Recognition](https://mlsanity.com/paper/1707.07012/) (2017)
- [SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size](https://mlsanity.com/paper/1602.07360/) (2016)
- [Transformer Transducer: A Streamable Speech Recognition Model with Transformer Encoders and RNN-T Loss](https://mlsanity.com/paper/2002.02562/) (2020)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs](https://mlsanity.com/paper/2210.09603/) (2022)
- [Automating Constraint-Aware Datapath Optimization using E-Graphs](https://mlsanity.com/paper/2303.01839/) (2023)
- [MCTS-GEB: Monte Carlo Tree Search is a Good E-graph Builder](https://mlsanity.com/paper/2303.04651/) (2023)
- [X-RLflow: Graph Reinforcement Learning for Neural Network Subgraphs Transformation](https://mlsanity.com/paper/2304.14698/) (2023)

## Built on by (most cited first, 1 of 1)

- [nncase: An End-to-End Compiler for Efficient LLM Deployment on Heterogeneous Storage Architectures](https://mlsanity.com/paper/2512.21571/) (2025)

## Abstract (arXiv)

The real-world effectiveness of deep neural networks often depends on their latency, thereby necessitating optimization techniques that can reduce a model's inference time while preserving its performance. One popular approach is to sequentially rewrite the input computation graph into an equivalent but faster one by replacing individual subgraphs. This approach gives rise to the so-called phase-ordering problem in which the application of one rewrite rule can eliminate the possibility to apply an even better one later on. Recent work has shown that equality saturation, a technique from compiler optimization, can mitigate this issue by first building an intermediate representation (IR) that efficiently stores multiple optimized versions of the input program before extracting the best solution in a second step. In practice, however, memory constraints prevent the IR from capturing all optimized versions and thus reintroduce the phase-ordering problem in the construction phase. In this paper, we present a tensor graph rewriting approach that uses Monte Carlo tree search to build superior IRs by identifying the most promising rewrite rules. We also introduce a novel extraction algorithm that can provide fast and accurate runtime estimates of tensor programs represented in an IR. Our approach improves the inference speedup of neural networks by up to 11% compared to existing methods.