【新智元导读】Google研究科学家提出一种可以用于计算图优化的端到端的可传递深度强化学习方法，在显著减少限制的情况下，能够大幅提高图形优化任务的性能，并在泛化性和多任务上表现优越。

当今的时代，有越来越多的app依靠在各种加速器上训练的大型复杂神经网络驱动。而通过将高层计算图映射到特定设备的低层可执行的ML编译器，可以达到简化此过程的目的。

 

为此，ML编译器需要解决很多优化问题，并且会牵扯到启发式算法，从而会带来许许多多的挑战，并在实际操作中常常无法拥有最佳解决方案。

近几年，使用基于学习的方法已经显示出不错的效果，但是这些方法依然具有局限性，从而导致在实践中并不可行。

 

而Google的研究人员在NeurIPS 2020的oral论文“ML编译器的可传递图优化“（《Transferable Graph Optimizers for ML Compilers》）中，提出了一种可以用于计算图优化的端到端的可传递深度强化学习方法。

该方法成功克服了以上提到的所有限制，并且具有良好的效果——

 

与Tensorflow默认优化相比，研究者证明，该方法在三个图形优化任务上的速度提高了33%-60%，在由多达80，000节点组成的不同表示图集中（包括Inception-v3, Transformer-XL, and WaveNet），在收敛速度提升15倍的同时，该模型的效果比专家优化平均提升了21%，此外，还比现有的技术高了18%。

重构决策问题，细化至节点级别

在ML编译器中，有三种经常出现的优化任务，这三个优化任务是什么呢？

 

第一个优化任务是设备放置。

第二个优化任务是操作调度。

第三个优化任务是操作融合。

 

Google的研究人员将其公式了化为可以使用学习策略解决的决策问题。其中，每个任务对应的决策问题，可以被重构为对计算图中的每个节点作决策。

注意力机制和GraphSAGE统统上场，三个关键步骤搭起整体框架

在论文中，研究科学家提出了GO。

 

GO是一个深层的强化学习框架，该框架可以用来解决上面提到的优化问题——可以单独一个个解决，也可以将优化问题联合起来解决。

该框架包含三个关键方面：

 

1 首先，研究者使用图神经网络（特别是GraphSAGE）来捕获计算图中编码的拓扑信息。

GraphSAGE的归纳网络利用了节点的属性信息，并将其概括为隐式图，从而在不产生大量训练成本的同时，为非显式数据做出决策。

2 其次，许多模型的计算图通常包含超过1万个节点。要在如此大规模的范围内有效解决优化问题，网络需要能够捕获节点之间的长期依赖关系。 GO的体系结构包括一个可扩展的注意力网络，该网络使用了分段循环，来捕获此类远程节点依赖性。

 

3 ML编译器需要解决来自不同应用领域的各种图形优化问题。如果单纯训练具有异构图的共享策略网络，那么其实并不能充分获取一些特定类型图的特征。为了克服此问题，GO使用了一种特征调制机制，该机制允许网络专门用于特定的图形类型，而无需增加参数数量。

 

此外，为了共同解决多个相关的优化任务，GO可以为每个任务添加额外的循环注意力层，并在不同任务之间共享参数。具有操作残留关联的循环注意力层，可以达到跟踪任务间的依存关系的目的。

疯狂加速15倍，泛化效果强，还不限于单任务

在真实硬件配置环境中，研究者给出了在设备放置任务上的单任务加速评估结果，接着，研究者将其泛化到隐式图中，并使用不同GO变量进行了分析。最后，研究者评估了在设备放置和设备调度上共同优化的多任务结合的评估效果：

 

加速性能：

 

为了评估这种架构的性能，研究者将GO应用到基于真实硬件评估的设备放置问题上。首先，他们针对每个工作负载分别对模型进行训练，并将其称为GO-one。

 

实验结果表明，GO-one始终优于专家放置，TensorFlow METIS放置和分层设备放置（HDP）。对于大型图形（如GNMT），GO可以实现优化的图形运行时间分别比HP和HDP快21.7％和36.5％。

 

总体而言，与HP和HDP相比，GO-one在14组不同图中的运行时间平均减少了20.5％和18.2％。

 

此外，重要的是，通过高效的端到端single-shot布局，相比HDP来说，GO可以将布局网络的收敛速度加快15倍。

 

泛化性能：

 

GO使用离线预训练，将图泛化为非显式图，然后对非显式图进行微调。

 

在预训练期间，研究者在训练图的异构子集中训练GO。在切换到下一批图之前，我们在每批图上训练1000步。然后，使用该预训练模型，在hold-out图上进行微调（GO泛化+微调），此步骤少于50个步并通常需要不到一分钟的时间。

 

对于hold-out图，GO泛化+微调在所有数据集中的表现，均优于专家级放置，HDP和GO-one的平均配置。

在对目标hold-out图没有进行任何微调的情况下，研究者还直接在预训练模型上进行了推理，并将其命名为GO-generalization-zeroshot。

 

这种未调整模型的性能仅比GO泛化+微调稍差，而稍好于专家放置和HDP。

 

这表明图形嵌入和学习到的的策略都可以有效传递，从而使模型可以泛化到隐式数据。

 

共同优化放置，调度和融合任务性能：

 

如果同时优化设备放置，调度和融合，那么，与未优化的单GPU情况相比，会直接提速30％-73％，与TensorFlow默认配置相比，可以提速33％-60％ 。

与单独优化每个任务相比，多任务GO（pl + sch + fu）平均速度优于单任务GO（p | sch | fu）（一次优化所有任务）7.8％。

 

此外，对于所有工作负载，与对其中两个任务进行优化并同时对第三个任务使用默认策略相比，对三个任务进行共同优化会拥有更快的运行时间。