
芯東西(公眾號:aichip001)
編輯 | ?Panken
芯東西8月10日報道,清華大學電子工程系方璐教授課題組、自動化系戴瓊海院士課題組首創全前向智能光計算訓練架構,研制了“太極-II”光訓練芯片,擺(bai)脫了(le)對離線訓練的依賴,實現了(le)光計算系統大規(gui)模神(shen)經網絡的高效精準訓練。相關研究成(cheng)果本周發表于國際頂級(ji)學術期刊Nature。
清華大學電子工程系發文介紹說,太極-II的面世,填補了智能光計算在大規模訓練這一核心拼圖的空白。
光計算具有高算力低功耗的特性,是加速智能計算的一大前沿方向。Nature審稿人在審稿評述中提到“本文中提出的想法非常新穎,此類光學神經網絡(ONN)的訓練過程是前所未有的。所提出的方法不僅有效,而且容易實現。因此,它有望成為訓練光學神經網絡和其他光學計算系統的廣泛采用的工具”。
清華(hua)(hua)大(da)學(xue)(xue)電子(zi)系(xi)為論文第一單(dan)位,方璐(lu)教(jiao)授、戴(dai)瓊(qiong)海教(jiao)授為論文的(de)通訊作者,清華(hua)(hua)大(da)學(xue)(xue)電子(zi)系(xi)博(bo)士(shi)生薛智威、博(bo)士(shi)后周天貺為共同一作,電子(zi)系(xi)博(bo)士(shi)生徐智昊、之(zhi)江實驗室虞紹良(liang)博(bo)士(shi)參與了本項(xiang)工(gong)作。本課(ke)題受到國(guo)家(jia)科(ke)技部、國(guo)家(jia)自然(ran)科(ke)學(xue)(xue)基(ji)金委(wei)、北京信息(xi)科(ke)學(xue)(xue)與技術(shu)國(guo)家(jia)研究(jiu)中(zhong)心、清華(hua)(hua)大(da)學(xue)(xue)-之(zhi)江實驗室聯合研究(jiu)中(zhong)心的(de)支持。
論文鏈接://www.nature.com/articles/s41586-024-07687-4
一、巧用對稱,助力光計算擺脫GPU依賴
光計(ji)算有(you)望提高機器學習應用的速度和能源效率。然(ran)而,目(mu)前有(you)效訓練這(zhe)些模型(xing)的方(fang)法(fa)受到計(ji)算機仿真(zhen)的限制(zhi)。
發表于國際學術頂刊Science的通用智能光計算芯片“太極”,首次將光計算從原理驗證推向了大規模實驗應用,以160TOPS/W的系統(tong)級(ji)能(neng)效為(wei)復雜智能(neng)任(ren)務的推理帶(dai)來曙光,但卻未能(neng)釋(shi)放光計算的“訓練之能(neng)”。
相較(jiao)于推理而言,模(mo)型(xing)(xing)訓(xun)練(lian)(lian)更訓(xun)練(lian)(lian)大量(liang)算力。電(dian)訓(xun)練(lian)(lian)架構要(yao)求(qiu)前向-反(fan)向傳播(bo)模(mo)型(xing)(xing)高度匹配,對光計算物理系統的精準(zhun)對齊提出了苛刻的要(yao)求(qiu),致(zhi)使梯(ti)度計算難、離線建模(mo)慢(man)、映射誤差大,禁錮了光訓(xun)練(lian)(lian)的規(gui)模(mo)與效率(lv)。
而方璐、戴瓊海課題組找到了“光子傳播對稱性”這把鑰匙,用全前向光訓練突破了電訓練架構對物理光計算的掣(che)肘。
▲全前(qian)向(xiang)智能(neng)光計算訓(xun)練架構(gou)(圖(tu)源:清華大學)
據論(lun)文第一作(zuo)者薛(xue)智(zhi)威介(jie)紹,在(zai)太極-II架構下,梯度(du)下降中的(de)反向(xiang)(xiang)傳(chuan)(chuan)播(bo)化為了光學系統(tong)的(de)前向(xiang)(xiang)傳(chuan)(chuan)播(bo),光學神經網絡(luo)的(de)訓(xun)(xun)練利用數據-誤差兩次(ci)前向(xiang)(xiang)傳(chuan)(chuan)播(bo)即可實現。兩次(ci)前向(xiang)(xiang)傳(chuan)(chuan)播(bo)具備天(tian)然的(de)對齊特性,保障了物(wu)理梯度(du)的(de)精確計算。如此(ci)實現的(de)訓(xun)(xun)練精度(du)高,便能(neng)夠(gou)支撐大規(gui)模的(de)網絡(luo)訓(xun)(xun)練。
物理光系(xi)統的(de)調(diao)制-傳播與神經網絡的(de)激活(huo)-連接相互映(ying)射,即調(diao)制模(mo)塊的(de)訓(xun)練可驅(qu)動任意(yi)網絡的(de)權重優化,從而保障了(le)訓(xun)練的(de)速度與能(neng)效。
由于不需(xu)要進行(xing)反向傳播,太極-II架構不再依賴電計算進行(xing)離線(xian)的(de)建模(mo)與(yu)訓練,大(da)規模(mo)神經網絡的(de)精準(zhun)高效光訓練得(de)以實(shi)現。
二、訓練有數百萬參數的光網絡,速度提升1個數量級
以光為計算(suan)媒介(jie),以光的(de)可控傳(chuan)播構建計算(suan)模型,光計算(suan)天(tian)然具有高速和低功耗(hao)的(de)特性,利用(yong)光的(de)全前向傳(chuan)播實(shi)現訓(xun)練(lian)能夠極大提升(sheng)光網絡訓(xun)練(lian)的(de)速度與(yu)能效。
系統實測結果表(biao)明,太極-II能夠對多種不同光學系統進行訓練,并(bing)在(zai)各種任務下均表(biao)現出了卓越(yue)的(de)性(xing)能。
1、大規模學習領域:突破了計算精度與效率的矛盾,將數百萬參數的光網絡訓練速度提升了1個數量級,代表性智能分類任務的準確率提升40%。
2、復雜場景智能成像:弱光環境(每像素光強度僅為亞光子)下實現了能量效率為5.40×10^6 TOPS/W的全光處理,系統級能效提升6個數量級。在非視域等復雜場景成像應用中,實現了千赫茲幀率的智能成像,效率提升2個數量級。
3、拓撲光子學領域:在(zai)不依賴任何模型先驗下(xia)可自動搜索(suo)非厄米奇異點,為高(gao)效精準解析復(fu)雜(za)拓撲(pu)系統提供(gong)了新思路(lu)。
▲通用智能光訓練賦能復雜系(xi)統(圖源:清華(hua)大學)
三、推動應用和理論進步,為AI大模型提供算力新動力
太(tai)極-II在拓(tuo)撲(pu)光子學領域亦展露出應用潛力,在不依賴(lai)任(ren)何模(mo)型(xing)先驗下可自動搜(sou)索非(fei)厄(e)米奇異點,為高(gao)效精準解析復(fu)雜拓(tuo)撲(pu)系(xi)統提供(gong)了新(xin)思路。
如(ru)(ru)兩(liang)儀(yi)分立,太(tai)極(ji)I和II分別實現了高能效AI推理與訓(xun)練;又如(ru)(ru)兩(liang)儀(yi)調和,太(tai)極(ji)I和II共同構成了大規模智能計算(suan)的完(wan)整生命周期。
“「定兩儀(yi)太極之(zhi)道(dao),合正反(fan)乾(qian)坤之(zhi)法」,我們(men)這樣形容太極系列這一組辯證協作架(jia)構,我們(men)相(xiang)信,它們(men)將合力(li)(li)為(wei)未來AI大模型(xing)注入算(suan)力(li)(li)發(fa)展的(de)新(xin)動(dong)力(li)(li),構建光算(suan)力(li)(li)的(de)新(xin)基座。”方璐說。
在原(yuan)理樣片的(de)基礎上,研究團隊正積(ji)極地向智能(neng)光芯(xin)片產業化邁進,在多種端側智能(neng)系統上進行了應用(yong)部(bu)署。
兩代太極(ji)芯片相繼揭示了智(zhi)(zhi)能(neng)光(guang)計(ji)算的(de)巨大(da)潛力(li)。經過包括太極(ji)系(xi)列在內的(de)光(guang)計(ji)算領域的(de)不懈(xie)努(nu)力(li),智(zhi)(zhi)能(neng)光(guang)計(ji)算平臺將有望以(yi)更低(di)的(de)資源(yuan)消耗和(he)更小的(de)邊際成本,為AI大(da)模型、通用人工智(zhi)(zhi)能(neng)、復雜智(zhi)(zhi)能(neng)系(xi)統的(de)高(gao)速(su)高(gao)能(neng)效(xiao)計(ji)算開辟新路(lu)徑(jing)。
來源:清華大學,Nature