计算图的神经拓扑排序
Neural Topological Ordering for Computation Graphs
摘要:
最近的工作机器学习的组合优化已经表明,基于学习的方法可以优于探试法在速度和性能方面。在本文中,我们考虑的问题,找到一个最佳的拓扑顺序上的有向无环图的内存最小化的问题,在编译器中出现的焦点。我们提出了一个端到端的机器学习为基础的方法,使用编码器-解码器框架的拓扑排序。我们的编码器是一种新的注意力为基础的图神经网络架构称为\emph{Topoformer},它使用不同的拓扑转换的DAG的消息传递。由编码器产生的节点嵌入被转换成节点优先级,其被解码器用于生成拓扑顺序上的概率分布。我们训练我们的模型的数据集的综合生成的图形称为分层图。我们表明,我们的模型优于,或与几个拓扑排序基线,同时显着更快的合成图与多达2k个节点。我们还训练和测试我们的模型上的一组现实世界的计算图,显示性能的改善。
英文摘要:
Recent works on machine learning for combinatorial optimization have shown that learning based approaches can outperform heuristic methods in terms of speed and performance. In this paper, we consider the problem of finding an optimal topological order on a directed acyclic graph with focus on the memory minimization problem which arises in compilers. We propose an end-to-end machine learning based approach for topological ordering using an encoder-decoder framework. Our encoder is a novel attention based graph neural network architecture called \emph{Topoformer} which uses different topological transforms of a DAG for message passing. The node embeddings produced by the encoder are converted into node priorities which are used by the decoder to generate a probability distribution over topological orders. We train our model on a dataset of synthetically generated graphs called layered graphs. We show that our model outperforms, or is on-par, with several topological ordering baselines while being significantly faster on synthetic graphs with up to 2k nodes. We also train and test our model on a set of real-world computation graphs, showing performance improvements.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。