訓(xùn)練下一代萬億級參數(shù)大模型的高效芯片誕生了！

最近，來自清華團(tuán)隊的研究人員開發(fā)了一種革命性的新型AI「光芯片」——「太極」（Taichi）。

不言而喻，「太極」最大的亮點是使用光，而不是電來處理數(shù)據(jù)。

與傳統(tǒng)堆疊PIC芯片方法不同，清華團(tuán)隊首創(chuàng)了分布式廣度智能光計算架構(gòu)，使得「太極」成為全球首款大規(guī)模干涉衍射異構(gòu)集成芯片。

「太極」具備了億級神經(jīng)元的芯片計算能力，可以顯著提高處理速度和能效。

它可以實現(xiàn)160 TOPS/W通用智能計算。

最新研究已于4月11日發(fā)表在Science期刊上。

論文地址：https://www.science.org/doi/10.1126/science.adl1203

更令人震驚的是，「太極」能效是英偉達(dá)H100的1000倍數(shù)。

研究人員表示，「太極」為大規(guī)模的光子計算和高級任務(wù)鋪平了道路，進(jìn)一步發(fā)掘了光子學(xué)在現(xiàn)代AGI中的靈活性和潛力。

ChatGPT耗電大有解了

當(dāng)前，越來越多跡象表明，LLM不會是通往AGI的最終路徑。

那是因為，基于Transformer架構(gòu)的大模型，通過token預(yù)測完成推理，需要消耗大量的算力。

此前ChatGPT日耗電50萬度，曾被網(wǎng)友們吵上熱搜。

若是能夠發(fā)明一種，節(jié)省大量能耗的芯片，LLM的性能或在未來實現(xiàn)更大的飛升。

而「太極」可能會使通用人工智能（AGI）成為現(xiàn)實。研究人員表示，

我們預(yù)計，「太極」將加速開發(fā)更強(qiáng)大的光學(xué)解決方案，為基礎(chǔ)模型和AGI新時代提供關(guān)鍵支持。

在將計算能力提升到AGI所需的水平方面，「太極」的模塊化設(shè)計可能是一個關(guān)鍵優(yōu)勢。

清華團(tuán)隊設(shè)計了一個擁有1396萬個人工神經(jīng)元的分布式「太極」網(wǎng)絡(luò)，超越了其他光芯片設(shè)計（147萬個神經(jīng)元）。

因此，「太極」實現(xiàn)了160.82 TOPS/W的能效。

與2022年一個團(tuán)隊實現(xiàn)的2.9 TOPS/W的能效相比，簡直相形見絀。

能效的大幅提升，對于AI計算的可持續(xù)發(fā)展，至關(guān)重要。

對此，Science表示：

通用人工智能（AGI）的飛速發(fā)展帶來了對下一代計算技術(shù)在性能和能效上的更高要求，而光子計算被認(rèn)為有望達(dá)到這些目標(biāo)。

但目前的光子集成電路，尤其是光學(xué)神經(jīng)網(wǎng)絡(luò)（ONN），在規(guī)模和計算能力上都非常有限，難以滿足現(xiàn)代AGI任務(wù)的需求。

來自清華的團(tuán)隊探索了一種新型的分布式衍射-干涉混合光子計算架構(gòu)，成功ONN的規(guī)模擴(kuò)展到了百萬神經(jīng)元級。他們在芯片上成功實現(xiàn)了一個擁有1396萬神經(jīng)元的ONN，能夠處理復(fù)雜的千類別級分類和AI生成內(nèi)容的任務(wù)。

可以說，這項研究是光子計算實際應(yīng)用的一個重要進(jìn)展，為各種AI應(yīng)用提供了支持。

創(chuàng)新性分布式計算架構(gòu)

根據(jù)論文介紹，清華團(tuán)隊為采用分布式計算的「太極」，構(gòu)建了一個深度較淺但寬度較廣的網(wǎng)絡(luò)結(jié)構(gòu)。

這種可重配置的衍射干涉混合光芯片，是實現(xiàn)多種先進(jìn)機(jī)器智能任務(wù)的關(guān)鍵組件，涵蓋了1000類別分類和內(nèi)容生成等應(yīng)用。

與傳統(tǒng)的深度計算層層堆疊的方法不同，「太極」將計算資源分配到多個獨立的集群中，為子任務(wù)單獨組織集群，最后為復(fù)雜的高級任務(wù)合成這些子任務(wù)。

