数据要素产业
Google:可迁移架构探索,用AI设计AI芯片
作者/文龙
摩尔定律的迫在眉睫和深度学习的不断使用推动了针对特定神经体系结构进行优化的定制芯片的设计。这些定制芯片也被称为加速器,可以显著地增加可用的算力,从而提高神经网络的性能以及广泛应用的可能性。
在谈到有关该公司是如何利用人工智能促进其定制芯片的研发时,Google Brain的总监杰夫迪恩(Jeff Dean)指出,在某些情况下,深度学习可以比人类做出更好的决策,像是规划芯片中的电路布局。
2月初,Google发表了论文《阿波罗:可迁移架构探索》(Apollo:Architecture Exploration),讲述了名为Apollo的研究项目,该项目提供了一种通用的芯片架构探索框架。
芯片架构指的是芯片功能组件的设计,它们如何相互作用以及软件程序员应如何访问这些组件。例如,经典的Intel x86处理器具有一定数量的片上内存、专用的数学逻辑单元和寄存器以一定的组合方式组成,这种组合方式赋予「英特尔架构」以意义。
定制芯片需要高效的架构探索方法。然而,由于搜索空间呈指数级增长,而目标函数又是一个黑匣子,优化架构的参数是一个艰巨的任务,通常需要专业知识。施加在参数上的约束进一步使有效配置的识别变得复杂。
为了解决上述问题,Apollo程序执行的是适合运行某种神经网络的架构的探索。
这样,目标就变成了,对于芯片架构来说正确的参数是什么,对于给定的神经网络任务,芯片应满足哪些标准?
该团队将目标函数定义为找到适合于给定的AI任务的某种正确的数学运算单元组合,并对超过4.52亿个参数进行排序检索,包括将使用多少个数学运算单元(处理器)以及对于给定模型而言最佳的参数存储量和激活存储量。
图示:人工智能加速器芯片的搜索空间,即芯片架构必须优化的功能块。许多AI芯片的特征是用于大量简单数学运算的并行,相同的处理器用于执行许多矢量矩阵乘法,这是神经网络处理的主力军。(来源:Google Blog)
他们选择了一个名为MobileNet的较为简单的AI任务进行测试,这是一种由资源节约型网络组成的卷积神经网络。实验演示了如何利用针对复杂的高维度的定制优化策略进行架构探索,并以合理的少量评估(约占搜索空间的0.0004%)产生大的改进(高达24.6%)。此外,他们还测试了执行诸如对象检测和语义分段之类任务的工作负载。
Apollo可以采用多种方法进行所谓的「黑盒优化」,并且可以比较每种方法在解决不同目标的效果,使这些方法适应于特定的任务。
该团队在架构探索的背景下探索了四种优化策略,这些优化策略实际上是为开发神经网络体系结构而设计的,包括随机采样、贝叶斯优化、为「设计」 DNA序列而研发的进化搜索以及基于总体的黑盒优化(P3BO)。
对这些优化进行比较的主要目的是为了有条不紊地指出最有效的方法。实验结果表明,进化方法和P3BO优于随机选择和贝叶斯优化。
图示:Apollo的优点是可以将各种现有的优化方法相结合,以了解它们如何在优化新型芯片设计的体系结构中相互叠加。(来源:Google Blog)
Apollo的搜索是高效的,可迁移的。
研究人员将P3BO与架构解决方案的半穷尽搜索集进行了比较,发现基于总体的黑盒优化方法能够发现并使用像是计算与内存这类的电路中权衡取舍的解决方案,而这些通常需要特定领域的知识。并且,基于总体的方法可以找到半穷尽搜索无法找到的解决方案。
因此,Apollo能够弄清楚不同的优化方法在芯片设计中的效果如何。但是,它还可以做更多的事情,那就是它可以运用「迁移学习」展示如何依次改进这些优化。
Yazdanbakhsh及其同事将Apollo称为「第一个可迁移的架构探索的基础框架」,它在不同的芯片上工作的越多,越能更好地探索可能的芯片体系结构,从而将学到的知识迁移到每个新任务上。研究表明,迁移学习可以有效地改善目标体系结构的探索,尤其是当优化约束具有更严格的界限时。
使用AI设计运行AI的芯片具有很好的对称性。
Apollo代表着一个有趣的发展,超越了之前所说的机器学习仅用于一些「布局与布线」之类的低层次的设计决策。相比之下,Apollo程序执行的是更高层次的「架构探索」,而不仅仅是创建类似于建筑中的平面设计图的产物。
当被问及Dean的描述时,主要作者阿米尔雅兹丹巴赫什(Amir Yazdanbakhsh)解释:「我会看到我们的工作和布局布线工作是正交且互补的。相比于计算堆栈中的布局布线,架构探索的层次要更高。」他认为「架构探索」是存在更高性能的改进余地的方式。
但是,为MobileNet或任何其他网络设计的芯片受到一定程度上的适用性的限制。
文章的作者之一Berkin Akin指出优化结果是芯片和神经网络共同优化的产物。「神经网络架构必须了解目标硬件架构,以优化整体系统性能和能效。」他表示,Apollo对于给定的工作负载可能就足够了,但在芯片和神经网络之间进行「共同优化」,将在未来产生其他好处。
「我们用ML为给定的工作负载找到最佳的硬件体系结构。当然,在某些情况下,可以灵活地同时优化硬件设计和神经网络体系结构。实际上,我们正在进行一些此类联合优化的工作,我们希望可以取得更好的折衷... 」
因此,虽然芯片设计受到AI的工作量的影响,但是新的芯片设计进程也会对神经网络的设计产生相当的影响,这类的辩证法在未来几年中可能会以有趣的方式发展。
论文链接:https://arxiv.org/abs/2102.01723
参考内容:
https://www.zdnet.com/google-amp/article/googles-deep-learning-finds-a-critical-path-in-ai-chips/
https://ai.googleblog.com/2021/02/machine-learning-for-computer.html