王欣然教授團(tuán)隊提出基于二維材料的高效稀疏神經(jīng)網(wǎng)絡(luò)硬件方案

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

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

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

圖1 硬件的索引瓶頸造成稀疏神經(jīng)網(wǎng)絡(luò)效率低下

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

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

圖2 “存內(nèi)稀疏”架構(gòu)設(shè)計

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

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

圖3 基于免索引硬件的片上稀疏訓(xùn)練

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

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