# X-RLflow: Graph Reinforcement Learning for Neural Network Subgraphs Transformation

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 50 references that are in the map.
- First posted on arXiv 28 Apr 2023.

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

- [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)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [Zero-Shot Text-to-Image Generation](https://mlsanity.com/paper/2102.12092/) (2021)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size](https://mlsanity.com/paper/1602.07360/) (2016)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Solving Rubik's Cube with a Robot Hand](https://mlsanity.com/paper/1910.07113/) (2019)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Transformer Transducer: A Streamable Speech Recognition Model with Transformer Encoders and RNN-T Loss](https://mlsanity.com/paper/2002.02562/) (2020)
- [MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems](https://mlsanity.com/paper/1512.01274/) (2015)
- [Learning to Perform Local Rewriting for Combinatorial Optimization](https://mlsanity.com/paper/1810.00337/) (2018)
- [A Closer Look at Invalid Action Masking in Policy Gradient Algorithms](https://mlsanity.com/paper/2006.14171/) (2020)
- [Model-Based Value Estimation for Efficient Model-Free Reinforcement Learning](https://mlsanity.com/paper/1803.00101/) (2018)
- [Imagination-Augmented Agents for Deep Reinforcement Learning](https://mlsanity.com/paper/1707.06203/) (2017)
- [Transferable Graph Optimizers for ML Compilers](https://mlsanity.com/paper/2010.12438/) (2020)
- [Equality Saturation for Tensor Graph Superoptimization](https://mlsanity.com/paper/2101.01332/) (2021)
- [AutoPhase: Juggling HLS Phase Orderings in Random Forests with Deep Reinforcement Learning](https://mlsanity.com/paper/2003.00671/) (2020)
- [PrefixRL: Optimization of Parallel Prefix Circuits using Deep Reinforcement Learning](https://mlsanity.com/paper/2205.07000/) (2022)
- [Implementing Reinforcement Learning Datacenter Congestion Control in NVIDIA NICs](https://mlsanity.com/paper/2207.02295/) (2022)

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

- [A Reinforcement Learning Environment for Automatic Code Optimization in the MLIR Compiler](https://mlsanity.com/paper/2409.11068/) (2024)
- [Optimizing Tensor Computation Graphs with Equality Saturation and Monte Carlo Tree Search](https://mlsanity.com/paper/2410.05534/) (2024)

## Abstract (arXiv)

Tensor graph superoptimisation systems perform a sequence of subgraph substitution to neural networks, to find the optimal computation graph structure. Such a graph transformation process naturally falls into the framework of sequential decision-making, and existing systems typically employ a greedy search approach, which cannot explore the whole search space as it cannot tolerate a temporary loss of performance. In this paper, we address the tensor graph superoptimisation problem by exploring an alternative search approach, reinforcement learning (RL). Our proposed approach, X-RLflow, can learn to perform neural network dataflow graph rewriting, which substitutes a subgraph one at a time. X-RLflow is based on a model-free RL agent that uses a graph neural network (GNN) to encode the target computation graph and outputs a transformed computation graph iteratively. We show that our approach can outperform state-of-the-art superoptimisation systems over a range of deep learning models and achieve by up to 40% on those that are based on transformer-style architectures.