[论文解读] Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs

Aditya Paliwal, Felix Gimeno+5 more|arXiv (Cornell University)|May 7, 2019

Metaheuristic Optimization Algorithms Research参考文献 42被引用 25

一句话总结

本文提出 REGAL，一种增强强化学习的遗传算法，通过训练图神经网络（GNN）来指导有偏随机键遗传算法（BRKGA）在优化神经网络计算图以实现最小运行时间和峰值内存使用量。通过学习输入特定的提议分布用于 BRKGA，REGAL 在未见过的真实 TensorFlow 图上仅用数秒即可获得高质量解，且无需微调，优于经典方法和基于学习的基线方法。

ABSTRACT

We present a deep reinforcement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler. Unlike earlier learning-based works that require training the optimizer on the same graph to be optimized, we propose a learning approach that trains an optimizer offline and then generalizes to previously unseen graphs without further training. This allows our approach to produce high-quality execution decisions on real-world TensorFlow graphs in seconds instead of hours. We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage. In comparison to an extensive set of baselines, our approach achieves significant improvements over classical and other learning-based methods on these two tasks.

研究动机与目标

• 开发一种快速、可泛化的神经网络计算图优化器，可在无需微调的情况下适用于多样化的真实世界模型。
• 在静态编译器中联合优化放置与调度决策，以最小化执行成本（运行时间或峰值内存）。
• 克服先前基于学习的优化器的局限性，后者需为每张图训练数小时，或缺乏对未见架构的泛化能力。
• 通过在大型图上数秒内获得高质量解，实现优化编译器在生产环境中的部署。

提出的方法

• 通过上下文 bandit 设置和 REINFORCE 算法训练图神经网络（GNN），以预测 BRKGA 突变生成步骤中每个节点的非均匀 beta 分布参数。
• BRKGA 算法使用这些学习到的、与输入相关的提议分布，而非默认的均匀分布，以引导搜索朝向更高品质的解。
• GNN 处理由计算图结构、节点特征（如操作类型、张量大小）和边依赖关系构成的信息，为每个节点生成与放置和调度相关的特定动作概率。
• 策略通过目标值（如运行时间或内存）作为稀疏奖励信号进行训练，无需显式监督。
• 使用粗粒度静态成本模型在 BRKGA 评估过程中估算候选解的执行成本，从而在训练期间实现快速反馈。
• 通过学习可迁移的决策模式，该方法在包括真实世界 TensorFlow 模型在内的多样化图拓扑上实现良好泛化。

实验结果

研究问题

• RQ1 学习到的策略是否能在无需微调的情况下，引导遗传算法在未见的计算图上找到更优解？
• RQ2 基于 GNN 的策略学习是否能提升静态编译器中联合放置与调度优化的解质量？
• RQ3 该方法是否能在数秒内而非数小时内获得高质量结果，从而实现实时在生产编译器中的应用？
• RQ4 该模型在多样化神经网络架构和工作负载上的泛化能力如何？

主要发现

• REGAL 在真实世界 TensorFlow 图上显著优于经典方法和基于学习的基线方法，在最小化运行时间和峰值内存使用方面均取得显著改进。
• 该方法在数秒内生成高质量解，而先前基于学习的方法（如 Mirhoseini et al. (2017, 2018)）需数小时。
• GNN 策略能有效泛化至未见的图拓扑，包括通过数据增强生成的图，且在不同架构上的性能方差较低。
• 策略学习到了非平凡的、与数据相关的策略——例如，为高内存节点分配较低的调度优先级，并避免对重型节点施加设备偏置——超越简单启发式方法。
• 在 TF Runtime 测试集上，REGAL 在 95% 的情况下优于 BRKGA 基线，且奖励值大于 -1 表明解的质量更优。