在剛剛過去的一天，來自清華的光電智能技術(shù)交叉創(chuàng)新團(tuán)隊(duì)突破智能光計(jì)算訓(xùn)練難題，相關(guān)論文登上 Nature。

論文共同一作是來自清華的薛智威、周天貺，通訊作者是清華的方璐教授、戴瓊海院士。此外，清華電子系徐智昊、之江實(shí)驗(yàn)室虞紹良也參與了這項(xiàng)研究。

• 論文地址：https://www.nature.com/articles/s41586-024-07687-4
• 論文標(biāo)題：Fully forward mode training for optical neural networks

隨著大模型的規(guī)模越來越大，算力需求爆發(fā)式增長，就拿 Sora 來說，據(jù)爆料，訓(xùn)練參數(shù)量約為 30 億，預(yù)計(jì)使用了 4200-10500 塊 H100 訓(xùn)了 1 個(gè)月。全球的科技大廠都在高價(jià)求購的「卡」，都是硅基的電子芯片。在此之外，還有一種將計(jì)算載體從電變?yōu)楣獾墓庾有酒夹g(shù)。它們利用光在芯片中的傳播進(jìn)行計(jì)算，具有超高的并行度和速度，被認(rèn)為是未來顛覆性計(jì)算架構(gòu)最有力的競(jìng)爭方案之一。

光計(jì)算領(lǐng)域也在使用 AI 輔助設(shè)計(jì)系統(tǒng)。然而，AI 也給光計(jì)算技術(shù)套上了「瓶頸」—— 光神經(jīng)網(wǎng)絡(luò)訓(xùn)練嚴(yán)重依賴基于數(shù)據(jù)對(duì)光學(xué)系統(tǒng)建模的方法。這導(dǎo)致研究人員難以修正實(shí)驗(yàn)誤差。更重要的是，不完善的系統(tǒng)加上光傳播的復(fù)雜性，幾乎不可能實(shí)現(xiàn)對(duì)光學(xué)系統(tǒng)的完美建模，離線模型與現(xiàn)實(shí)之間總是難以完全同步。

而機(jī)器學(xué)習(xí)常用的「梯度下降」和「反向傳播」，來到了光學(xué)領(lǐng)域，也不好使了。為了使基于梯度的方法有效，光學(xué)系統(tǒng)必須非常精確地校準(zhǔn)和對(duì)齊，以確保光信號(hào)能夠正確地在系統(tǒng)中反向傳播，離線模型往往很難實(shí)現(xiàn)這點(diǎn)。

來自清華大學(xué)的研究團(tuán)隊(duì)抓住了光子傳播具有對(duì)稱性這一特性，將神經(jīng)網(wǎng)絡(luò)訓(xùn)練中的前向與反向傳播都等效為光的前向傳播。該研究開發(fā)了一種稱為全前向模式（FFM，fully forward mode）學(xué)習(xí)的方法，研究人員不再需要在計(jì)算機(jī)模型中建模，可以直接在物理光學(xué)系統(tǒng)上設(shè)計(jì)和調(diào)整光學(xué)參數(shù)，再根據(jù)測(cè)量的光場(chǎng)數(shù)據(jù)和誤差，使用梯度下降算法有效地得出最終的模型參數(shù)。借助 FFM，大多數(shù)機(jī)器學(xué)習(xí)操作都可以有效地并行進(jìn)行，從而減輕了 AI 對(duì)光學(xué)系統(tǒng)建模的限制。

FFM 學(xué)習(xí)表明，訓(xùn)練具有數(shù)百萬個(gè)參數(shù)的光神經(jīng)網(wǎng)絡(luò)可以達(dá)到與理想模型相當(dāng)?shù)臏?zhǔn)確率。

此外，該方法還支持通過散射介質(zhì)進(jìn)行全光學(xué)聚焦，分辨率達(dá)到衍射極限；它還可以以超過千赫茲的幀率平行成像隱藏在視線外的物體，并可以在室溫下進(jìn)行光強(qiáng)弱至每像素亞光子的全光處理。 

最后，研究證明了 FFM 學(xué)習(xí)可以在沒有分析模型的情況下自動(dòng)搜索非厄米異常點(diǎn)。FFM 學(xué)習(xí)不僅有助于將學(xué)習(xí)過程提高幾個(gè)數(shù)量級(jí)，還可以推動(dòng)深度神經(jīng)網(wǎng)絡(luò)、超靈敏感知和拓?fù)涔鈱W(xué)等應(yīng)用和理論領(lǐng)域的發(fā)展。

深度 ONN 上的并行 FFM 梯度下降

