清華大學(xué)突破智能光計(jì)算訓(xùn)練難題

清華大學(xué)電子工程系方璐教授課題組自動(dòng)化系戴瓊海院士課題組另辟蹊徑首創(chuàng)了全前向智能光計(jì)算訓(xùn)練架構(gòu)，研制了“太極-II”光訓(xùn)練芯片，實(shí)現(xiàn)了光計(jì)算系統(tǒng)大規(guī)模神經(jīng)網(wǎng)絡(luò)的高效精準(zhǔn)訓(xùn)練，該研究成果以“光神經(jīng)網(wǎng)絡(luò)全前向訓(xùn)練”為題，發(fā)表于《自然》期刊上。

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)的廣泛采用的工具�！�

巧用對稱，助力光計(jì)算擺脫GPU依賴

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

相較于模型推理而言，模型訓(xùn)練更需要大規(guī)模算力。然而，現(xiàn)有的光神經(jīng)網(wǎng)絡(luò)訓(xùn)練嚴(yán)重依賴GPU進(jìn)行離線建模并且要求物理系統(tǒng)精準(zhǔn)對齊。正因如此，光學(xué)訓(xùn)練的規(guī)模受到了極大的限制，光高性能計(jì)算的優(yōu)勢仿佛被禁錮在無形的枷鎖之中。

在這個(gè)時(shí)候，方璐、戴瓊海課題組找到了“光子傳播對稱性”這把鑰匙，將神經(jīng)網(wǎng)絡(luò)訓(xùn)練中的前向與反向傳播都等效為光的前向傳播。

據(jù)論文第一作者、電子系博士生薛智威介紹，在太極-II架構(gòu)下，梯度下降中的反向傳播化為了光學(xué)系統(tǒng)的前向傳播，光學(xué)神經(jīng)網(wǎng)絡(luò)的訓(xùn)練利用數(shù)據(jù)-誤差兩次前向傳播即可實(shí)現(xiàn)。兩次前向傳播具備天然的對齊特性，保障了物理梯度的精確計(jì)算。如此實(shí)現(xiàn)的訓(xùn)練精度高，便能夠支撐大規(guī)模的網(wǎng)絡(luò)訓(xùn)練。

由于不需要進(jìn)行反向傳播，太極-II架構(gòu)不再依賴電計(jì)算進(jìn)行離線的建模與訓(xùn)練，大規(guī)模神經(jīng)網(wǎng)絡(luò)的精準(zhǔn)高效光訓(xùn)練終于得以實(shí)現(xiàn)。

圖片:557A6DEBDA4B29834A1F5E457BA_DBE99787_3AC3D.png

全前向智能光計(jì)算訓(xùn)練架構(gòu)

高效精準(zhǔn)，智能光訓(xùn)練事事可為

以光為計(jì)算媒介，以光的可控傳播構(gòu)建計(jì)算模型，光計(jì)算天然具有高速和低功耗的特性，利用光的全前向傳播實(shí)現(xiàn)訓(xùn)練能夠極大的提升光網(wǎng)絡(luò)訓(xùn)練的速度與能效。

論文研究表明，太極-II能夠?qū)Χ喾N不同光學(xué)系統(tǒng)進(jìn)行訓(xùn)練，并在各種任務(wù)下均表現(xiàn)出了卓越的性能。

大規(guī)模學(xué)習(xí)領(lǐng)域：突破了計(jì)算精度與效率的矛盾，將數(shù)百萬參數(shù)的光網(wǎng)絡(luò)訓(xùn)練速度提升了1個(gè)數(shù)量級，代表性智能分類任務(wù)的準(zhǔn)確率提升40%。

復(fù)雜場景智能成像：弱光環(huán)境下（每像素光強(qiáng)度僅為亞光子）實(shí)現(xiàn)了能量效率為5.40×10^6TOPS/W的全光處理，系統(tǒng)級能效提升6個(gè)數(shù)量級。在非視域場景下實(shí)現(xiàn)了千赫茲幀率的智能成像，效率提升2個(gè)數(shù)量級。

拓?fù)涔庾訉W(xué)領(lǐng)域：在不依賴任何模型先驗(yàn)下可自動(dòng)搜索非厄米奇異點(diǎn)，為高效精準(zhǔn)解析復(fù)雜拓?fù)湎到y(tǒng)提供了新思路。

圖片:59C03606438E95A03FD3A67B67D_9985FD9B_7A0E0.png

通用智能光訓(xùn)練賦能復(fù)雜系統(tǒng)

攜手太極，推動(dòng)AI光算力揚(yáng)帆遠(yuǎn)航

太極-II的面世，繼太極I芯片之后進(jìn)一步揭示了智能光計(jì)算的巨大潛力。

如兩儀分立，太極I和II分別實(shí)現(xiàn)了大規(guī)模智能光計(jì)算的高效推理與訓(xùn)練；

又如兩儀調(diào)和，太極I和II共同構(gòu)成了大規(guī)模智能計(jì)算的完整生命周期。

方璐表示：“‘定兩儀太極之道，合正反乾坤之法’，我們這樣形容太極系列這一組辯證協(xié)作架構(gòu)，我們相信，它們將合力為未來AI大模型注入算力發(fā)展的新動(dòng)力，構(gòu)建光算力的新基座”。

在原理樣片的基礎(chǔ)上，研究團(tuán)隊(duì)正積極地向智能光芯片產(chǎn)業(yè)化邁進(jìn)，在多種端側(cè)智能系統(tǒng)上進(jìn)行了應(yīng)用部署。

可以預(yù)見，經(jīng)過太極系列在內(nèi)的光計(jì)算領(lǐng)域的不懈努力，智能光計(jì)算平臺(tái)將有望以更低的資源消耗和更小的邊際成本，為人工智能大模型、通用人工智能、復(fù)雜智能系統(tǒng)的高速高能效計(jì)算開辟新路徑。

清華大學(xué)電子系為論文第一單位，方璐教授、戴瓊海教授為論文的通訊作者，清華大學(xué)電子系博士生薛智威、博士后周天貺為共同一作，電子系博士生徐智昊、之江實(shí)驗(yàn)室虞紹良博士參與了本項(xiàng)工作。本課題受到國家科技部、國家自然科學(xué)基金委、北京信息科學(xué)與技術(shù)國家研究中心、清華大學(xué)-之江實(shí)驗(yàn)室聯(lián)合研究中心的支持。

相關(guān)鏈接：https://www.nature.com/articles/s41586-024-07687-4