編輯 | 蘿蔔皮
光子神經網路使用光子而不是電子執行受大腦啟發的計算，以實現顯著提高的計算效能。然而，現有架構只能處理具有規則結構的資料，而無法泛化到歐幾里得空間之外的圖結構資料。
清華大學的研究人員提出了衍射圖神經網路 (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