1月8日，王欣然教授領導的合作團隊在國際頂級期刊《Nature Electronics》以“An index-free sparse neural network using two-dimensional semiconductor ferroelectric field-effect transistors”為題發(fā)表最新研究進展，報道了基于二維材料的高效稀疏神經網絡硬件解決方案。

稀疏性 (Sparsity)是人腦中的神經突觸的本征屬性。在大腦發(fā)育過程中，超過一半的突觸會以細粒度和非結構化的方式被剪枝(Pruning)，這是人腦具有高能效的關鍵因素。受此啟發(fā)，稀疏神經網絡(Sparse neural network)早在上世紀九十年代就被提出，已成為人工智能輕量化的主流路徑。

雖然稀疏神經網絡中90%的權重可以被剪枝，但剪枝過程中需要反復與外部存儲設備進行索引（Indexing），消耗了整個系統(tǒng)90%以上的能耗和時延。類比于分離計算和存儲的馮諾依曼瓶頸，索引過程成為稀疏神經網絡硬件的瓶頸(圖1)。目前產業(yè)界主流方案為英偉達自A100芯片開始提出的粗粒度、結構化剪枝(Structured pruning)，其本質是通過一定程度的精度犧牲來減輕索引開銷，并未在根本上解決索引瓶頸。

圖1 硬件的索引瓶頸造成稀疏神經網絡效率低下

為了解決稀疏神經網絡軟硬件不適配的問題，王欣然教授領導的合作團隊借鑒了神經生物學模型：支持突觸產生、剪枝、重新生長等動態(tài)行為的，并非神經元本身，而是環(huán)繞在神經元和突觸周圍的星形膠質細胞 (Astrocytes)和小膠質細胞 (Microglial)。受此啟發(fā)，稀疏性信息在位置上也需要盡可能接近權重信息，并直接參與稀疏運算。

在本工作中，團隊首次提出了“存內稀疏”計算架構(In-Memory Sparsity)，其把稀疏神經網絡的訓練過程抽象為稀疏矩陣和權重矩陣的Hadamard乘積，并集成在一個單元內部(圖2)。整個稀疏網絡的硬件基于二硫化鉬(MoS?)鐵電晶體管技術，每個單元包含兩個鐵電晶體管，其中模擬鐵電晶體管用于存儲權重數據，而數字鐵電晶體管用于編碼稀疏性信息，直接決定權重是否被修剪。稀疏性信息被提前編程而免除了外部索引，從而大大降低了稀疏神經網絡訓練的開銷。

圖2 “存內稀疏”架構設計

為支撐存內稀疏架構芯片級并行計算，團隊開發(fā)了以矢量近似更新算法(Vectorial Approximate Updating, VAU)為核心的軟件-硬件協(xié)同優(yōu)化方法(Software-Hardware Co-Optimization, SHCO)。其要點在于：摒棄傳統(tǒng)更新方案中完全精確但低效的逐個單元更新、以及高效卻失準的行列更新，而是在預先的稀疏性編碼的基礎上，以行列為最小單元進行近似更新。結果表明，VAU算法在稀疏硬件上實現(xiàn)的精確度可以和理論值媲美，證明了軟硬件協(xié)同優(yōu)化的必要性。基于存內稀疏架構和軟硬件協(xié)同優(yōu)化方法，團隊制備了硬件陣列，完成了多種稀疏訓練過程的片上實測：包括預訓練、剪枝、過剪、以及重新生長，最后在75%的稀疏率下實現(xiàn)了精度為98.4%的EMNIST手寫字母分類(圖3)。

為了證明硬件方案的可拓展性，團隊基于NeuroSim仿真工具，將經典卷積神經網絡VGG8-Net部署在三種不同的硬件架構上：密集硬件、傳統(tǒng)稀疏硬件、以及本文提出的免索引稀疏硬件。仿真證明，存內稀疏架構的免索引稀疏硬件，首次基于極細粒度和非結構化稀疏性，實現(xiàn)一個數量級的能耗和時延收益。

圖3 基于免索引硬件的片上稀疏訓練

綜上所述，受人腦啟發(fā)，團隊首次提出了稀疏神經網絡的“存內稀疏”計算架構，并基于二維半導體鐵電晶體管技術進行了免索引單元開發(fā)和陣列級片上演示，打破了稀疏神經網絡硬件長期面臨的瓶頸。本工作也充分展示了二維半導體等新材料、新器件技術賦能人工智能硬件的巨大潛力：二維材料具有低溫后道工藝兼容的特點，可以與成熟的硅基電路進行三維單片集成，突破先進封裝技術在互聯(lián)密度方面的局限，進一步提升近存、存算芯片的能效。

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