用于推理和原位訓練的薄膜鈮酸鋰中的120GOPS光子張量核心

  近日，中山大學(xué)研究團隊與不列顛哥倫比亞大學(xué)、女王大學(xué)、拉瓦爾大學(xué)、加拿大國家研究院、清華大學(xué)團隊攜手，提出了一種完全集成的光子張量核心，僅由兩個(gè)薄膜鈮酸鋰(TFLN)調制器、一個(gè)III-V激光器和一個(gè)電荷積分光接收器組成。該光子張量核心能夠以120GOPS的計算速度實(shí)現整個(gè)神經(jīng)網(wǎng)絡(luò )層，同時(shí)還允許靈活調整輸入和輸出的數量。相關(guān)結果以“120 GOPS Photonic tensor core in thin-film lithium niobate for inference and in situ training”為題在Nature Communications上發(fā)表。該張量核心支持快速原位訓練，權重更新速度為60GHz。它通過(guò)原位訓練成功地對112×112像素的圖像進(jìn)行了分類(lèi)和聚類(lèi)。此外，該團隊研究人員首次在光上實(shí)現了負數與負數的乘法，為聚類(lèi)AI任務(wù)的訓練提供了新的方案。

文章鏈接：https://doi.org/10.1038/s41467-024-53261-x

  光子張量核心(Photonic Tensor Core)是一種新型的并行計算核心，它利用光子進(jìn)行計算，具有極高的計算速度和并行處理能力。在光子張量核心中，并行卷積處理是一種重要的計算模式，它可以極大地提高卷積神經(jīng)網(wǎng)絡(luò )的計算效率和性能。其基本原理是利用光子的高速傳輸特性和并行處理能力，將卷積運算分布到多個(gè)光子通道中進(jìn)行，從而實(shí)現并行計算。具體來(lái)說(shuō)，光子張量核心將輸入數據和卷積核分別分配到不同的光子通道中，每個(gè)光子通道獨立進(jìn)行卷積運算，并將結果匯總起來(lái)得到最終的卷積結果。

  近幾年，人工智能(AI)正越來(lái)越多地被整合到各個(gè)領(lǐng)域，包括自動(dòng)駕駛汽車(chē)、智能建筑和智能工廠(chǎng)，如下圖1a所示。人工智能系統的核心是張量核心處理器，該核心處理器將表現出幾個(gè)關(guān)鍵特征：首先是高速、大規模矩陣向量乘法;其次是快速更新體重，加快訓練速度，促進(jìn)“即時(shí)”或在線(xiàn)學(xué)習，這對自動(dòng)駕駛汽車(chē)等應用尤其有益;最后是低能耗和緊湊的外形。

  然而，找到一個(gè)同時(shí)滿(mǎn)足所有這些要求的張量核心處理器是具有挑戰性的。由于焦耳熱、電磁串擾和寄生電容，傳統的數字計算機難以滿(mǎn)足矩陣代數所需的速度和能效。相比于傳統的電子計算核心，光子張量核心具有更高的計算速度和更低的功耗。這是因為在光子計算中，光子的傳輸速度非?？?，而且光子通道可以同時(shí)處理多個(gè)數據，從而實(shí)現真正的并行計算。此外，光子張量核心還可以通過(guò)復用光子通道來(lái)進(jìn)一步提高計算效率和降低功耗。

  盡管如此，開(kāi)發(fā)一個(gè)能夠進(jìn)行大規模矩陣向量乘法的集成光子張量核心(IPTC)，其輸入和輸出大小可調，同時(shí)具有快速的權重更新，仍然是一個(gè)重大挑戰。

  針對上述問(wèn)題，該團隊基于TFLN調制器在寬波長(cháng)范圍內工作的能力，介紹了一種集成了薄膜鈮酸鋰(TFLN)光子學(xué)和電荷積分光接收器的IPTC(圖1b)，此TFLN芯片是使用由一個(gè)360 nm厚、x切割、y傳播的LN薄膜組成，該薄膜位于一個(gè)500μm厚的石英手柄上，兩者之間有一個(gè)2μmSiO2層。使用光學(xué)光刻對光學(xué)器件進(jìn)行圖案化，并使用電感耦合等離子體進(jìn)行蝕刻。然后，在光學(xué)器件的頂部沉積一層1μm厚的SiO2包覆層。然后用剝離工藝對金和加熱器電極進(jìn)行圖案化。

  這種完全集成的處理器僅包括兩個(gè)TFLN調制器、一個(gè)III-V激光器和一個(gè)電荷積分光接收器。此處理器可以以高計算速度執行大規模矩陣向量乘法，如圖1c所示。通過(guò)調整電荷積分光接收器的積分時(shí)間，可以靈活地修改矩陣矢量乘法的扇形尺寸。利用TFLN調制器的高調制速度和電荷積分光接收器的快速累積操作，張量核心實(shí)現了120 GOPS的計算速度。此外，該處理器的張量核心具有60 GHz的權重更新速度，可以實(shí)現快速的原位訓練。

  圖1 集成光子張量核心(IPTC)的概念。a頂部：人工智能(AI)AI系統的應用和功能要求處理器能夠適應各種AI任務(wù)，底部：IPTC的示意圖。b基于傳統波分復用(WDM)的IPTC的示意圖。c設備的性能與幾種最先進(jìn)的光子張量核心的性能的比較

  圖2a顯示了該設備的原型照片。此外，圖2b–e分別提供了所制造的TFLN芯片、倒裝芯片光電探測器、調制器行波電極和激光器的放大顯微照片。該團隊使用倒裝芯片鍵合技術(shù)，在平衡檢測方案中，將兩個(gè)光電探測器(標記為PD1和PD2)固定在兩個(gè)光柵耦合器上方，如圖2c所示。激光器和TFLN芯片使用光子引線(xiàn)鍵合連接，其形狀可以調整以匹配波導刻面的實(shí)際位置(見(jiàn)圖2e)。如圖2c右側所示，該研究團隊還通過(guò)光子引線(xiàn)鍵合將TFLN芯片與光纖陣列連接，用于校準偏置電壓和延遲時(shí)間，并協(xié)助涉及兩個(gè)負數的乘法。圖2f顯示了TFLN芯片、激光器和光電探測器的相對高度。

  圖2g顯示了從波長(cháng)為1307.22 nm的激光器耦合到TFLN芯片的光的光電流-電壓(L-I-V)曲線(xiàn)。由于周期性電容加載行波電極(見(jiàn)圖2d)，因此，該團隊所使用調制器的3-dB電光帶寬比60 GHz更寬(見(jiàn)圖2h)。對于恒定的輸入光功率，積分器的輸出電壓隨積分時(shí)間線(xiàn)性增加(見(jiàn)圖2i)。在平衡檢測方案中，當PD1接收到的光功率低于PD2接收到的功率時(shí)，積分器的輸出電壓變化為正，當它高于PD2接收的光功率時(shí)，集成商的輸出電壓波動(dòng)為負。這意味著(zhù)所提出的光接收器可以在矩陣向量乘法中執行加法和減法運算。

  圖2 封裝設備的原型。a整個(gè)設備的照片。b混合集成芯片的顯微照片。c–e分別是倒裝芯片光電探測器(PD)、調制器行波電極和激光器的放大顯微照片。f設備側視顯微照片。g從激光器耦合到TFLN芯片中的光的光電流-電壓曲線(xiàn)。h調制器的電光帶寬(S21參數)。i當輸入光功率固定在一定值時(shí)，光接收器的輸出電壓隨積分時(shí)間而變化。

  該研究團隊通過(guò)利用Python控制所有的設備，在兩個(gè)向量之間執行點(diǎn)積運算，圖3a顯示了通過(guò)設備的數據流示意圖。通過(guò)隨機改變兩個(gè)矢量，使用設備記錄了3780個(gè)光子點(diǎn)積測量值。每個(gè)矢量的維度設置為131072，這兩個(gè)矢量分別由兩個(gè)調制器以60 GB的調制率進(jìn)行調制，從而實(shí)現了120 GOPS的計算速度和60 GHz的權重更新速度。最初校準了兩個(gè)向量之間的時(shí)間延遲，以確保第一個(gè)向量的每個(gè)元素都能正確地乘以第二個(gè)向量的對應元素。測量的輸出電壓(即點(diǎn)積結果)在-1和+1之間縮放，作為預期點(diǎn)積結果的函數，如圖3b所示。與預期的點(diǎn)積結果相比，測量結果的誤差標準偏差為0.03(6.04位)，超過(guò)了執行AI任務(wù)所需的4位精度。

  圖3 使用設備進(jìn)行點(diǎn)積操作的實(shí)驗結果。a設備工作原理示意圖。b設備以120GOPS的計算速度在兩個(gè)131072維向量之間執行點(diǎn)積運算的結果。

  圖像分類(lèi)

  該團隊構建了一個(gè)多層感知器(見(jiàn)圖4a)，并在大規模手寫(xiě)數字數據庫上對其進(jìn)行了測試。每個(gè)手寫(xiě)數字圖像有112×112像素，被展平成12544×1的向量作為第一層的輸入。第一和第二隱藏層中的節點(diǎn)數量分別設置為70和300，并且泄漏ReLU函數用于非線(xiàn)性激活函數。

  圖4 設備對手寫(xiě)數字進(jìn)行分類(lèi)的結果。a多層感知器神經(jīng)網(wǎng)絡(luò )的框圖。b現場(chǎng)訓練示意圖。c與僅在中央處理單元(CPU，藍色虛線(xiàn))上運行的方案相比，現場(chǎng)訓練(實(shí)線(xiàn))方案的驗證精度隨時(shí)間變化。d、 e使用大規模數據庫從理論上計算混淆矩陣和實(shí)驗混淆矩陣。

  分類(lèi)是一項監督學(xué)習AI任務(wù)，需要標記數據來(lái)訓練模型。該團隊構建的多層感知器模型使用IPTC執行前向傳播的原位訓練方案(見(jiàn)圖4b)用2000個(gè)標記的數字圖像進(jìn)行訓練。同時(shí)，電子計算機處理非線(xiàn)性函數和反向傳播。權重向量通過(guò)隨機梯度下降法進(jìn)行更新，允許迭代訓練單個(gè)樣本。重復從前向傳播到反向傳播的訓練過(guò)程，直到收斂。圖4c顯示了與僅在中央處理單元(CPU)上運行相比，現場(chǎng)訓練方案的驗證精度隨時(shí)間的變化。

  500幅圖像的混淆矩陣(圖4d，e)顯示，生成的預測準確率為91.8%，而CPU計算的數值結果準確率為92%。該團隊所設計的IPTC達到了接近理論的精度，表明原位訓練方案使系統能夠固有地考慮硬件的非理想性，包括制造變化和噪聲。

  圖像聚類(lèi)

  監督學(xué)習可以成功地解決現實(shí)世界的挑戰，但它也有一些缺點(diǎn)。主要局限之一是，它需要大量準確標記的數據來(lái)訓練模型。創(chuàng )建這樣一個(gè)數據庫是一項耗時(shí)且資源密集的任務(wù)，可能并不總是可行的。相比之下，無(wú)監督學(xué)習可以對未標記的數據進(jìn)行操作，以發(fā)現其底層結構，為提取數據特征提供了一種替代方法。

  該團隊通過(guò)利用主成分分析對大規模手寫(xiě)數字進(jìn)行聚類(lèi)，展示了該團隊研發(fā)設備在無(wú)監督學(xué)習AI任務(wù)中的潛力，主成分分析是最常用的無(wú)監督學(xué)習模型之一。主成分分析通過(guò)將高維數據幾何投影到有限數量的主成分(PC)上，以獲得數據的最佳摘要，從而簡(jiǎn)化了高維數據。該團對發(fā)研發(fā)設備的收斂速度與CPU的收斂速度相當(見(jiàn)圖5b)。

  圖5 使用設備對手寫(xiě)數字進(jìn)行聚類(lèi)的結果。a設備工作原理示意圖。c、 d分別是基于投影到前三個(gè)主部件(PC)上的分數的每個(gè)手寫(xiě)指針的3D坐標的前視圖和后視圖。

  此外，為了使用設備可視化手寫(xiě)數字的聚類(lèi)結果，圖5c和d顯示了PC1-PC3上的投影，占特征的28.7%。盡管只使用了前三臺PC，但未標記的手寫(xiě)數字仍然可以很好地聚類(lèi)。

  此外，該團隊為了展示解決方案的可擴展性，提出了一種端到端的光子神經(jīng)網(wǎng)絡(luò )，該網(wǎng)絡(luò )結合了TDM和WDM方法的優(yōu)點(diǎn)，如圖6所示。該網(wǎng)絡(luò )能夠同時(shí)執行多個(gè)AI任務(wù)，從輸入層到輸出層，延遲為納秒，所有這些都不依賴(lài)于數字處理器的幫助。

  圖6 一個(gè)光子神經(jīng)網(wǎng)絡(luò )的示意圖，旨在顯示所提出的集成光子張量核心的可擴展性，該網(wǎng)絡(luò )采用時(shí)分復用(TDM)和波分復用(WDM)相結合的混合方法。

  總之，該團隊已經(jīng)通過(guò)實(shí)驗證明，研發(fā)的IPTC可以執行大規模的矩陣向量乘法，具有靈活可調的扇入和扇出尺寸，并有助于快速更新權重。此IPTC具有處理兩個(gè)負數之間乘法能力的開(kāi)創(chuàng )性，能夠通過(guò)現場(chǎng)訓練處理監督和無(wú)監督學(xué)習AI任務(wù)。

  我們身處于一個(gè)數據、信息量爆炸性增長(cháng)的時(shí)代，一個(gè)由人工智能(AI)引領(lǐng)的、更加智能的時(shí)代。但是，持續增加的數據量在為 AI 提供源源不斷的“動(dòng)力”的同時(shí)，也對用于 AI 的電子計算硬件提出了更多的挑戰，無(wú)論是在計算速度，還是在功耗方面，都已經(jīng)成為嚴重制約 AI 發(fā)展的主要瓶頸之一。隨著(zhù) AI 的興起，傳統的電子計算方法逐漸達到了其性能極限，并且滯后于可處理數據的快速增長(cháng)。在各種類(lèi)型的 AI 系統中，人工神經(jīng)網(wǎng)絡(luò )由于優(yōu)異的性能而被廣泛應用于 AI 任務(wù)，這些網(wǎng)絡(luò )使用多層相互連接的人工神經(jīng)網(wǎng)絡(luò )進(jìn)行復雜的數學(xué)運算。為了加速人工神經(jīng)網(wǎng)絡(luò )的處理，人們已經(jīng)做出了各種努力來(lái)設計和實(shí)現特定的計算系統，通過(guò)將電子電路和數千或數百萬(wàn)個(gè)光子處理器集成到一個(gè)合適的架構中，一種同時(shí)利用光子和電子處理器的混合光電框架，或許在不久的將來(lái)可以徹底改變 AI 硬件。未來(lái)，這種硬件將在通信、數據中心營(yíng)運和云計算等領(lǐng)域具有十分重要的應用。