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.

翻译：张量图超优化系统通过对神经网络执行一系列子图替换，以寻找最优计算图结构。这种图变换过程天然属于序列化决策框架，而现有系统通常采用贪婪搜索方法，由于无法容忍性能的暂时性损失，因而不能充分探索整个搜索空间。本文通过探索强化学习这一替代搜索方法，解决了张量图超优化问题。我们提出的X-RLflow方法能够学习执行神经网络数据流图重写，即每次替换一个子图。X-RLflow基于无模型强化学习智能体，利用图神经网络对目标计算图进行编码，并迭代输出变换后的计算图。实验表明，本方法在多种深度学习模型上优于现有最先进的超优化系统，对于基于Transformer架构的模型，性能提升最高可达40%。