圖 2a 展示了使用 FFM 學(xué)習(xí)的自由空間 ONN（optical neural networks，光學(xué)神經(jīng)網(wǎng)絡(luò)）的自我訓(xùn)練過程。為了驗(yàn)證 FFM 學(xué)習(xí)的有效性，研究者首先使用基準(zhǔn)數(shù)據(jù)集訓(xùn)練了一個(gè)單層 ONN 以進(jìn)行對(duì)象分類。

圖 2b 可視化了在 MNIST 數(shù)據(jù)集上的訓(xùn)練結(jié)果，可以看到，實(shí)驗(yàn)和理論光場(chǎng)之間的結(jié)構(gòu)相似性指數(shù)（SSIM）超過了 0.97，這意味著相似度很高（圖 2c）。值得注意的是，由于系統(tǒng)不完善的原因，光場(chǎng)和梯度的理論結(jié)果并不能精準(zhǔn)地代表物理結(jié)果。因此，這些理論結(jié)果不應(yīng)被視為基本事實(shí)。

接下來，研究者探究了用于 Fashion-MNIST 數(shù)據(jù)集分類的多層 ONN，具體如圖 2d 所示。

通過將層數(shù)從 2 層增加到 8 層，他們觀察到，計(jì)算機(jī)訓(xùn)練網(wǎng)絡(luò)的實(shí)驗(yàn)測(cè)試結(jié)果平均達(dá)到了 44.0% (35.1%)、52.4%（8.8%）、58.4%（18.4%）和 58.8%（5.5%）的準(zhǔn)確率（兩倍標(biāo)準(zhǔn)差）。這些結(jié)果低于 92.2%、93.8%、96.0% 和 96.0% 的理論準(zhǔn)確率。通過 FFM 學(xué)習(xí)，準(zhǔn)確率數(shù)值分別提升到了 86.5%、91.0%、92.3% 和 92.5%，接近理想的計(jì)算機(jī)準(zhǔn)確率。

圖 2e 描述了 8 層 ONN 的輸出結(jié)果。隨著層數(shù)增加，計(jì)算機(jī)訓(xùn)練的實(shí)驗(yàn)輸出逐漸偏離目標(biāo)輸出并最終對(duì)對(duì)象做出誤分類。相比之外，F(xiàn)FM 設(shè)計(jì)的網(wǎng)絡(luò)可以準(zhǔn)確地進(jìn)行正確分類。除了計(jì)算密集型數(shù)據(jù)和誤傳播之外，損失和梯度計(jì)算還可以通過現(xiàn)場(chǎng)光學(xué)和電子處理來執(zhí)行。

研究者進(jìn)一步提出了非線性 FFM 學(xué)習(xí)，如圖 2f 所示。在數(shù)據(jù)傳播中，輸出在饋入到下一層之前被非線性地激活，記錄非線性激活的輸入并計(jì)算相關(guān)梯度。在誤差傳播過程中，輸入在傳播之前與梯度相乘。

利用 FFM 進(jìn)行全光學(xué)成像和處理

圖 3a 展示了點(diǎn)掃描散射成像系統(tǒng)的實(shí)現(xiàn)原理。一般來說，在自適應(yīng)光學(xué)中，啟發(fā)式優(yōu)化方法已經(jīng)用于焦點(diǎn)優(yōu)化。

研究者分析了不同的 SOTA 優(yōu)化方法，并利用粒子群優(yōu)化（PSO）進(jìn)行比較，如圖 3b 所示。出于評(píng)估的目的，這里采用了兩種不同類型的散射介質(zhì)，分別是隨機(jī)相位板（稱為 Scatterer-I）和透明膠帶（稱為 Scatterer-II）。基于梯度的 FFM 學(xué)習(xí)表現(xiàn)出更高的效率，在兩種散射介質(zhì)的實(shí)驗(yàn)中經(jīng)過 25 次迭代后收斂，收斂損耗值分別為 1.84 和 2.07。相比之下，PSO 方法需要至少 400 次迭代后才能進(jìn)行收斂，最終損耗值為 2.01 和 2.15。

圖 3c 描述了 FFM 自我設(shè)計(jì)的演變過程，展示了最開始隨機(jī)分布的強(qiáng)度逐漸分布圖逐漸收斂到一個(gè)緊密的點(diǎn)，隨后在整個(gè) 3.2 毫米 × 3.2 毫米成像區(qū)域來學(xué)習(xí)設(shè)計(jì)的焦點(diǎn)。

