清華首款AI光芯片登上Science，全球首創架構邁向AGI

新智元報道

編輯：桃子 好睏

【新智元導讀】巨耗算力大模型，離通往AGI目標又近了一步。清華團隊首創AI光芯片架構，研製全新「太極」實現了160 TOPS/W通用智能計算，能效竟是H100的1000倍。

訓練下一代萬億級參數大模型的高效芯片誕生了！

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

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

與傳統堆疊PIC芯片方法不同，清華團隊首創了分佈式廣度智能光計算架構，使得「太極」成爲全球首款大規模干涉衍射異構集成芯片。

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

它可以實現160 TOPS/W通用智能計算。

最新研究已於4月11日發表在Science期刊上。

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

更令人震驚的是，「太極」能效是英偉達H100的1000倍數。

研究人員表示，「太極」爲大規模的光子計算和高級任務鋪平了道路，進一步發掘了光子學在現代AGI中的靈活性和潛力。

ChatGPT耗電大有解了

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

那是因爲，基於Transformer架構的大模型，通過token預測完成推理，需要消耗大量的算力。

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

若是能夠發明一種，節省大量能耗的芯片，LLM的性能或在未來實現更大的飛昇。

而「太極」可能會使通用人工智能（AGI）成爲現實。研究人員表示，

在將計算能力提升到AGI所需的水平方面，「太極」的模塊化設計可能是一個關鍵優勢。

清華團隊設計了一個擁有1396萬個人工神經元的分佈式「太極」網絡，超越了其他光芯片設計（147萬個神經元）。

因此，「太極」實現了160.82 TOPS/W的能效。

與2022年一個團隊實現的2.9 TOPS/W的能效相比，簡直相形見絀。

能效的大幅提升，對於AI計算的可持續發展，至關重要。

對此，Science表示：

創新性分佈式計算架構

根據論文介紹，清華團隊爲採用分佈式計算的「太極」，構建了一個深度較淺但寬度較廣的網絡結構。

這種可重配置的衍射干涉混合光芯片，是實現多種先進機器智能任務的關鍵組件，涵蓋了1000類別分類和內容生成等應用。

與傳統的深度計算層層堆疊的方法不同，「太極」將計算資源分配到多個獨立的集羣中，爲子任務單獨組織集羣，最後爲複雜的高級任務合成這些子任務。

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

這表明光網絡具有用比電子系統更少的層來實現相同變換的潛力。

「太極」 的分佈式架構深度淺而寬，旨在以可持續和高效的方式擴展計算能力。

在CIFAR-10數據集中，具有四個分佈式層的「太極」 實現了與16層電子VGG-16網絡相當的精度。

圖1. 「太極」：一個配備分佈式計算架構的大規模光子芯片，專爲百萬神經元級芯片網絡模型設計

圖1（B）中展示了「太極」芯片，包括用於大規模輸入和輸出數據的雙衍射單元，以及用於可重構特徵嵌入和硬件多路複用的MZI陣列的可調矩陣乘法。

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

圖1. 「太極」：一個配備分佈式計算架構的大規模光芯片，專爲百萬神經元級芯片網絡模型設計

接下來，再細看「太極」的設計結構。

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

1. 輸入衍射編碼器（DE）（藍色標註）採用8×8光柵耦合器陣列進行二維信息接收。總共對64個通道的輸入進行了編碼，並將有效信息通過衍射調製權重壓縮爲8個通道。

2. 干涉特徵嵌入（IE）（紫色標註）採用Mach-Zehnder調製器（MZM）陣列進行任意矩陣乘法。

3. 相對於衍射解碼器，輸出繞射解碼器（DD）（藍色標註）是反向的。

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

每個DE處理一個8×8的分佈式patch，原始1024個通道的輸入數據被編碼爲32個通道。

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

通過調整分佈式DE、IE和DD模塊的數量，形成不同的特徵嵌入通道數量和輸出通道數量，可重構和可擴展的DE-IE-DD框架可以適應不同的patch大小和任務難度。

圖2（C）展示了具有TEU羣集的分佈式架構。圖2D中，研究者繪製了不同不確定性水平下的層數D和穩健性Lip（F）之間的關係。

圖2. 構建「太極」的示意圖

（A）「太極」的執行單元（TEUs）。

（B）多個TEUs根據計算分配協議協同工作，組成TEU集羣。這些TEU集羣採用滑動窗口機制處理較大的輸入數據。

（C）複雜任務被分解成多個簡單任務，每個簡單任務由一系列TEU集羣（標記爲「路徑」）負責處理。

（D）理論性能分析表明，隨着每層網絡的錯誤率增加，理想的層數（深度）在物理系統中會減少。然而，採用多路徑的計算分配可以有效擴展網絡規模，提升計算能力。

圖像分類，90%+準確率

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

七條路徑的二值化準確率平均達到94%。

結合四條基本路徑的子結果，最終的準確率達到了76.68%，已經超過了現有的芯片架構。