具體地說，光學(xué)衍射層的完全連通特性，可以提供比傳統(tǒng)DNN中的卷積層更大的變形能力。

這表明光網(wǎng)絡(luò)具有用比電子系統(tǒng)更少的層來實現(xiàn)相同變換的潛力。

「太極」 的分布式架構(gòu)深度淺而寬，旨在以可持續(xù)和高效的方式擴(kuò)展計算能力。

在CIFAR-10數(shù)據(jù)集中，具有四個分布式層的「太極」 實現(xiàn)了與16層電子VGG-16網(wǎng)絡(luò)相當(dāng)?shù)木取?/span>

圖 1. 「太極」：一個配備分布式計算架構(gòu)的大規(guī)模光子芯片，專為百萬神經(jīng)元級芯片網(wǎng)絡(luò)模型設(shè)計

圖1（B）中展示了「太極」芯片，包括用于大規(guī)模輸入和輸出數(shù)據(jù)的雙衍射單元，以及用于可重構(gòu)特征嵌入和硬件多路復(fù)用的MZI陣列的可調(diào)矩陣乘法。

這些組件是「太極」（TEUs）的基本芯片上的執(zhí)行單元，利用了光學(xué)衍射和干涉的強(qiáng)大變形能力。

圖 1. 「太極」：一個配備分布式計算架構(gòu)的大規(guī)模光芯片，專為百萬神經(jīng)元級芯片網(wǎng)絡(luò)模型設(shè)計

接下來，再細(xì)看「太極」的設(shè)計結(jié)構(gòu)。

下圖A中展示了「太極」整體布局，分為三個部分：

1. 輸入衍射編碼器（DE）（藍(lán)色標(biāo)注）采用8×8光柵耦合器陣列進(jìn)行二維信息接收。總共對64個通道的輸入進(jìn)行了編碼，并將有效信息通過衍射調(diào)制權(quán)重壓縮為8個通道。

2. 干涉特征嵌入（IE）（紫色標(biāo)注）采用Mach-Zehnder調(diào)制器（MZM）陣列進(jìn)行任意矩陣乘法。

3. 相對于衍射解碼器，輸出繞射解碼器（DD）（藍(lán)色標(biāo)注）是反向的。

圖2（B）便是由20個DES、4個IE，以及4個DES被部署為新的TEU，來處理32×32的patch。

每個DE處理一個8×8的分布式patch，原始1024個通道的輸入數(shù)據(jù)被編碼為32個通道。

接下來的4個IE計算特征嵌入，最后4個DD將嵌入解碼為256個通道輸出。

通過調(diào)整分布式DE、IE和DD模塊的數(shù)量，形成不同的特征嵌入通道數(shù)量和輸出通道數(shù)量，可重構(gòu)和可擴(kuò)展的DE-IE-DD框架可以適應(yīng)不同的patch大小和任務(wù)難度。

圖2（C）展示了具有TEU群集的分布式架構(gòu)。圖2D中，研究者繪制了不同不確定性水平下的層數(shù)D和穩(wěn)健性Lip（F）之間的關(guān)系。

圖 2. 構(gòu)建「太極」的示意圖

（A）「太極」的執(zhí)行單元（TEUs）。

（B）多個TEUs根據(jù)計算分配協(xié)議協(xié)同工作，組成TEU集群。這些TEU集群采用滑動窗口機(jī)制處理較大的輸入數(shù)據(jù)。

（C）復(fù)雜任務(wù)被分解成多個簡單任務(wù)，每個簡單任務(wù)由一系列TEU集群（標(biāo)記為「路徑」）負(fù)責(zé)處理。

（D）理論性能分析表明，隨著每層網(wǎng)絡(luò)的錯誤率增加，理想的層數(shù)（深度）在物理系統(tǒng)中會減少。然而，采用多路徑的計算分配可以有效擴(kuò)展網(wǎng)絡(luò)規(guī)模，提升計算能力。

圖像分類，90%+準(zhǔn)確率

為了測試性能，研究人員首先取CIFAR-10數(shù)據(jù)集，并將每條路徑設(shè)置為6層。這是實際噪聲水平下的最佳規(guī)模，每層16-8-8-4-4-1 TEU。

七條路徑的二值化準(zhǔn)確率平均達(dá)到94%。

結(jié)合四條基本路徑的子結(jié)果，最終的準(zhǔn)確率達(dá)到了76.68%，已經(jīng)超過了現(xiàn)有的芯片架構(gòu)。

