Google传奇Jeff Dean最新演讲：如何构建未来的机器学习芯片

允中 编译整理

量子位 出品 | 公众号 QbitAI

如何构建未来的机器学习加速芯片？

Google大脑负责人Jeff Dean是最有资格回答这个问题的人之一。昨天，还是在NIPS大会期间，Jeff Dean详细介绍了Google在AI芯片方面的最新研究。

以下是Jeff Dean最新演讲的主要内容。

众所周知，深度学习需要大量的计算资源支持，深度学习正在改变我们设计电脑的方式。例如，降低计算精度也是OK的。

Google在新机器学习方面的成果之一，就是TPU。这个专用的AI加速芯片，主要用来执行神经网络的推理计算。Google的搜索、神经机器翻译、语音图像识别，以及大名鼎鼎的AlphaGo背后，都是TPU在提供计算支持。

第一代TPU在推理方面取得了巨大的进步，但是训练怎么办？

于是Google又研发了第二代TPU。第二代TPU被设计用来同时执行训练和推理计算。第二代TPU的架构如下图所示：

Google还用64块TPU组成阵列（TPU Pod），这进一步提升了计算效力。具体有多厉害？例如，训练Resnet-50达到75%以上精确度，单个第二代TPU要耗时一整天，而阵列只需要22分钟，速度提升31倍，不需要任何额外代码优化。

成绩属于过去，未来挑战依然严峻。2009年以来，arXiv上机器学习论文的增长速度，已经超过摩尔定律。

所以，接下来要思考的问题还是：应该如何构建未来的机器学习加速器？如果现在开始着手，如何设计一个两年内能投入使用，五年内不会过时的AI芯片？

需要考虑的问题包括精度、稀疏和嵌入、Batch大小、训练算法等等。但首先是整个系统都应该有所改变。传统的low-level系统代码（操作系统、编译器、存储系统）还没有广泛利用当今的机器学习。

对于更高性能的机器学习模型来说，并行性非常重要。但是在多个计算设备上获得良好的性能，是并不是一件易事。

为什么这样？

因为Learned Index结构，不是传统的索引结构。这部分实际上是Google最新的研究成果。

在这个研究中，Google从假设现在所有的索引结构都可以用其他类型的模型来代替，包括深度学习模型，这被称为Learned Index。核心思想是，一个模型可以学习查询的排序顺序或者结构，并且利用这个信号来有效预测记录的位置。

Google还从理论上分析了Learned Index在哪些条件下，表现优于传统的指标结构，描述了Learned Index结构设计中的主要挑战。

初步的结果显示，在神经网络环境中，这能让经过缓存优化的B-Tree提速70%，同时节省了一个数量级的内存占用。

另外，GPU/TPU让Learned Index更加可行。当然，GPU/TPU还面临高调用延迟等挑战，但是使用批量请求等技术，可以分摊调用成本。

重要的是，Google认为通过学习模型取代数据管理系统核心组件的想法，对未来的系统设计有着深远的影响。

如果对这部分研究的细节感兴趣，可以查看Google的论文The Case for Learned Index Structures，地址在此：https://arxiv.org/abs/1712.01208 。