近期机器学习 (ML) 进步的关键因素之一是定制加速器的开发，例如Google TPU和Edge TPU，它们显著提高了可用的计算能力，从而解锁了各种功能，例如AlphaGo、RankBrain、WaveNets和Conversational Agents 。这种提升可以提高神经网络训练和推理的性能，为视觉、语言、理解和自动驾驶汽车等广泛应用带来新的可能性。
为了维持这些进步，硬件加速器生态系统必须继续在架构设计方面进行创新，并适应快速发展的 ML 模型和应用。这需要对许多不同的加速器设计点进行评估，每个设计点不仅可以提高计算能力，还可以揭示一项新功能。这些设计点通常由各种硬件和软件因素（例如，内存容量、不同级别的计算单元数量、并行性、互连网络、流水线、软件映射等）参数化。这是一项艰巨的优化任务，因为搜索空间呈指数级增长1 ，而目标函数（例如，更低的延迟和/或更高的能效）通过模拟或综合来评估在计算上成本高昂，从而使得识别可行的加速器配置具有挑战性。
在“阿波罗：可转移架构探索”中，我们介绍了我们在机器学习驱动的定制加速器设计方面的研究进展。虽然最近的 工作已经证明了利用机器学习改进低级布局规划过程（其中硬件组件在硅片上进行空间布局和连接）的良好结果，但在这项工作中，我们专注于将机器学习融入高级系统规范和架构设计阶段，这是影响芯片整体性能的关键因素，控制高级功能的设计元素就是在此阶段建立的。我们的研究展示了机器学习算法如何促进架构探索，并在一系列深度神经网络中推荐高性能架构，涉及的领域涵盖图像分类、对象检测、OCR和语义分割。
架构搜索空间和工作负载
架构探索的目标是发现一组适用于一组工作负载的可行加速器参数，以便在用户定义的一组可选约束下最小化所需的目标函数（例如，运行时的加权平均值）。但是，架构搜索的流形通常包含许多点，这些点没有从软件到硬件的可行映射。其中一些设计点是先验已知的，可以通过用户将它们公式化为优化约束来绕过（例如，在面积预算2约束的情况下，总内存大小不得超过预定义的限制）。但是，由于架构和编译器的相互作用以及搜索空间的复杂性，一些约束可能无法正确地公式化到优化中，因此编译器可能无法找到目标硬件的可行软件映射。这些不可行的点在优化问题中不容易公式化，并且通常在执行整个编译器传递之前是未知的。因此，架构探索的主要挑战之一是有效地避开不可行的点，以便以最少数量的循环精确架构模拟有效地探索搜索空间。
下图显示了目标 ML 加速器的整体架构搜索空间。加速器包含一个 2D 处理元件 (PE) 阵列，每个 PE 以单指令多数据(SIMD) 方式执行一组算术计算。每个 PE 的主要架构组件是处理核心，其中包括用于 SIMD 操作的多个计算通道。每个 PE 在其所有计算核心中都有共享内存（PE 内存），主要用于存储模型激活、部分结果和输出，而各个核心则具有主要用于存储模型参数的内存。每个核心都有多个带有多路乘法累加(MAC) 单元的计算通道。每个周期的模型计算结果要么存储回 PE 内存以供进一步计算，要么卸载回 DRAM。
用于架构探索的基于模板的 ML 加速器概述。
优化策略
在本研究中，我们在架构探索的背景下探讨了四种优化策略：
随机：随机地对架构搜索空间进行均匀采样。
Vizier：使用贝叶斯优化在搜索空间中进行探索，其中目标函数的评估成本高昂（例如硬件模拟，可能需要数小时才能完成）。使用来自搜索空间的采样点集合，贝叶斯优化形成一个替代函数，通常用高斯过程表示，该替代函数近似于搜索空间的流形。在替代函数值的指导下，贝叶斯优化算法在探索和利用的权衡中决定是从流形中的有希望区域采样更多（利用）还是从搜索空间中看不见的区域采样更多（探索）。然后，优化算法使用这些新采样的点并进一步更新替代函数以更好地模拟目标搜索空间。Vizier 将 预期改进作为其核心获取功能。在这里我们使用Vizier (safe)，它是约束优化的变体，可指导优化过程以避免建议不满足给定约束的试验。
进化：使用k 个个体组成的种群 进行进化搜索，其中每个个体的基因组对应于一系列离散化的加速器配置。通过锦标赛选择从种群中为每个个体选择两个父母，以一定的交叉率重组它们的基因组，并以一定的概率对重组后的基因组进行变异，从而生成新的个体。
基于种群的黑盒优化(P3BO)： 使用一组优化方法，包括进化和基于模型的方法，这些方法已被证明可以提高样本效率和稳健性。采样数据在集成中的优化方法之间交换，优化器根据其性能历史加权以生成新配置。在我们的研究中，我们使用了 P3BO 的变体，其中优化器的超参数使用进化搜索动态更新。
加速器搜索空间嵌入
为了更好地直观显示每种优化策略在加速器搜索空间中的有效性，我们使用t 分布随机邻居嵌入 (t-SNE) 将探索的配置映射到优化范围内的二维空间中。所有实验的目标（奖励）定义为每个加速器区域的吞吐量（推理/秒）。在下图中，x轴和y轴表示嵌入空间的 t-SNE 组件（嵌入 1 和嵌入 2）。星形和圆形标记分别显示不可行（零奖励）和可行设计点，可行点的大小与其奖励相对应。
正如预期的那样，随机策略以均匀分布的方式搜索空间，最终在设计空间中找到很少的可行点。
可视化显示通过随机优化策略（最大奖励 = 0.96）探索的设计点（~4K）的 t-SNE。动画的最后一帧突出显示了最大奖励点（红色十字标记）。
与随机采样方法相比，Vizier默认优化策略在探索搜索空间和寻找奖励更高的设计点之间取得了良好的平衡（1.14 vs. 0.96）。然而，这种方法往往会陷入不可行区域，虽然它确实找到了几个奖励最高的点（用红十字标记表示），但在探索的最后迭代中，它发现的可行点很少。
如上所述，采用Vizier（默认）优化策略（最大奖励 = 1.14）。动画的最后一帧突出显示最大奖励点（红色十字标记）。
另一方面，进化优化策略在优化的早期就找到了可行解决方案，并在这些解决方案周围聚集可行点集群。因此，这种方法主要导航可行区域（绿色圆圈），并有效地避开不可行点。此外，进化搜索能够找到更多具有最大回报的设计选项（红色十字）。高回报解决方案的多样性为设计师提供了灵活性，可以探索具有不同设计权衡的各种架构。
如上所述，采用进化优化策略（最大奖励 = 1.10）。动画的最后一帧突出显示最大奖励点（红色十字标记）。
最后，基于种群的优化方法 (P3BO) 以更有针对性的方式探索设计空间（奖励点高的区域），以找到最佳解决方案。P3BO 策略在具有更严格约束（例如，不可行点数量较多）的搜索空间中找到奖励最高的设计点，显示出其在具有大量不可行点的搜索空间中导航的有效性。
如上所述，采用P3BO优化策略（最大奖励 = 1.13）。动画的最后一帧突出显示最大奖励点（红色十字标记）。
不同设计约束下的架构探索
我们还研究了不同面积预算约束（6.8 平方毫米、5.8 平方毫米和 4.8 平方毫米）下每种优化策略的优势。以下小提琴图显示了优化结束时（每次运行 10 次，每次 4K 次试验后）研究的优化策略中最大可实现奖励的完整分布。较宽的部分表示在特定给定奖励下观察到可行架构配置的概率更高。这意味着我们倾向于在奖励较高（性能较高）的点上产生更大宽度的优化算法。
架构探索中表现最佳的两种优化策略是进化和 P3BO，这两种策略都能在多次运行中提供具有高回报和稳健性的解决方案。研究不同的设计约束，我们观察到，随着面积预算约束的收紧，P3BO 优化策略会产生更多高性能解决方案。例如，当面积预算约束设置为 5.8 平方毫米时， P3BO 找到的回报（吞吐量/加速器面积）为 1.25 的设计点优于所有其他优化策略。当面积预算约束设置为 4.8 平方毫米时，也观察到了同样的趋势，在多次运行中，可以发现回报略好，并且稳健性更高（变化性更小）。
小提琴图显示了在 6.8 mm 2的面积预算下，经过 4K 次试验评估后，十次优化策略中可实现的最大奖励的完整分布。P3BO 和进化算法产生了大量高性能设计（更宽的部分）。x 轴和 y 轴分别表示所研究的优化算法和基线加速器上的加速（奖励）的几何平均值。
如上所述，面积预算为 5.8 mm 2。
如上所述，面积预算为 4.8 mm 2。
结论
虽然阿波罗计划迈出了更好地理解加速器设计空间和构建更高效硬件的第一步，但发明具有新功能的加速器仍然是一个未知领域和新前沿。我们相信这项研究是一条令人兴奋的前进之路，可以进一步探索机器学习驱动的架构设计和跨计算堆栈的协同优化技术（例如编译器、映射和调度），从而为下一代应用程序发明具有新功能的高效加速器。
致谢
这项工作由 Amir Yazdanbakhsh、Christof Angermueller 和 Berkin Akin 完成。 我们还要感谢 Milad Hashemi、Kevin Swersky、James Laudon、Herman Schmit、Cliff Young、Yanqi Zhou、Albin Jones、Satrajit Chatterjee、Ravi Narayanaswami、Ray (I-Jui) Sung、Suyog Gupta、Kiran Seshadri、Suvinay Subramanian、Matthew Denton 和Vizier 团队的帮助和支持。
评论