對于所有七條路徑，最終結(jié)果提高到93.65%，與目前流行的電子神經(jīng)網(wǎng)絡(luò)的性能相當(dāng)。

圖3（E）是七條路徑的整個測試集的混淆矩陣，圖3（B）列出了「太極」、傳統(tǒng)芯片網(wǎng)絡(luò)體系結(jié)構(gòu)、自由空間光計算體系結(jié)構(gòu)和電子對應(yīng)體系結(jié)構(gòu)之間的精度性能基準(zhǔn)。

圖3（D）則展示了額外的路徑如何幫助糾正錯誤的分類案例。

以青蛙圖像為例，將七條路徑的路徑輸出繪制為直方圖（基本路徑為紫線，額外路徑為藍(lán)線）。

在計算路徑輸出與每個類別的理想標(biāo)簽之間的相似度時，如果只采用基本路徑（即錯誤地將青蛙視為一艘船），則會做出錯誤的決定，但如果將所有七條路徑放在一起考慮，錯誤就會得到糾正。

圖 3. 用于1000類別分類的大規(guī)模光芯片

（A）CIFAR-10的多路徑二進(jìn)制標(biāo)簽，其中數(shù)據(jù)集中的每個對象在每條路徑上被標(biāo)記為「0」或「1」。單路徑（傳統(tǒng)方法）的分類準(zhǔn)確率有限，但多路徑（提議的方法）的分類準(zhǔn)確率隨參數(shù)數(shù)量增加而提高。

（B）對比傳統(tǒng)芯片上的光學(xué)、自由空間光學(xué)、基于電子的最先進(jìn)（SOTA）架構(gòu)以及「太極」在不同路徑數(shù)量下的CIFAR-10分類準(zhǔn)確率。

（C）層數(shù)對10類別分類準(zhǔn)確率的影響，展示了實驗數(shù)據(jù)（條形圖）和理論預(yù)測（曲線）。

（D）在CIFAR-10數(shù)據(jù)集中，一個樣本通過「太極」的路徑輸出顯示，最少的路徑數(shù)量可能導(dǎo)致錯誤判斷，但增加路徑數(shù)量可以糾正錯誤。

（E）使用七條路徑的CIFAR-10混淆矩陣。

（F）在mini-ImageNet數(shù)據(jù)集上進(jìn)行100類別分類任務(wù)的模擬（藍(lán)色）與實驗（紫色）結(jié)果。

（G）在Omniglot數(shù)據(jù)集上進(jìn)行1623類別分類任務(wù)的模擬（藍(lán)色）與實驗（紫色）結(jié)果。

為了進(jìn)一步挖掘「太極」的潛力，研究人員通過為更高級的任務(wù)部署更多路徑來擴(kuò)展規(guī)模。

在每條路徑中，層數(shù)保持不變，但每層將包含更多TEU（每層16-16-8-8-4-4-1 TEU）。

在100個類別的mini-ImageNet數(shù)據(jù)集上，每條路徑的平均二值準(zhǔn)確率在數(shù)值計算中為92.97%，在光學(xué)實驗中為88.05%。

在七條基本路徑和八條額外路徑的情況下，100個類別的總正確率在數(shù)值模擬中為92.76%，在實際芯片測試中為87.34%。

其中，圖3（F）是每個類別的正確樣本計數(shù)顯示為直方圖。

音樂家藝術(shù)家，全能模仿

研究人員將每個音符的生成視為一個分類問題，從47個可能的音調(diào)中進(jìn)行選擇，前后各有16個音符作為輸入。

對于訓(xùn)練，團(tuán)隊使用了接受率為95%的MCMC方法，來優(yōu)化生成的音樂片段的風(fēng)格。

隨著訓(xùn)練的進(jìn)行，網(wǎng)絡(luò)給出了一個在頻率（音高）域中的音符分布，來表示音樂風(fēng)格。

經(jīng)過訓(xùn)練，網(wǎng)絡(luò)中的參數(shù)被固定下來，以適應(yīng)巴赫的音樂生成風(fēng)格。

清華團(tuán)隊通過一個獨立訓(xùn)練的網(wǎng)絡(luò)對生成的結(jié)果進(jìn)行評估，該網(wǎng)絡(luò)給出了一個體現(xiàn)結(jié)果的巴赫風(fēng)格概率的「巴赫指數(shù)」輸出。

