編輯 | 蘿蔔皮
光子神經網絡使用光子而不是電子執行受大腦啟發的計算，以實現顯著提高的計算性能。然而，現有架構只能處理具有規則結構的數據，而無法泛化到歐幾里得空間之外的圖結構數據。
清華大學的研究人員提出了衍射圖神經網絡 (DGNN)，這是一種基於衍射光子計算單元 (DPU) 和片上光學器件的全光學圖表示學習架構，以解決這一限制。具體來說，圖節點屬性被編碼到帶狀光波導中，由 DPU 轉換，並由光耦合器聚合以提取其特徵表示。
DGNN 在光速光學消息通過圖結構期間，捕獲節點鄰域之間的複雜依賴關係。他們通過基準數據庫演示了 DGNN 在節點和圖級分類任務中的應用，並獲得了卓越的性能。該研究團隊的工作為使用深度學習設計用於高效處理大規模圖形數據結構的專用集成光子電路開闢了一個新方向。
該研究以「All-optical graph representation learning using integrated diffractive photonic comPuting units」為題，於 2022 年 6 月 15 日發佈在《Science Advances》。
深度學習技術在廣泛的人工智能（AI）應用中取得了巨大進步，包括計算機視覺、語音識別、自然語言處理、自動駕駛汽車、生物醫學科學等。其核心是利用多層神經網絡從大數據中學習層次和複雜的抽象，這是由集成電子計算平台（如中央處理器、圖形處理單元（GPU）、張量處理單元和現場可編程門陣列）的不斷發展推動的。
然而，電子計算性能已接近物理極限，難以跟上人工智能發展需求的增長，這是需要大規模深度神經模型的廣泛應用中的普遍困境。
近年來，光子計算的研究越來越受到關注，以光子為計算介質，利用其高並行性、最低功耗和光速信號處理的先進特性來構建光子神經網絡。已經提出了許多光子神經網絡架構來促進複雜的神經啟發計算，例如衍射神經網絡、光學干涉神經網絡、光子尖峰神經網絡和光子儲存計算。
現有架構在處理具有矢量或網格狀圖像形式的規則結構的數據方面最為成功。然而，各種科學領域分析的數據超出了這種基本的歐幾里得領域。作為典型代表，編碼複雜系統中實體（即節點）之間豐富關係（即邊）的圖結構數據在現實世界中無處不在，小到化學分子，大到大腦網絡。
為了處理圖結構數據，圖神經網絡（GNN）已經發展成為一類廣泛的新方法，能夠集成局部節點特徵和圖拓撲結構進行表示學習。
在這些模型中，基於消息傳遞的 GNN 具有靈活性和效率的主要優勢，它通過在圖節點處生成神經消息並將邊緣傳遞給其鄰居進行特徵更新。它已成功應用於許多基於圖的應用，包括分子屬性預測、藥物發現、基於骨架的人體動作識別、時空預測等。然而，如何有效地利用光子計算來使基於圖的深度學習受益仍然很大程度上尚未探索。
在這裡，研究人員提出了衍射 GNN (DGNN)，這是一種新穎的光子 GNN 架構，可以在圖形結構數據上執行光學消息傳遞。DGNN 建立在用於生成光學節點特徵的集成衍射光子計算單元 (DPU) 的基礎之上。每個 DPU 包括用金屬線實現的連續衍射層，以將節點屬性轉換為光學神經信息，其中部署條形光波導以對輸入節點屬性進行編碼並輸出轉換結果。從節點鄰域發送的光學神經信息使用光耦合器聚合。
在 DGNN 架構中，DPU 可以水平級聯以擴大感受野，從而從任意大小的相鄰節點中捕獲複雜的依賴關係。此外，DPU 還可以垂直堆疊以提取更高維的光學節點特徵，以提高其學習能力，這受到許多現代深度學習模型（例如 Transformer 和圖注意力網絡）中使用的多頭策略的啟發。
在這種可擴展的光學消息傳遞方案的基礎上，研究人員首先演示了半監督節點分類任務，其中 DGNN 提取的光學節點特徵被輸入光學或電子輸出分類器以確定節點類別。結果表明，該光學 DGNN 在合成圖模型和三個真實世界圖基準數據集上的電子 GNN 方面實現了具有競爭力甚至更出色的分類性能。
此外，DGNN 還支持圖級分類，其中額外的 DPU 用於將全光節點特徵聚合為圖級表示進行分類。基於骨架的人類動作識別的結果證明了該架構對圖分類任務的有效性。
圖示：光學 DGNN 的架構。（來源：論文）
稀缺的訓練標籤
研究人員分析了 DGNN 在訓練標籤大小有限的情況下的有效性，這是半監督學習中的常見情況。在相同的架構設置下，他們比較了 DGNN 在不同尺寸訓練標籤下電子模型基線的性能，包括每類 1、5、10、15、20 和 25 個標籤。
通過對每個尺寸的訓練標籤進行 10 次評估，繪製了帶有誤差條的測試精度條形圖。對衍射調製層進行二值化有助於克服網絡訓練期間的局部最小問題，並提高分類精度。DGNN 架構優於所有標籤稀缺設置的所有基線，特別是在較小的訓練集大小時，例如每個類只有一個標籤，這表明相對於其他電子計算方法具有更高的泛化能力。
圖示：具有稀缺訓練標籤的 Amazon Photo 上的分類。（來源：論文）
帶有錐形輸出波導的 DPU
錐形波導用於將較大區域的輸出光場耦合到集成 DPU 的輸出端口。錐形輸出波導的改進耦合效率使光電探測器能夠接收更多的光功率並提高光電轉換過程中的信噪比 (SNR)。更高的信噪比為分類器提供了更高質量的輸入信號，保證了分類任務的穩定性。
錐形和單模波導的輸出能量分佈和模型性能的定量評估，研究人員使用 FDTD 來評估合成 SBM 圖測試節點上光學特徵的功率分佈和訓練好的 DGNN-E 模型。錐形輸出波導的起始纖芯寬度經過優化，設置為 2 μm，而不是單模輸出波導中使用的 500 nm。對於每個測試圖節點，通過計算兩個端口的輸出功率相對於輸入光源功率的比例得到DPU的功率傳輸率，從而得到所有圖節點上傳輸率的頻率直方圖。具有錐形輸出波導的 DPU 的平均功率傳輸率為 2.01%，比單模輸出波導的 0.36% 高約 5.6 倍。
利用 DPU 的估計功率傳輸率，研究人員評估了片上光電探測器的光電流 SNR，公式在材料和方法中詳述，在不同的輸入光源功率下。同時，進一步評估了 DGNN-E 模型在 SNR 方面的測試精度，在 top-k 相鄰節點設置為 16 下，通過將光電探測器噪聲包括到節點特徵中並重新訓練電子分類器（圖 S6D）。增加輸入光源功率和 DPU 的功率傳輸率提高了光電流 SNR，並在合成 SBM 圖上實現了更穩定的模型性能。
圖示：三個基準圖數據庫上的半監督節點分類。（來源：論文）
在這項工作中，採用單輪消息傳遞的 PPRGo 模型直接捕獲高階鄰域信息。DGNN 的計算能量效率是在輸入光源功率為 10 mW 的基礎上計算的，在錐形和單模輸出波導中，DGNN 的光電流信噪比分別達到 34.6 和 20.2 dB，相應的模型測試精度分別為 94.4%和 92.3%。
DPU的計算精度
量化比特決定了 DPU 的計算精度，可以從光電流 SNR 中推斷出來。在數字信號處理中，在模數轉換器的量化過程中引入了量化誤差。假設信號具有覆蓋所有量化級別的均勻分佈，則信量化噪聲比可以表示為 SQNR = 20log10(2Q) = 6.02Q dB，其中 Q 表示量化位數。因此，使用具有 10 mW 輸入光源功率的錐形輸出波導的 34.6 dB 光電流 SNR 對應於約 6 個量化位。
計算密度和能源效率
值得注意的是，一旦 DGNN 架構設計經過優化和物理製造，用於計算節點和圖形表示的片上光學器件以及推理過程中的光學輸出分類器都是無源的。這種基於圖的 AI 任務的推理過程以光速處理，僅受輸入數據調製和輸出檢測率的限制，與電子 GNN 相比消耗的能量很少。
具體而言，假設 DGNN 使用 MSG（·）將每個節點的 n 維屬性轉換為 m 維光學神經信息，使用 AGG（·）聚合 k 個節點的光學特徵，並為 C 類分類任務堆疊 P 個頭部。因此，每個節點的 MSG（·）模塊包含每個節點的 n×m 權重矩陣，每個頭部的 AGG（·）模塊包含 m 維向量的 k 個節點之和，分類器包含 mP×C 權重矩陣。
因此，DGNN 的每個推理周期包含用於特徵提取的 (2nmk + mk)P 操作 (OP) 和用於分類的 2mPC 操作，即具有 (2nk + k + 2C)mP 的總操作。
圖示：DGNN 在動作識別任務上的圖分類。（來源：論文）
考慮到基於現有硅光子代工廠的 30GHz 數據調製和光電探測速率，DGNN 的計算速度為 (6nk + 3k + 6C)mP × 10^10 OPs/s。假設典型光源功率為 10 mW，DGNN 的能量效率為 (6nk + 3k + 6C)mP × 10^12 OPs/J。
對於n = 20，m = 2，k = 8，P = 4，C = 8的節點分類設置，計算速度為 82.6 TOP/s（Tera-Operations/s），能效為 8.26 POP/s（ Peta-Operations/s) 每瓦。對於下圖中計算區域大小為 61.5 μm x 45 μm 的 DPU 模塊，使用 3×2 權重矩陣執行 MSG(·) 函數，計算密度為每平方毫米 130 TOP/s。
假設每個 MZI 的大小為 100 μm x 100 μm，則使用片上 MZI 光子器件的相同 3×2 權重矩陣的相應實現將需要 300 μm x 200 μm 的計算區域大小，大約是 21.7 倍 。
請注意，最先進的 GPU Tesla V100 的能效和計算密度分別為每瓦 100 GOP/s 和每平方毫米 37 GOP/s（Giga-Operations/s）。DGNN 架構在能效上實現了四個數量級以上的提升，在計算速度上實現了三個數量級以上的提升。
圖示：合成圖上的半監督節點分類。（來源：論文）
架構的可擴展性
所提出的 DGNN 架構只執行一次 AGG(·) 以直接考慮高階節點特徵，避免了在提取遠程鄰域信息時的指數鄰域擴展問題，並有利於學習更大圖的可擴展性。原則上，架構的頭數可以縮放到任意大小，並且基本DPU模塊（例如圖1C中的）可以水平堆疊並與Y耦合器和條形波導互連，以聚合來自任意大小鄰域的光學神經消息。
此外，該架構具有通過進一步堆疊 DPU 模塊擴展多輪光學消息傳遞的靈活性。DPU 模塊可以通過增加每層金屬層和元原子的數量來放大，DPU 的輸入和輸出數量可以通過額外的光調製器和波導交叉來放大。
圖示：放大 DGNN 節點的神經消息維度。（來源：論文）
該架構的工作波長可以從單波長擴展到多波長，以進一步提高計算吞吐量。通過重新訓練輸出分類器可以減輕系統誤差的累積。
圖示：使用二進制調製訓練 DGNN-E。（來源：論文）
此外，原位訓練方法還可以通過開發具有可編程調製係數的片上 DPU 模塊（例如，使用一維氧化銦錫進行調製）來解決系統錯誤並提高訓練效率。
限制和未來的工作
在本研究中，DGNN 中的光學特徵聚合是使用 2×1 光學 Y 耦合器實現的，其組合比為 50:50，不支持為不同的相鄰節點分配不同的權重，即加權和。儘管平均特徵聚合在節點和圖級分類任務中已經取得了顯著的性能，但使用加權和的消息傳遞可以進一步提高模型容量，並且可以使用片上調幅器（例如相變材料）來實現。
另一個限制是所提出的 DGNN 架構使用線性模型進行光學消息傳遞。儘管現有的工作已經證明了實現光學非線性激活函數的可能性，但非線性操作在 GNN 中並不重要，正如之前的工作所研究的那樣。這可以通過 DGNN 在現實世界的基準數據集中實現的卓越模型性能來證明。
例如，在大型稀缺訓練標籤下，DGNN 在 Amazon Photo 上幾乎達到了最先進的性能，並在稀缺標籤設置下顯著優於電子 GNN。因此，在 DGNN 中包含非線性激活函數留待未來工作，作為進一步增強模型學習能力的潛力。
綜上所述，研究人員希望該工作將激發高級光學深度學習架構的未來發展，該架構具有超越歐幾里得域的集成光子電路，用於高效的圖形表示學習。
論文鏈接：https://www.science.org/doi/10.1126/sciadv.abn7630
