鱼与熊掌不可兼得？二维存算一体器件实现AI训练-推理一体化

近日，南京大学王欣然教授团队与清华大学吴华强教授团队合作，提出基于二维半导体铁电晶体管的新型存内计算器件架构，通过调节铁电势阱，实现了同时满足AI训练和推理需求的底层器件，并展示了高达103 TOPS/W级别的能效潜力。该成果突破了边缘端人工智能硬件的关键瓶颈之一。

AI (人工智能)的历史，可以归纳为软件和硬件的协同发展史，即“强大硬件推动软件发展”和“复杂软件催生硬件革命”的交替迭代。如今，我们正处在“软催硬”的周期内。以最近大热的现象级AI应用ChatGPT为例，其算力可以达到每秒1019次计算操作（即107 TOPS），而维持其夸张算力的基础是上万块最先进的GPU以及77,160 kWh日耗电量——超过了全球40%国家的日发电量！因此，基于冯诺依曼架构的GPU硬件将难以支撑AI技术的可持续性发展。实际上，人脑算力与GPT相当，功耗却仅为20W，能效比GPT采用的GPU高5个数量级。因此，发展高能效的类脑计算器件和架构势在必行。
受大脑存算一体特征的启发，存内计算架构近年来蓬勃发展，显著降低了“冯诺依曼”架构中数据传输的能耗与延迟，大幅提高了计算并行度和能效。在此基础上，更高维度的需求——训练和推理一体化，即“学”和“用”在同一硬件上实现（如图1所示的主动学习能力)，被认为是通向智慧终端的核心挑战。

图1 训练-推理一体化更接近人脑的计算方式

其中的瓶颈就在于，推理和训练过程对于存储器硬件本身的要求截然不同。从物理本质上看，两者好似鱼和熊掌一般不可兼得。如图2所示，推理的实际过程，是数据在权重矩阵中做批量的乘加运算，它要求存储器具有高势垒，器件上体现为长久的数据保持(Long data retention)，而不需要有很高精度的多态(Multi-state)；而训练过程，则是对权重矩阵中每一个节点进行更新，这要求存储器具有较低的势垒，器件上体现为极高的耐久性(High endurance)、极快的操作速度(High speed)和较小的翻转功耗(Low switch energy)。长久以来，由于这种制衡关系的存在，大家难以找到一种普适的存储器来实现训练-推理一体，这严重制约了AI硬件的发展。

图2 推理和训练过程对于存储器性能的不同要求

在这项工作中，研究人员巧妙利用了铁电材料能量双阱曲线的可调性，并从“材料-物理-器件-电路-系统”的层面，逐级证明了这种可调性设计用于训练-推理一体特征的可行性。研究人员提出了一种“复式（Duplex）”的二维半导体铁电晶体管结构(如图3所示)，通过设计浮栅两侧铁电电容和介电电容的面积比例(AFE/ADE)，调控了铁电两端的分压比例，从而实现了对铁电能量双阱曲线的形态调节。AFE/ADE较大的栅极，能垒较低、铁电部分翻转、翻转功耗低(22.7 fJ/(bit μm2))、速度快(<4.8ns)、耐久性高 (>1013次翻转)，更适合用于AI训练；而AFE/ADE较小的栅极，能垒较高、铁电翻转较为彻底、数据保持长久(10年)，适合用于AI推理。


图3 满足训练-推理一体的复式器件结构

在电路层面(如图4所示)，研究人员使用1个复式结构和2个开关晶体管组成了2T-1D的训练推理一体(Training-Inference-In-One, TIIO)单元，单元之间通过行列线组成crossbar结构的阵列，该阵列可以实现全模拟(即输入信号未通过数字化)的非线性分类任务。其中的训练和推理，都在同一硬件中原位完成，分别实现了100%的训练精确度和99.86%的推理精确度。所演示的单元设计和操作时序，原则上适用于任意结构和种类的神经网络。
在系统层面，我们基于对边缘智能的主流应用场景之一——自动驾驶，通过仿真演示了TIIO架构在单目深度识别中的成功应用。该架构在计算精度媲美GPU的前提下，显著提高了能效。在22nm先进节点下，TIIO核心阵列的能效极限可达1151 TOPS/W（训练）和111.86 TOPS/W（推理）。

图4 基于TIIO阵列的原位训练和推理演示

面向未来，研究人员提出的TIIO架构旨在实现无云端依赖的边缘智能(Edge intelligence)，即具有本地学习、快速决策、数据安全、低能耗和高能效的“更加聪明的终端”。另外，由于其集成工艺与CMOS工艺兼容，可以与硅基电路异质集成，构建面向原位学习的三维集成电路。最后，该架构还将促进迁移学习(Transfer learning)和终身学习(Lifelong learning)等核心AI方法和场景的深入发展。
成果以“An in-memory computing architecture based on a duplex two-dimensional material structure for in situ machine learning”为题，于2023年3月20日在线发表于Nature Nanotechnology。论文链接：https://www.nature.com/articles/s41565-023-01343-0。
该工作由南京大学、南京邮电大学、苏州实验室、清华大学、香港理工大学等单位合作完成。南京大学王欣然教授、南京邮电大学于志浩教授、清华大学高滨副教授和吴华强教授为论文的共同通讯作者，南京大学博士生宁宏凯、温恒迪、南京邮电大学于志浩、清华大学张清天工程师和高滨为论文的共同第一作者。该研究得到了国家自然科学基金、国家重点研发计划、江苏省前沿引领技术基础研究专项等经费资助，以及南京大学微制造与集成工艺中心的工艺支持、北京市未来芯片技术高精尖创新中心的仿真支持。