芯東西(公眾號:aichip001)
編輯 | ?Panken

芯東西8月10日報道,清華大學電子工程系方璐教授課題組、自動化系戴瓊海院士課題組首創全前向智能光計算訓練架構,研制了“太極-II”光訓練芯片,擺脫了對離線訓練的依賴,實現了光計算系統大規模神經網絡的高效精準訓練。相關研究成果本周發表于國際頂級學術期刊Nature。

清華大學電子工程系發文介紹說,太極-II的面世,填補了智能光計算在大規模訓練這一核心拼圖的空白

光計算具有高算力低功耗的特性,是加速智能計算的一大前沿方向。Nature審稿人在審稿評述中提到“本文中提出的想法非常新穎,此類光學神經網絡(ONN)的訓練過程是前所未有的。所提出的方法不僅有效,而且容易實現。因此,它有望成為訓練光學神經網絡和其他光學計算系統的廣泛采用的工具”。

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

論文鏈接://www.nature.com/articles/s41586-024-07687-4

一、巧用對稱,助力光計算擺脫GPU依賴

光計算有望提高機器學習應用的速度和能源效率。然而,目前有效訓練這些模型的方法受到計算機仿真的限制。

發表于國際學術頂刊Science的通用智能光計算芯片“太極”首次將光計算從原理驗證推向了大規模實驗應用,以160TOPS/W的系統級能效為復雜智能任務的推理帶來曙光,但卻未能釋放光計算的“訓練之能”。

相較于推理而言,模型訓練更訓練大量算力。電訓練架構要求前向-反向傳播模型高度匹配,對光計算物理系統的精準對齊提出了苛刻的要求,致使梯度計算難、離線建模慢、映射誤差大,禁錮了光訓練的規模與效率。

而方璐、戴瓊海課題組找到了“光子傳播對稱性”這把鑰匙,用全前向光訓練突破了電訓練架構對物理光計算的掣肘。

▲全前向智能光計算訓練架構(圖源:清華大學)

據論文第一作者薛智威介紹,在太極-II架構下,梯度下降中的反向傳播化為了光學系統的前向傳播,光學神經網絡的訓練利用數據-誤差兩次前向傳播即可實現。兩次前向傳播具備天然的對齊特性,保障了物理梯度的精確計算。如此實現的訓練精度高,便能夠支撐大規模的網絡訓練。

物理光系統的調制-傳播與神經網絡的激活-連接相互映射,即調制模塊的訓練可驅動任意網絡的權重優化,從而保障了訓練的速度與能效。

由于不需要進行反向傳播,太極-II架構不再依賴電計算進行離線的建模與訓練,大規模神經網絡的精準高效光訓練得以實現。

二、訓練有數百萬參數的光網絡,速度提升1個數量級

以光為計算媒介,以光的可控傳播構建計算模型,光計算天然具有高速和低功耗的特性,利用光的全前向傳播實現訓練能夠極大提升光網絡訓練的速度與能效。

系統實測結果表明,太極-II能夠對多種不同光學系統進行訓練,并在各種任務下均表現出了卓越的性能。

1、大規模學習領域:突破了計算精度與效率的矛盾,將數百萬參數的光網絡訓練速度提升了1個數量級,代表性智能分類任務的準確率提升40%

2、復雜場景智能成像:弱光環境(每像素光強度僅為亞光子)下實現了能量效率為5.40×10^6 TOPS/W的全光處理,系統級能效提升6個數量級。在非視域等復雜場景成像應用中,實現了千赫茲幀率的智能成像,效率提升2個數量級

3、拓撲光子學領域:在不依賴任何模型先驗下可自動搜索非厄米奇異點,為高效精準解析復雜拓撲系統提供了新思路。

▲通用智能光訓練賦能復雜系統(圖源:清華大學)

三、推動應用和理論進步,為AI大模型提供算力新動力

太極-II在拓撲光子學領域亦展露出應用潛力,在不依賴任何模型先驗下可自動搜索非厄米奇異點,為高效精準解析復雜拓撲系統提供了新思路。

如兩儀分立,太極I和II分別實現了高能效AI推理與訓練;又如兩儀調和,太極I和II共同構成了大規模智能計算的完整生命周期。

“「定兩儀太極之道,合正反乾坤之法」,我們這樣形容太極系列這一組辯證協作架構,我們相信,它們將合力為未來AI大模型注入算力發展的新動力,構建光算力的新基座。”方璐說。

在原理樣片的基礎上,研究團隊正積極地向智能光芯片產業化邁進,在多種端側智能系統上進行了應用部署。

兩代太極芯片相繼揭示了智能光計算的巨大潛力。經過包括太極系列在內的光計算領域的不懈努力,智能光計算平臺將有望以更低的資源消耗和更小的邊際成本,為AI大模型、通用人工智能、復雜智能系統的高速高能效計算開辟新路徑。

來源:清華大學,Nature