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

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

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

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

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

  然而，找到一個同時滿足所有這些要求的張量核心處理器是具有挑戰(zhàn)性的。由于焦耳熱、電磁串擾和寄生電容，傳統(tǒng)的數(shù)字計算機難以滿足矩陣代數(shù)所需的速度和能效。相比于傳統(tǒng)的電子計算核心，光子張量核心具有更高的計算速度和更低的功耗。這是因為在光子計算中，光子的傳輸速度非?？欤夜庾油ǖ揽梢酝瑫r處理多個數(shù)據(jù)，從而實現(xiàn)真正的并行計算。此外，光子張量核心還可以通過復用光子通道來進一步提高計算效率和降低功耗。

  盡管如此，開發(fā)一個能夠進行大規(guī)模矩陣向量乘法的集成光子張量核心(IPTC)，其輸入和輸出大小可調(diào)，同時具有快速的權(quán)重更新，仍然是一個重大挑戰(zhàn)。

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

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

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

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

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

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

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

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

  圖像分類

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

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

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

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

  圖像聚類

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

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

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

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

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

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

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

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