Google Brain 总监 Jeff Dean 曾指出,在某些情况下,人工智能的深度学习形式在如何布置芯片中的电路方面可以比人类做出更好的决策。
本月,谷歌在 arXiv 文件服务器上发表了一篇名为“ Apollo: Transferable Architecture Exploration”的论文。Apollo 计划代表了一种有趣的发展,它超越了 Jeff Dean 一年前所讲的东西。相比之下,Apollo 计划执行的是“架构探索”而不是平面图。
Apollo 计划是超越“布局和路线”的“架构探索”
这篇论文的主要作者 Amir Yazdanbakhsh 说:“在计算堆栈中,架构探索比布局路线的探索要高级得多。”
在 Jeff Dean 当时给出的例子中,机器学习可以用于一些低层次的设计决策,即“布局和路线”。在位置和路径上,芯片设计者使用软件来确定构成芯片操作的电路的布局,类似于设计建筑物的平面图。
而芯片的体系结构是芯片功能元素的设计,包括如何交互,以及软件程序员应该如何访问这些功能元素。
例如,经典的 Intel x86 处理器具有一定数量的片上存储器、专用的算术逻辑单元和一些寄存器等等。这些部分组合在一起的方式赋予了所谓的英特尔架构的意义。
当被问及 Jeff Dean 的描述时,Yazdanbakhsh 表示,“我们的工作和布局规划项目与 Jeff Dean 所讲的是相互配合且互补的。”
Yazdanbakhsh 说:“我认为,在架构探索中,存在更高的性能改进余地。”
Yazdanbakhsh 和他的同事们称 Apollo 为“第一个可转移的架构探索基础结构”,这是第一个在探索可能的芯片架构方面做得更好的程序,它在不同的芯片上工作得越多,越能更好地探索可能的芯片体系结构,从而将学到的知识转移到每个新任务上。
探索不同开发方法,根据工作负载进行调整
Yazdanbakhsh 和团队正在开发的芯片本身就是用于 AI 的芯片,称为加速器。该芯片与 Nvidia A100“ Ampere” GPU,Cerebras Systems WSE 芯片以及目前投放市场的许多其他启动部件属于同一类。因此,使用 AI 设计运行 AI 的芯片具有很好的对称性。
鉴于任务是设计一个 AI 芯片,Apollo 程序正在探索的架构是适合运行神经网络的架构。这意味着大量的线性代数,许多简单的数学单位,执行矩阵乘法和求和的结果。
团队将挑战定义为找到适合给定 AI 任务的这些数学块的正确组合之一。他们选择了一个相当简单的 AI 任务,即一个称为 MobileNet 的卷积神经网络,由谷歌的 Andrew g. Howard 和他的同事在 2017 年设计。
此外,他们还使用几个内部设计的网络测试工作负载,这些网络用于完成目标检测和语义分割等任务。
通过这种方式,目标就变成了: 对于一个给定的神经网络任务,芯片的结构的正确参数是什么?
整个搜索过程包括对超过 4.52 亿个参数进行排序,其中包括将使用多少被称为处理器元素的数学单元,以及对于给定的模型有多少参数内存和激活内存是最优的。
Apollo 是一个框架,这意味着它可以采用文献中为所谓的黑盒优化开发的各种方法,并且它可以根据特定的工作负载调整这些方法,并比较每种方法在解决目标方面的表现。
在另一个很好的对称性中,Yazdanbakhsh 采用了一些优化方法,这些方法实际上是为开发神经网络架构而设计的。包括 2019 年由 Quoc v. Le 和他在谷歌的同事开发的所谓的进化方法; 基于模型的强化学习方法,以及由 Christof Angermueller 和其他人在 Google 上为“设计”DNA 序列而开发的所谓的基于人口的方法集合; 以及贝叶斯优化方法。
进化和基于模型的方法优于随机选择和其他方法
Apollo 包含了对称性的主要层次,将为神经网络设计和生物合成设计的方法结合起来,设计可能反过来用于神经网络设计和生物合成的电路。
比较所有这些优化,这正是 Apollo 框架的亮点所在。它的整个存在目的是有条不紊地采用不同的方法,并且告诉人们什么方法最有效。Apollo 试验结果详细说明了进化和基于模型的方法如何优于随机选择和其他方法。
但 Apollo 最引人注目的发现是,如何运行这些优化方法,可以使过程比暴力搜索法更有效率。例如,他们将基于群体的集成方法与他们所说的对体系结构方法的解决方案集的半穷尽搜索进行了比较。
Yazdanbakhsh 和他的同事们看到的是,基于人群的方法能够发现利用电路中权衡取舍的解决方案,比如计算和内存,这通常需要特定领域的知识。由于基于人群的方法是一种习得的方法,因此它能够找到半详尽搜索无法找到的解决方案。
实际上,P3BO(基于总体的黑盒优化)发现的设计比半详尽的 3k 样本搜索空间更好。这种设计使用了非常小的内存(3MB),以支持更多的计算单元。这利用了视觉工作负载的计算密集型特性,这在原始的半详尽搜索空间中是不包含的。这表明手动搜索空间工程需要半详尽的方法,而基于学习的优化方法利用大的搜索空间减少手动工作。
因此,Apollo 能够计算出不同的优化方法在芯片设计中的表现。但是,它还可以做更多的事情,可以运行所谓的迁移学习来展示如何反过来改进这些优化方法。
通过运行优化策略以将芯片设计改进一个设计点,例如以毫米为单位的最大芯片尺寸,这些实验的结果便可以作为输入输入到随后的优化方法中。Apollo 团队发现,各种优化方法通过利用初始或种子优化方法的最佳结果,提高了它们在类似区域约束电路设计这样的任务中的性能。
所有这些必须被这样一个事实所包括: 为 MobileNet 或任何其他网络或工作负载设计芯片受设计过程对给定工作负载的适用性的限制。
事实上,该论文的作者之一,Berkin Akin 帮助开发了 MobileNet 版本 MobileNet Edge,他指出优化是芯片和神经网络优化的产物。
Berkin Akin 在去年和同事共同完成的一篇论文中写道:“神经网络架构必须了解目标硬件架构,以优化整体系统性能和能源效率。”
当从神经网络结构设计中分离出来时,硬件设计的价值有多大?
Berkin Akin 认为,Apollo 对于给定的工作量来说可能已经足够了,但是芯片和神经网络之间的协同优化,将在未来产生其他的好处。
他说:“在某些情况下,我们需要为给定的一组固定神经网络模型设计硬件。这些模型可以是来自硬件目标应用程序域的已经高度优化的代表性工作负载的一部分,也可以是定制加速器用户所需的一部分。在这项工作中,我们正在处理这种性质的问题,我们使用机器学习来为给定的工作负载套件找到最好的硬件架构。然而,在某些情况下,联合优化硬件设计和神经网络体系结构是具有灵活性的。事实上,我们有一些正在进行的工作,这样的联合优化,我们希望能够产生更好的权衡。”
最后的结论是,即使芯片设计正在受到人工智能新工作负载的影响,芯片设计的新进程可能对神经网络的设计产生可衡量的影响,这种辩证法可能在未来几年以有趣的方式发展。