[论文解读] Neural Topological Ordering for Computation Graphs

Mukul Gagrani, Corrado Rainone+8 more|arXiv (Cornell University)|Jul 13, 2022

Advanced Graph Neural Networks被引用 4

一句话总结

本文提出了一种新颖的端到端深度学习框架 Topoformer，用于在有向无环图（DAG）中学习最优拓扑顺序，以最小化编译器工作负载中的峰值内存使用量。该方法采用非自回归、基于注意力的图神经网络，并结合拓扑感知的消息传递机制，实现了显著快于基线方法的推理速度，同时达到最先进或具有竞争力的性能表现，包括相比某些基线方法快达4000倍的运行时间。

ABSTRACT

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.

研究动机与目标

• 解决在有向无环图（DAG）中寻找最优拓扑顺序以最小化编译器流水线中峰值内存使用量的组合优化问题。
• 开发一种可扩展的端到端深度学习框架，以学习到的策略替代启发式方法，用于拓扑顺序生成。
• 提出一种新颖的图神经网络架构 Topoformer，支持在 DAG 上进行拓扑感知的全局消息传递。
• 创建一个公开可用的分层、类似神经网络的计算图合成数据集，以支持机器学习组合优化中的可复现基准测试。
• 在合成图和真实计算图上均展示性能提升，同时保持快速的推理速度。

提出的方法

• 该方法采用非自回归（NAR）解码器，以高效参数化拓扑顺序的概率分布，从而实现快速推理。
• 编码器为 Topoformer，一种新颖的基于注意力的图神经网络，其通过基于节点对之间拓扑关系（如祖先、后代、无关）的不同可学习参数执行消息传递。
• Topoformer 通过七种不同的拓扑变换处理 DAG，包括前向和后向边传播，以捕捉全局结构依赖关系。
• 由 Topoformer 生成的节点嵌入被转换为节点优先级，以指导解码器生成拓扑序列。
• 模型在分层计算图的合成数据集上进行训练，这些图被设计为模拟真实神经网络工作负载，并支持可扩展的数据生成。
• 训练采用监督学习，以近似动态规划（DP）解作为目标，监督策略网络的学习。

实验结果

研究问题

• RQ1 可学习的、端到端的深度学习框架是否能在最小化 DAG 的峰值内存使用量方面优于传统启发式方法？
• RQ2 与局部或仅基于边的消息传递相比，图神经网络架构中拓扑感知的全局消息传递在 DAG 上的效果如何？
• RQ3 非自回归解码方案是否能以显著降低的推理时间实现与自回归基线相当的性能表现？
• RQ4 分层 DAG 的合成公开数据集在多大程度上能够支持机器学习组合优化中的可复现基准测试？
• RQ5 所提出的方法是否能泛化到真实世界的编译器工作负载，而不仅限于合成图？

主要发现

• 在 500 个节点的合成图上，完整版 Topoformer 模型在使用束搜索（beam search）时，与近似 DP 的平均 % 差距为 3.08 ± 0.51，优于 GAT（前向+后向）及其他基线模型。
• 在 1000 个节点的图上，使用束搜索的 Topoformer 实现了 0.05 ± 0.36 的 % 差距，显著优于 GAT（前向+后向）及其他模型。
• 与仅使用前向+后向的变体相比，完整版 Topoformer 将与近似 DP 的 % 差距降低了 1.5–2.5 个百分点，证明了全局消息传递的优势。
• Topoformer 在 500 个节点图上的束搜索推理平均耗时仅为 1.62 秒，达到最先进性能，而其基线模型的推理时间最高可达其 4000 倍。
• 在真实计算图上，该模型持续优于强基线模型，验证了其在工业级编译器工作流中的实际相关性。
• 所提出的合成图生成算法支持可扩展、可复现的基准测试，且已公开提供给研究社区使用。