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

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

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

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

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

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

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

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

  這種完全集成的處理器僅包括兩個(gè)TFLN調(diào)制器、一個(gè)III-V激光器和一個(gè)電荷積分光接收器。此處理器可以以高計(jì)算速度執(zhí)行大規(guī)模矩陣向量乘法，如圖1c所示。通過調(diào)整電荷積分光接收器的積分時(shí)間，可以靈活地修改矩陣矢量乘法的扇形尺寸。利用TFLN調(diào)制器的高調(diào)制速度和電荷積分光接收器的快速累積操作，張量核心實(shí)現(xiàn)了120 GOPS的計(jì)算速度。此外，該處理器的張量核心具有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芯片、倒裝芯片光電探測(cè)器、調(diào)制器行波電極和激光器的放大顯微照片。該團(tuán)隊(duì)使用倒裝芯片鍵合技術(shù)，在平衡檢測(cè)方案中，將兩個(gè)光電探測(cè)器(標(biāo)記為PD1和PD2)固定在兩個(gè)光柵耦合器上方，如圖2c所示。激光器和TFLN芯片使用光子引線鍵合連接，其形狀可以調(diào)整以匹配波導(dǎo)刻面的實(shí)際位置(見圖2e)。如圖2c右側(cè)所示，該研究團(tuán)隊(duì)還通過光子引線鍵合將TFLN芯片與光纖陣列連接，用于校準(zhǔn)偏置電壓和延遲時(shí)間，并協(xié)助涉及兩個(gè)負(fù)數(shù)的乘法。圖2f顯示了TFLN芯片、激光器和光電探測(cè)器的相對(duì)高度。

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

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

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

  圖像分類

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

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

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

  500幅圖像的混淆矩陣(圖4d，e)顯示，生成的預(yù)測(cè)準(zhǔn)確率為91.8%，而CPU計(jì)算的數(shù)值結(jié)果準(zhǔn)確率為92%。該團(tuán)隊(duì)所設(shè)計(jì)的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)建這樣一個(gè)數(shù)據(jù)庫(kù)是一項(xiàng)耗時(shí)且資源密集的任務(wù)，可能并不總是可行的。相比之下，無(wú)監(jiān)督學(xué)習(xí)可以對(duì)未標(biāo)記的數(shù)據(jù)進(jìn)行操作，以發(fā)現(xiàn)其底層結(jié)構(gòu)，為提取數(shù)據(jù)特征提供了一種替代方法。

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

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

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

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

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

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

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