中國科學(xué)家新突破：未來的AI模型，或許可以用“光”來訓(xùn)練了！

撰文 | 馬雪薇

前言

未來的人工智能（AI）模型，或許可以用“光”來訓(xùn)練了。

長期以來，傳統(tǒng)的電子計(jì)算推動(dòng)著 AI 模型規(guī)模的不斷擴(kuò)展，但大規(guī)模意味著高算力，也意味著高能耗，因此，亟需一種高效低能耗的計(jì)算方法來解決未來 AI 模型所面臨的能源問題。

近年間，具有高算力低功耗特性的智能光計(jì)算逐步登上了算力發(fā)展的舞臺。通用智能光計(jì)算芯片“太極”的問世便是其中的一個(gè)縮影，它首次將光計(jì)算從原理驗(yàn)證推向了大規(guī)模實(shí)驗(yàn)應(yīng)用，以 160TOPS/W 的系統(tǒng)級能效為大規(guī)模復(fù)雜任務(wù)的“推理”帶來了曙光，但未能夠釋放智能光計(jì)算的“訓(xùn)練之能”。

訓(xùn)練和推理是 AI 大模型核心能力的兩大基石，缺一不可。相較于模型推理而言，模型訓(xùn)練更需要大規(guī)模算力。

然而，現(xiàn)有的光神經(jīng)網(wǎng)絡(luò)訓(xùn)練嚴(yán)重依賴電計(jì)算進(jìn)行離線建模并且要求物理系統(tǒng)精準(zhǔn)對齊。正因如此，光學(xué)訓(xùn)練的規(guī)模受到了極大的限制。

由清華大學(xué)電子工程系教授方璐、自動(dòng)化系戴瓊海院士及其科研團(tuán)隊(duì)提出的全前向智能光計(jì)算訓(xùn)練架構(gòu)太極-II，擺脫了原有光計(jì)算系統(tǒng)對電計(jì)算離線建模的依賴，實(shí)現(xiàn)了大規(guī)模神經(jīng)網(wǎng)絡(luò)的在線訓(xùn)練，并有潛力支撐未來 AI 大模型的高速、低功耗訓(xùn)練。

相關(guān)研究論文以“Fully forward mode training for optical neural networks”為題，已發(fā)表在權(quán)威科學(xué)期刊 Nature 上。清華大學(xué)電子系博士生薛智威、博士后周天貺為共同一作。

據(jù)介紹，研究團(tuán)隊(duì)利用“光子傳播對稱性”，將神經(jīng)網(wǎng)絡(luò)訓(xùn)練中的前向與反向傳播都等效為光的前向傳播。在物理系統(tǒng)上實(shí)現(xiàn)計(jì)算密集的訓(xùn)練過程。

據(jù)清華大學(xué)官方報(bào)道，Nature 審稿人在審稿評述中指出“本文中提出的想法非常新穎，此類光學(xué)神經(jīng)網(wǎng)絡(luò)（ONN）的訓(xùn)練過程是前所未有的。所提出的方法不僅有效，而且容易實(shí)現(xiàn)。因此，它有望成為訓(xùn)練光學(xué)神經(jīng)網(wǎng)絡(luò)和其他光學(xué)計(jì)算系統(tǒng)的廣泛采用的工具?！?/p>

太極-II 架構(gòu)是如何工作的？

太極-II 架構(gòu)是一種在物理系統(tǒng)上直接進(jìn)行光學(xué)神經(jīng)網(wǎng)絡(luò)訓(xùn)練的方法。它利用空間對稱性和洛倫茲互易性，消除了反向傳播的需求，提高了訓(xùn)練效率，且可以應(yīng)用于大規(guī)模學(xué)習(xí)、復(fù)雜場景智能成像、拓?fù)涔庾訉W(xué)等領(lǐng)域。原理如下：

將光學(xué)系統(tǒng)參數(shù)化，映射到神經(jīng)網(wǎng)絡(luò)上。

利用空間對稱性和洛倫茲互易性，實(shí)現(xiàn)數(shù)據(jù)傳播和誤差傳播的精準(zhǔn)對齊。

基于測量數(shù)據(jù)/誤差傳播的輸出光場計(jì)算梯度，并進(jìn)行參數(shù)更新。

