王欣然教授团队提出基于二维材料的高效稀疏神经网络硬件方案

1月8日，王欣然教授领导的合作团队在国际顶级期刊《Nature Electronics》以“An index-free sparse neural network using two-dimensional semiconductor ferroelectric field-effect transistors”为题发表最新研究进展，报道了基于二维材料的高效稀疏神经网络硬件解决方案。

稀疏性 (Sparsity)是人脑中的神经突触的本征属性。在大脑发育过程中，超过一半的突触会以细粒度和非结构化的方式被剪枝(Pruning)，这是人脑具有高能效的关键因素。受此启发，稀疏神经网络(Sparse neural network)早在上世纪九十年代就被提出，已成为人工智能轻量化的主流路径。

虽然稀疏神经网络中90%的权重可以被剪枝，但剪枝过程中需要反复与外部存储设备进行索引（Indexing），消耗了整个系统90%以上的能耗和时延。类比于分离计算和存储的冯诺依曼瓶颈，索引过程成为稀疏神经网络硬件的瓶颈(图1)。目前产业界主流方案为英伟达自A100芯片开始提出的粗粒度、结构化剪枝(Structured pruning)，其本质是通过一定程度的精度牺牲来减轻索引开销，并未在根本上解决索引瓶颈。

图1 硬件的索引瓶颈造成稀疏神经网络效率低下

为了解决稀疏神经网络软硬件不适配的问题，王欣然教授领导的合作团队借鉴了神经生物学模型：支持突触产生、剪枝、重新生长等动态行为的，并非神经元本身，而是环绕在神经元和突触周围的星形胶质细胞 (Astrocytes)和小胶质细胞 (Microglial)。受此启发，稀疏性信息在位置上也需要尽可能接近权重信息，并直接参与稀疏运算。

在本工作中，团队首次提出了“存内稀疏”计算架构(In-Memory Sparsity)，其把稀疏神经网络的训练过程抽象为稀疏矩阵和权重矩阵的Hadamard乘积，并集成在一个单元内部(图2)。整个稀疏网络的硬件基于二硫化钼(MoS₂)铁电晶体管技术，每个单元包含两个铁电晶体管，其中模拟铁电晶体管用于存储权重数据，而数字铁电晶体管用于编码稀疏性信息，直接决定权重是否被修剪。稀疏性信息被提前编程而免除了外部索引，从而大大降低了稀疏神经网络训练的开销。

图2 “存内稀疏”架构设计

为支撑存内稀疏架构芯片级并行计算，团队开发了以矢量近似更新算法(Vectorial Approximate Updating, VAU)为核心的软件-硬件协同优化方法(Software-Hardware Co-Optimization, SHCO)。其要点在于：摒弃传统更新方案中完全精确但低效的逐个单元更新、以及高效却失准的行列更新，而是在预先的稀疏性编码的基础上，以行列为最小单元进行近似更新。结果表明，VAU算法在稀疏硬件上实现的精确度可以和理论值媲美，证明了软硬件协同优化的必要性。基于存内稀疏架构和软硬件协同优化方法，团队制备了硬件阵列，完成了多种稀疏训练过程的片上实测：包括预训练、剪枝、过剪、以及重新生长，最后在75%的稀疏率下实现了精度为98.4%的EMNIST手写字母分类(图3)。

为了证明硬件方案的可拓展性，团队基于NeuroSim仿真工具，将经典卷积神经网络VGG8-Net部署在三种不同的硬件架构上：密集硬件、传统稀疏硬件、以及本文提出的免索引稀疏硬件。仿真证明，存内稀疏架构的免索引稀疏硬件，首次基于极细粒度和非结构化稀疏性，实现一个数量级的能耗和时延收益。

图3 基于免索引硬件的片上稀疏训练

综上所述，受人脑启发，团队首次提出了稀疏神经网络的“存内稀疏”计算架构，并基于二维半导体铁电晶体管技术进行了免索引单元开发和阵列级片上演示，打破了稀疏神经网络硬件长期面临的瓶颈。本工作也充分展示了二维半导体等新材料、新器件技术赋能人工智能硬件的巨大潜力：二维材料具有低温后道工艺兼容的特点，可以与成熟的硅基电路进行三维单片集成，突破先进封装技术在互联密度方面的局限，进一步提升近存、存算芯片的能效。

原文链接：https://www.nature.com/articles/s41928-024-01328-4