圖 3d 比較了使用 FFM 和 PSO 分別優(yōu)化的焦點(diǎn)的半峰全寬（FWHM）和峰值信噪比（PSNR）指標(biāo)。使用 FFM，平均 FWHM 為 81.2 μm，平均 PSNR 為 8.46 dB，最低 FWHM 為 65.6 μm。當(dāng)使用 3.2mm 寬的方形孔徑和 0.388m 的傳播距離時(shí)，通過 FFM 學(xué)習(xí)設(shè)計(jì)的焦點(diǎn)尺寸接近衍射極限 64.5 μm。相比之下，PSO 優(yōu)化產(chǎn)生的 FWHM 為 120.0 μm，PSNR 為 2.29 dB。

在圖 4a 中，利用往返隱藏對(duì)象的光路之間的空間對(duì)稱性，F(xiàn)FM 學(xué)習(xí)可以實(shí)現(xiàn)動(dòng)態(tài)隱層對(duì)象的全光學(xué)現(xiàn)場(chǎng)重建和分析。圖 4b 展示了 NLOS 成像，在學(xué)習(xí)過程中，輸入波峰被設(shè)計(jì)用來將對(duì)象中所有網(wǎng)格同步映射到它們的目標(biāo)位置。

現(xiàn)場(chǎng)光子集成電路與 FFM

FFM 學(xué)習(xí)方法可以推廣到集成光系統(tǒng)的自設(shè)計(jì)中。圖 5a 展示了 FFM 學(xué)習(xí)實(shí)現(xiàn)過程。其中矩陣的對(duì)稱性允許誤差傳播矩陣和數(shù)據(jù)傳播矩陣之間對(duì)等。因此，數(shù)據(jù)和誤差傳播共享相同的傳播方向。圖 5b 展示了對(duì)稱核心實(shí)現(xiàn)和封裝芯片實(shí)驗(yàn)的測(cè)試設(shè)置。

研究者構(gòu)建的神經(jīng)網(wǎng)絡(luò)用于對(duì)鳶尾花（Iris）數(shù)據(jù)進(jìn)行分類，輸入處理為 16 × 1 向量，輸出代表三種花的類別之一。訓(xùn)練期間矩陣編程的保真度如圖 5c 中所示，三個(gè)對(duì)稱矩陣值的時(shí)間漂移分別產(chǎn)生了 0.012%、0.012% 和 0.010% 的標(biāo)準(zhǔn)偏差。

在這種不確定下，研究者將實(shí)驗(yàn)梯度與模擬值進(jìn)行比較。如圖 5d 所示，實(shí)驗(yàn)梯度與理想模擬值的平均偏差為 3.5%。圖 5d 還說明了第 80 次學(xué)習(xí)迭代時(shí)第二層的設(shè)計(jì)梯度，而整個(gè)神經(jīng)網(wǎng)絡(luò)的誤差在圖 5e 中進(jìn)行了可視化。在第 80 次迭代中，F(xiàn)FM 學(xué)習(xí)（計(jì)算機(jī)模擬訓(xùn)練）的梯度誤差為 3.50%（5.10%）、3.58%（5.19%）、3.51%（5.24%）、3.56%（5.29%）和 3.46%（5.94%）。設(shè)計(jì)精度的演變?nèi)鐖D 5f 所示。理想模擬和 FFM 實(shí)驗(yàn)都需要大約 100 個(gè) epoch 才能收斂。在三種對(duì)稱率配置下，實(shí)驗(yàn)性能與模擬性能相似，網(wǎng)絡(luò)收斂到 94.7%、89.2% 和 89.0% 的準(zhǔn)確率。FFM 方法實(shí)現(xiàn)了 94.2%、89.2% 和 88.7% 的準(zhǔn)確率。相比之下，計(jì)算機(jī)設(shè)計(jì)的網(wǎng)絡(luò)表現(xiàn)出 71.7%、65.8% 和 55.0% 的實(shí)驗(yàn)準(zhǔn)確率。

基于這篇論文的成果，研究團(tuán)隊(duì)也推出了「太極 - II」光訓(xùn)練芯片。「太極 - II」的研發(fā)距離上一代「太極」僅過了 4 個(gè)月，相關(guān)成果也登上了 Science。

論文鏈接：https://www.science.org/doi/10.1126/science.adl1203

值得一提的是，作為全球首款大規(guī)模干涉衍射異構(gòu)集成芯片的「太極」，其計(jì)算能力可以比肩億級(jí)神經(jīng)元的芯片。論文的實(shí)驗(yàn)結(jié)果顯示，「太極」的能效是英偉達(dá) H100 的 1000 倍。這種強(qiáng)大的計(jì)算能力基于研究團(tuán)隊(duì)首創(chuàng)的分布式廣度智能光計(jì)算架構(gòu)。

更多細(xì)節(jié)，請(qǐng)參考原論文。