圖｜太極-II 智能光計(jì)算訓(xùn)練架構(gòu)。a, 傳統(tǒng)上，光學(xué)AI 系統(tǒng)是通過離線建模和優(yōu)化來設(shè)計(jì)的，這導(dǎo)致其設(shè)計(jì)效率和系統(tǒng)性能受限。b, 一般的光學(xué)系統(tǒng)，包括自由空間系統(tǒng)和集成光子系統(tǒng)，都包含了調(diào)制區(qū)域（深綠色）和傳播區(qū)域（淺綠色），在這些區(qū)域中，折射率分別是可調(diào)和固定的。c, 光學(xué)系統(tǒng)中的這些區(qū)域可以映射到神經(jīng)網(wǎng)絡(luò)表示中的權(quán)重和神經(jīng)元連接，這使得可以在輸入和輸出之間構(gòu)建一個(gè)可微分的神經(jīng)網(wǎng)絡(luò)表征（左上）。利用空間對稱的互易性，數(shù)據(jù)和誤差計(jì)算共享前向物理傳播和測量，并在設(shè)計(jì)區(qū)域內(nèi)計(jì)算在線梯度以更新折射率（右上和左下）。通過在線梯度下降，光學(xué)系統(tǒng)逐漸收斂（右下）。

效果怎么樣？

太極-II 架構(gòu)在多個(gè)領(lǐng)域均展現(xiàn)出巨大的應(yīng)用潛力。

1. 光學(xué)神經(jīng)網(wǎng)絡(luò)（ONN）

深層 ONN 構(gòu)建：太極-II 架構(gòu)能夠?qū)崿F(xiàn)具有數(shù)百萬參數(shù)的深層 ONN，并達(dá)到與理想模型相當(dāng)?shù)臏?zhǔn)確率。這對于光學(xué)計(jì)算至關(guān)重要，因?yàn)樯顚泳W(wǎng)絡(luò)可以處理更復(fù)雜的任務(wù)并實(shí)現(xiàn)更高的性能。

非線性 ONN：太極-II 架構(gòu)支持非線性光學(xué)神經(jīng)網(wǎng)絡(luò)的訓(xùn)練，通過在數(shù)據(jù)傳播中引入非線性激活函數(shù)，并利用記錄的函數(shù)輸入輸出來計(jì)算梯度，實(shí)現(xiàn)高效精準(zhǔn)的訓(xùn)練過程。

圖 | 深度 ONN 的并行梯度下降訓(xùn)練。

2. 復(fù)雜場景智能成像

穿散射成像：太極-II 架構(gòu)能夠透過散射介質(zhì)實(shí)現(xiàn)接近衍射極限的聚焦成像，達(dá)到更高的分辨率。這對于顯微成像和宏觀成像等領(lǐng)域具有重要意義。

非視域場景成像：太極-II 架構(gòu)能夠?qū)崿F(xiàn)毫秒級的并行成像，并實(shí)現(xiàn)對非視域的目標(biāo)進(jìn)行全光處理。這對于非視域場景下的動(dòng)態(tài)目標(biāo)成像和識別具有重要意義。

圖 | 利用 太極-II 架構(gòu)，透過散射介質(zhì)達(dá)到衍射極限分辨率。

3. 集成光子系統(tǒng) (PIC)

自設(shè)計(jì) PIC 網(wǎng)絡(luò)：太極-II 架構(gòu)能夠?qū)崿F(xiàn)集成光子系統(tǒng)的自設(shè)計(jì)，并達(dá)到與理論相當(dāng)?shù)臏?zhǔn)確率。這對于構(gòu)建高性能的光子計(jì)算系統(tǒng)具有重要意義。

非厄米系統(tǒng)解析：太極-II 架構(gòu)能夠自動(dòng)搜索非厄米系統(tǒng)的奇異點(diǎn)，并無需物理模型。這對于研究和解析復(fù)雜拓?fù)湎到y(tǒng)具有重要意義。

圖 | 基于太極-II 架構(gòu)的集成光子系統(tǒng)在線訓(xùn)練。

用“光”拉動(dòng) AI 算力增長

太極-II 架構(gòu)將光學(xué)系統(tǒng)視為可學(xué)習(xí)的神經(jīng)網(wǎng)絡(luò)，并通過目標(biāo)導(dǎo)向的優(yōu)化方法在物理系統(tǒng)中進(jìn)行設(shè)計(jì)和訓(xùn)練，突破了傳統(tǒng)光學(xué)設(shè)計(jì)方法的局限性，實(shí)現(xiàn)了光學(xué)系統(tǒng)的自主學(xué)習(xí)和高性能設(shè)計(jì)。

太極-II 架構(gòu)可以應(yīng)用于各種光學(xué)系統(tǒng)，包括自由空間系統(tǒng)和集成光子系統(tǒng)，并能夠支撐多種智能任務(wù)，例如圖像處理、模式識別、機(jī)器學(xué)習(xí)和深度學(xué)習(xí)等。

可以預(yù)見，智能光計(jì)算平臺將有望以更低的資源消耗和更小的邊際成本，為人工智能大模型、通用人工智能、復(fù)雜智能系統(tǒng)的高速高能效計(jì)算開辟新路徑。