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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

  圖像分類

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

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

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

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

  圖像聚類

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

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

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

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

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

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

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

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