對於所有七條路徑，最終結果提高到93.65%，與目前流行的電子神經網絡的性能相當。

圖3（E）是七條路徑的整個測試集的混淆矩陣，圖3（B）列出了「太極」、傳統芯片網絡體系結構、自由空間光計算體系結構和電子對應體系結構之間的精度性能基準。

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

以青蛙圖像爲例，將七條路徑的路徑輸出繪製爲直方圖（基本路徑爲紫線，額外路徑爲藍線）。

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

圖3. 用於1000類別分類的大規模光芯片

（A）CIFAR-10的多路徑二進制標籤，其中數據集中的每個對象在每條路徑上被標記爲「0」或「1」。單路徑（傳統方法）的分類準確率有限，但多路徑（提議的方法）的分類準確率隨參數數量增加而提高。

（B）對比傳統芯片上的光學、自由空間光學、基於電子的最先進（SOTA）架構以及「太極」在不同路徑數量下的CIFAR-10分類準確率。

（C）層數對10類別分類準確率的影響，展示了實驗數據（條形圖）和理論預測（曲線）。

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

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

（F）在mini-ImageNet數據集上進行100類別分類任務的模擬（藍色）與實驗（紫色）結果。

（G）在Omniglot數據集上進行1623類別分類任務的模擬（藍色）與實驗（紫色）結果。

爲了進一步挖掘「太極」的潛力，研究人員通過爲更高級的任務部署更多路徑來擴展規模。

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

在100個類別的mini-ImageNet數據集上，每條路徑的平均二值準確率在數值計算中爲92.97%，在光學實驗中爲88.05%。

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

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

音樂家藝術家，全能模仿

研究人員將每個音符的生成視爲一個分類問題，從47個可能的音調中進行選擇，前後各有16個音符作爲輸入。

對於訓練，團隊使用了接受率爲95%的MCMC方法，來優化生成的音樂片段的風格。

隨着訓練的進行，網絡給出了一個在頻率（音高）域中的音符分佈，來表示音樂風格。

經過訓練，網絡中的參數被固定下來，以適應巴赫的音樂生成風格。

清華團隊通過一個獨立訓練的網絡對生成的結果進行評估，該網絡給出了一個體現結果的巴赫風格概率的「巴赫指數」輸出。

圖4（D）演示了生成過程。隨機噪聲作爲初始輸入，其巴赫指數爲6.61%。隨着迭代的進行，音調圖中形成了模式，巴赫指數增加。

經過500次迭代，生成結果的巴赫指數達到95.17%，具有典型的巴赫風格。

在這種情況下，訓練和生成被獨立地處理爲總共4個聲音。

最後，「太極」創作了一個高度巴赫風格的合成四聲合唱，如圖4（B）所示。

圖4. 大規模光芯片用於多樣化內容生成

（A）配備TEU集羣的音樂生成網絡。

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

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

（D）使用巴赫指數進行迭代音樂生成，該指數用來評估生成音樂與巴赫風格的相似度。

再來看圖像生成，不同藝術家和風格的圖像，被用來訓練下一代神經網絡。

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

首先使用較大的Scale 1，生成粗紋理。而較小的Scale 2，然後用於生成精細紋理，從而獲得具有多尺度紋理的風格化圖像。

爲了評估結果，研究人員對預訓練的VGG-16網絡進行了微調，以得出藝術家風格分類結果。

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

輸入的圖像是風格化的，保留了場景中的對象形狀，並添加了藝術紋理。

圖4. 大規模光芯片用於多樣化內容生成

（E）配備TEU集羣的圖像生成網絡。

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

此外，研究人員還進行了字體風格遷移的擴展實驗，以進一步展示「太極」 chiplets的高級內容生成能力。

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

成果討論

在這項工作中，團隊設計了一種具有靈活分佈式計算架構的大規模衍射-干涉混合型光子AI芯片——「太極」。

在光芯片方面，「太極」深入探索了光子學的大規模並行連接，相較於其他TOPS/W級別框架，展現了更優的計算效率。

未來，藉助直接激光寫入（DLW）和相變材料（PCM），所有權重都能被重新配置，從而提升系統的靈活性。此外，芯片上的激光源、調製器和探測器也可以被整合到同一平臺上，並通過晶圓鍵合技術實現高級集成。

在分佈式計算架構方面，這種計算和任務分配方法不僅限於「太極」使用，還能幫助現有的光子集成電路（PIC）擴展其處理更高級任務的能力。

在現代通用人工智能（AGI）領域，處理更復雜任務的趨勢是不可逆的。而「太極」展示了光子計算在處理多樣化複雜任務中的巨大潛力，使光學計算的實際應用成爲可能。

團隊認爲，「太極」將加速更爲強大的光學解決方案的開發，爲基礎模型和新一代通用人工智能的發展提供關鍵支持。

參考資料：

https://www.science.org/doi/10.1126/science.adl1203

https://www.tsinghua.edu.cn/info/1175/110690.htm