圖4（D）演示了生成過程。隨機(jī)噪聲作為初始輸入，其巴赫指數(shù)為6.61%。隨著迭代的進(jìn)行，音調(diào)圖中形成了模式，巴赫指數(shù)增加。

經(jīng)過500次迭代，生成結(jié)果的巴赫指數(shù)達(dá)到95.17%，具有典型的巴赫風(fēng)格。

在這種情況下，訓(xùn)練和生成被獨立地處理為總共4個聲音。

最后，「太極」創(chuàng)作了一個高度巴赫風(fēng)格的合成四聲合唱，如圖4（B）所示。

圖 4. 大規(guī)模光芯片用于多樣化內(nèi)容生成

（A）配備TEU集群的音樂生成網(wǎng)絡(luò)。

（B）展示了巴赫風(fēng)格原始音樂與生成的四聲部音高模式的對比。

（C）展示了生成的巴赫音樂的音符分布情況。

（D）使用巴赫指數(shù)進(jìn)行迭代音樂生成，該指數(shù)用來評估生成音樂與巴赫風(fēng)格的相似度。

再來看圖像生成，不同藝術(shù)家和風(fēng)格的圖像，被用來訓(xùn)練下一代神經(jīng)網(wǎng)絡(luò)。

研究人員采用不同的比例來生成不同級別的紋理，如下圖所示。

首先使用較大的Scale 1，生成粗紋理。而較小的Scale 2，然后用于生成精細(xì)紋理，從而獲得具有多尺度紋理的風(fēng)格化圖像。

為了評估結(jié)果，研究人員對預(yù)訓(xùn)練的VGG-16網(wǎng)絡(luò)進(jìn)行了微調(diào)，以得出藝術(shù)家風(fēng)格分類結(jié)果。

然后，作者在小圖像（來自MNIST數(shù)據(jù)集的手寫數(shù)字「4」）和大規(guī)模真實場景圖像下測試「太極」。

輸入的圖像是風(fēng)格化的，保留了場景中的對象形狀，并添加了藝術(shù)紋理。

圖 4. 大規(guī)模光芯片用于多樣化內(nèi)容生成

（E）配備TEU集群的圖像生成網(wǎng)絡(luò)。

（F）展示了三種不同藝術(shù)家風(fēng)格的圖像生成結(jié)果。輸入到「太極」的圖像包括帶有隨機(jī)噪聲的手寫數(shù)字「4」和真實場景，目的是生成符合指定藝術(shù)家風(fēng)格的風(fēng)格化圖像。使用一個獨立的分類網(wǎng)絡(luò)（風(fēng)格概率）來識別生成圖像的風(fēng)格。

此外，研究人員還進(jìn)行了字體風(fēng)格遷移的擴(kuò)展實驗，以進(jìn)一步展示「太極」 chiplets的高級內(nèi)容生成能力。

通過這些額外的實驗，他們驗證了「太極」不僅具有模仿藝術(shù)家風(fēng)格的能力，而且能夠從2D圖像中提取更高層次的語義信息。

成果討論

在這項工作中，團(tuán)隊設(shè)計了一種具有靈活分布式計算架構(gòu)的大規(guī)模衍射-干涉混合型光子AI芯片——「太極」。

在光芯片方面，「太極」深入探索了光子學(xué)的大規(guī)模并行連接，相較于其他TOPS/W級別框架，展現(xiàn)了更優(yōu)的計算效率。

未來，借助直接激光寫入（DLW）和相變材料（PCM），所有權(quán)重都能被重新配置，從而提升系統(tǒng)的靈活性。此外，芯片上的激光源、調(diào)制器和探測器也可以被整合到同一平臺上，并通過晶圓鍵合技術(shù)實現(xiàn)高級集成。

在分布式計算架構(gòu)方面，這種計算和任務(wù)分配方法不僅限于「太極」使用，還能幫助現(xiàn)有的光子集成電路（PIC）擴(kuò)展其處理更高級任務(wù)的能力。

在現(xiàn)代通用人工智能（AGI）領(lǐng)域，處理更復(fù)雜任務(wù)的趨勢是不可逆的。而「太極」展示了光子計算在處理多樣化復(fù)雜任務(wù)中的巨大潛力，使光學(xué)計算的實際應(yīng)用成為可能。

團(tuán)隊認(rèn)為，「太極」將加速更為強(qiáng)大的光學(xué)解決方案的開發(fā)，為基礎(chǔ)模型和新一代通用人工智能的發(fā)展提供關(guān)鍵支持。