你知道吗?对于一块普通芯片来讲,要想处理大脑皮层表面重建这样一个复杂数学任务,至少也得花费几十分钟。由北京大学杨玉超教授团队和中国科学院上海微系统与信息技术研究所宋志棠研究员团队打造的新芯片,把这一时间压到了零点几秒。
这是一颗 40 纳米制程的芯片,运行一次神经动力系统的单次迭代时间仅需要 2.12 毫秒,比此前已知最快的同类芯片快了几十倍之多。
它的核心能力藏在一种名为相变存储器的器件里,研究人员把这种存储器做出了新花样,一边用它来存储数据,一边用它来调整计算的步长。研究中,他们利用这颗芯片进行了大脑皮层和灰质表面的 3D 重建,结果显示误差低于千分之一毫米,速度比此前最好的 GPU 快了将近五百倍之多。
对于神经动力系统来说,“神经”说的是它里面嵌入了神经网络,可被用于学习数据里的规律。“动力系统”指的是一套数学方程,可被用来描述一个东西随时间变化的规律。当这两样叠加在一起,就成为了一个既能学习又能演化的计算模型。
这样的模型很适合用来重建 3D 物体的表面,因为物体的形状是千变万化的,使用传统方法一帧接一帧地算,不仅慢而且还容易出错。
神经动力系统则能够把形状变化当作是一个连续流动的过程,它先从一张粗糙的模板出发,然后让模板慢慢地发生变形,直到能够贴合上真实的物体表面。从数学角度来看这套方法非常严格,可以保证在变形的过程中,曲面不会自己穿破自己,这样一来出来的结果永远是光滑、完整且没有破洞的 3D 网格,这一特性对于医学影像、计算机图形学和 AR/VR 等场景非常重要。
但是神经动力系统存在一个巨大的问题,计算起来太慢了。它需要反复地做数值积分,每一步都需要调整步长,步长太大就会导致结果不稳,步长太小则会导致算不完,这时控制器就要不停地试错。
这个试错的过程会涉及到大量的读、写、乘、加,传统芯片在进行这些操作的时候要来回搬运数据,不仅费时而且费电。本次研究人员换了一种思路,他们利用相变存储器来做步长的调整。
相变存储器,是一种依靠热量来改变自身电阻的材料,能够通电加热让晶体状态发生改变,这时电阻就会跟着变。研究人员发现这种电阻变化并不是死板的,它拥有一种叫作漂移的现象,也就是电阻会慢慢地自己走一点,以前学界觉得这是缺陷,但是本次研究人员觉得这可能是机会。
研究中,他们把电阻漂移加以控制,让它按着预先设计的方向和速度前进。当步长要变大的时候就把电阻调整到某个值,当步长要变小的时候就调整到另一个值。在这个过程之中,完全不需要额外的计算电路,也不需要来回地搬运数据,当电阻自己走完这段路程,步长就能够自动调好。
这个设计方案把步长搜索的硬件面积压缩到了传统方案的十分之一。在芯片上还有一片相变存储阵列,它专门用来运行神经网络,每一步要算的乘法和加法直接可以在存储器里完成。数据不需要搬出去,电流一通结果就会出来。
在芯片上,相变存储器阵列运行着神经网络,同一片材料又负责步长漂移,这两种功能是互不干扰的,配合得十分和谐。这样一来,让相变存储器利用这种方式变成了计算的一部分,不再只是存数据的地方。
在传统芯片的设计方案里,步长调整需要依靠一堆数字电路完成,计数器、比较器、乘法器和加法器这些东西堆在一起,不仅非常占地方,而且每次调整都要读数据、算数据、写数据,来回得折腾好几轮。
相变存储器的漂移效应则能把这几步全给绕过去,电阻会自己发生变化,变化的速度和方向都是受控的,这就相当于直接在物理层面完成了步长的调节。
据了解,为了达到上述控制精度,研究人员在材料和工艺上做了很多优化。他们在相变材料里掺了一些碳,碳的作用在于可以细化晶粒,能够让材料在反复相变之后仍能保持稳定的电学性能。
测试结果显示,这颗芯片的写擦次数能够达到十的十次方,换算成为实际使用时间,能够连续跑上好几年都没问题。当温度从零度变化到七十度时,电阻分布只会发生整体移动,不同电阻档位之间一直保持着清晰的间隔,完全不会混在一起。这一点的重要性在于,由于芯片在实际设备里会遇到各种各样的温度环境,比如手机发烫、冬天户外和服务器机房里散热不良等,所以每一种情况都必须保持稳定。
相变存储器的另一个优点是能够存多级电阻。对于普通存储器来说,它们只能存 0 和 1 两种状态。而本次芯片能够稳稳地存 16 种不同的电阻值,在阵列里排列成一个差分结构,每一个存储单元能够表达正负各八个等级。这就等于在同样的面积里放进去了更多的权重信息,神经网络运行起来也就更快。
为了让器件达到上述密度和稳定度,研究人员在工艺上做了不少努力,掺了一些碳来对晶粒进行细化,从而让材料在反复相变之后仍能保持稳定的性能,这让芯片的写擦次数可以达到十的十次方,连续工作寿命可以用上好几年。当温度从零度变化到七十度,电阻分布也只发生线性移动,没有出现明显的交叠或失控现象。
研究人员做了几个测试来验证芯片的性能。最硬核的一个测试是大脑皮层的 3D 重建,这项任务在医学影像里非常耗时。采用传统工具 FreeSurfer 运行一次要花两三个小时,即便是 16 核的服务器也得两个半小时才能跑完。而假如换上普通的 GPU,使用同样的神经动力系统算法,最快也要花费将近两秒才能算完一次。
本次新芯片只用了 426 毫秒,比 GPU 快了 50 倍之多,重建出来的大脑皮层的表面误差非常小,灰质和白质的平均距离误差分别仅有 0.245 和 0.376 毫米,既没有破洞也没有交叉,能被直接拿去 3D 打印做脑科手术的导航模型。
他们还使用这颗芯片做了更加复杂的 3D 流形网格生成,运行一次单次迭代只需要 2.12 毫秒,比此前已知最快的同类芯片快了 36 倍,功耗却仅有后者的二十四分之一。进行一次完整的表面重建,这颗芯片消耗的能量大约是给手机充一次电的千分之一。那些以前得用服务器吭哧吭哧算半小时的活,这颗芯片只需眨几次眼的工夫就能跑完,所使用的电量还少得几乎可以忽略不计。
研究人员还设计了一个时间交错机制,让步长漂移的工作轮流在不同的存储行上开展。这样依赖,每一行的工作负担都能均匀分布,整个阵列的寿命被大大延长。
对于单个相变存储器来说,它的写擦次数上限是十的十次方,通过这种轮流调度机制,让整个芯片的实际使用寿命远远超出了单器件的极限。这种工程角度的考虑说明研究人员从一开始就在考虑如何把技术从实验室推向真实环境,而不只是发表一篇论文。
和传统的神经动力系统加速器相比,这颗芯片在面积上的优势也十分显著。在相同任务之下,采用传统方案的时候,需要 0.7 平方毫米的乘法累加电路以及 0.26 平方毫米的权重缓存,加起来大约是一个平方毫米。
本次研究人员使用的相变存储方案把权重直接存在存储阵列里,乘法累加直接在阵列里面完成,步长调整也由漂移效应进行代劳,这让整颗芯片的计算相关面积只有 0.28 平方毫米。面积小了之后,功耗也会跟着降,芯片的发热减少了,散热方案也很简单,对于封装和系统集成非常友好。
不过,这颗芯片目前还只是一个原型,但是它的意义已经超出了几组漂亮的数字。在传统计算机之中,存储和计算是分开的,数据在两者之间来回“搬家”,这个来回过程成了算力的主要瓶颈。本次研究人员使用的相变存储器直接在存储的位置做计算,“搬家”的路省了,“搬家”的电也省了。
此前,神经动力系统对于算力的要求一直很高,许多任务只可以在大型服务器上运行,现在使用一颗专用芯片就能拿下,速度还快了上百倍之多。在磁存储时代来临之前,很多研究人员说相变存储器这种材料不好用,漂移不够稳定、噪声也比较大。
而本次研究人员绕过了这些缺陷,把缺陷变成了功能。让漂移不再是噪声,摇身一变成了步长搜索的引擎。让电阻不再是固定的值,脱胎换骨成了可调的计算参数,让存算一体得到了一次具体的验证。
大脑皮层重建的案例只是一个开头,在物理世界之中需要做高保真表面建模的地方还有很多,比如冠脉血管的 3D 重建、无人驾驶的实时环境建模、文物数字化保护以及游戏里的动态角色形变,每一个场景都需要在极短的时间内算出光滑、完整、没有错误的 3D 表面。
本次芯片展示了这种计算到底可以快到什么程度,也展示了存算一体这条路到底可以走多深。补充成果的论文发表在《科学》,这一成果的重要性之一就在于,相变存储器用在数字计算机里折腾了许多年,但在消费市场始终没火起来。而将其用在存算一体的模拟计算里,它的漂移和电阻多级性反而成了一种好用的工具。
本次论文的共同通讯作者还有北京大学朱毅鑫研究员和北京大学陶耀宇副研究员,第一作者是北京大学博士后蔡磊(现任职于北京化工大学讲师)、中国科学院上海微系统与信息技术研究所解晨晨研究员以及北京大学博士后闫龙皞。
预计这款芯片的改进方向还有很多,比如阵列规模可以继续扩大、外围电路可以再优化、跟大模型的接口也可以做得更紧密,许多以前做不到的事情,未来会一个接一个地变成现实。
参考资料:
相关论文https://www.science.org/eprint/WEY75M4YUHGJVGTEX5YC/full?activationRedirect=/doi/full/10.1126/science.aee6277
http://shmmc.kjtj.cas.cn/zj/201505/t20150519_493975.html
https://www.ece.pku.edu.cn/info/1045/2542.htm
https://www.ai.pku.edu.cn/info/1137/2306.htm
https://www.ai.pku.edu.cn/info/1136/1864.htm
https://www.linkedin.com/in/%E6%99%A8%E6%99%A8-%E8%A7%A3-4221a5b4/
https://ic.pku.edu.cn/szdw/bsh/index.htm
https://www.linkedin.com/in/%E6%99%A8%E6%99%A8-%E8%A7%A3-4221a5b4/
排版:胡巍巍
注:封面/首图由 AI 辅助生成
热门跟贴