清华团队提出新型光计算架构，光训练速度提升1个数量级

近期，清华大学电子工程系方璐教授课题组、自动化系戴琼海教授课题组创新性地设计了全前向智能光计算训练架构，开发出一种名为“太极-II”的通用光训练芯片。

该架构摆脱了对电计算离线训练的依赖，而且能够为智能系统的高效光训练提供支撑。

图丨图中为课题负责人方璐教授，图左和图右分别为论文共同一作薛智威与周天贶（来源：该团队）

研究团队通过建立光子传播对称性模型，未采用传统的电训练反向传播范式，化“反”为“前”。

这种新范式挣脱了前向-反向光场传播严格对齐的束缚，基于通用的光计算原位训练系统，突破了电训练架构对物理光计算的掣肘。

该研究为光计算领域提供了大规模神经网络精准、高效的训练方法，为智能光计算系统的设计及开发提供了新思路，为光的计算能力开拓了新边界。

审稿人对该研究评价称：“本文中提出的想法非常新颖，此类光学神经网络的训练过程是前所未有的。（作者们）所提出的方法不仅有效，而且容易实现。因此，有望成为训练光学神经网络和其他光学计算系统的、广泛采用的工具。”

图丨太极-II 智能光训练概念示意图（来源：该团队）

近日，相关论文以《光神经网络全前向训练》（Fully forward mode training for optical neural networks）为题，发表于 Nature[1]。

清华大学博士生薛智威、博士后周天贶是共同一作，方璐教授、戴琼海教授担任共同通讯作者。清华大学博士生徐智昊、之江实验室虞绍良博士参与了本项研究。

图丨相关论文（来源：Nature）

释放智能光计算的“训练之能”

近年来，随着 ChatGPT、Sora、Llama 等 AI 大模型的快速发展和广泛应用，也对提升算力提出了越来越高的要求。

训练一般的神经网络需要几百到几千万的参数量，而大模型的参数量则多至几十亿级别。

一直以来，电子计算芯片为模型规模的不断发展提供支撑，但不可忽视的问题在于，电子计算架构下的高算力，也意味着高能耗。

以训练 GPT-3 为例，根据美国斯坦福人工智能研究所发布的《‌2023 年人工智能指数报告》‌，训练大模型一次所需的耗电量为 1287 兆瓦。

因此，在现阶段发展 AI 不仅仅需要解决技术问题、算力问题，还意味着面临着能源问题，此前的范式已不能很好地解决当下难题，因此发展新兴智能计算范式应运而生。

光具有干涉、衍射等多维计算模态，以光为计算载体，可以用光的可控传播来构建计算模型。相较于电计算，在完成相同计算的前提下，光计算能够以更快的速度和更低的能耗，来实现相关神经网络。

这意味着，光计算以其算力高和能耗低的优势作为智能计算的“潜力股”，为后摩尔时代带来新的希望。

推理和训练是 AI 大模型核心能力的两个重要阶段，正是看到了这一点，该团队在推理和训练方面进行了同期的并行研究。

今年 4 月，他们在 Science 报道了通用智能光计算芯片“太极”[2]，第一次推动光计算从原理验证走向了大规模实验应用。其系统级能效为 160TOPS/W，让人们看到复杂智能任务的推理的更多可能性。

然而，初代“太极”（太极-I）的相关研究以推理为核心，尚未将智能光计算的“训练之能”释放，本次研究中的太极-II 重点则在训练。

摆脱对电计算离线训练的依赖，突破现有支撑训练光网络的规模上限

与模型推理相比，模型训练对算力的需求更为迫切，然而现有光神经网络的训练严重依赖电计算离线建模优化。

电训练架构要求与前向-反向传播模型高度匹配，因此对光计算物理系统的精准对齐提出了更高的要求。其往往会导致梯度计算具有挑战、离线建模慢、映射误差大，极大地限制了光训练的规模和效率。

为解决上述问题，研究人员提出了正反互易和光学共通的解决思路。

·正反互易：全前向智能光计算训练架构。

受物理学对称性的启发，研究人员建立了“空间互易-时间反演”的双对称光传播模型，提出了全前向智能光计算在线训练架构。

薛智威解释说道：“将梯度下降中的反向传播化为光学系统的前向传播，两次前向传播实际上走的是完全相同的路径，因而具备天然对齐的特性，从而保障了物理梯度的计算精度。”

该架构突破了对电计算离线建模训练的依赖，将神经网络训练等效映射为光的前向传播，其高速低功耗特性极大地提升了训练的速度与能效，为支撑大规模网络训练奠定了基础。

·光学共通：通用智能光训练赋能复杂系统。

从波动光学的基本原理出发，该团队提出了多尺度光学系统的通用可微分神经表征，以调制和传播来构建任意光学系统。

研究人员建立了物理光系统的调制-传播与神经网络的激活-连接之间的映射关系，即调制模块的训练可驱动任意网络的权重优化，从而保障了训练的速度与能效。

薛智威表示：“通过这种新型光学系统，我们为复杂物理系统的在线训练提供了‘光速’解决方案。”

图丨全前向智能光计算训练架构（来源：Nature）

系统实测结果表明，太极-II 智能光训练架构在大规模学习、复杂场景智能成像、拓扑光子学等方面都展现了出色的性能。

具体来说：

在大规模学习领域，为“不容易兼得”的计算精度与效率提供解决方案。相较于此前数百万参数的光网络训练速度，太极-II 的光训练速度快 1 个数量级，而代表性智能分类任务的准确率则高出 40%。

在复杂场景智能成像领域，达到千赫兹帧率的智能成像效果，其成像效率提高了 2 个数量级。

在拓扑光子学领域，太极-II 在不基于其他模型先验条件下，能自动搜索非厄米奇异点，让高效、精准解析复杂拓扑系统成为可能。

图丨通用智能光训练赋能复杂系统（来源：Nature）

该研究始于 2021 年年末，最初，他们完成了线性网络，但当推进非线性大规模网络时，遇到了较大挑战。

虽然二者从理论上相对一致，但在实际过程中完成非线性网络并不容易，有时候课题组成员接连好几周都没有任何实验进展。

目前，薛智威在清华大学电子系博士三年级在读，据介绍，这是他博士期间的第一项工作，这项研究花了整整 3 年的时间认真打磨才最终开花结果。

回忆研究的实验过程，薛智威说道：“在北京冬日的一个凌晨，经过了很长时间的优化和调试后，此前一直不能很好工作的系统终于能够工作了，这意味着在实验上已经接近理论。我记得从实验室走出来的时候，感觉风都是甜的。”

有望为后摩尔时代的 AI 大模型提供算力支撑

受物理光学特性启发，太极-II 提出了一种不基于电训练架构的技术路径。

以全前向光学传播模型来求解大规模网络训练难题，克服了计算精度差、训练速度慢、能量效率低的瓶颈，进而支撑多尺度复杂光学系统的高效、高精度在线训练。

据了解，当下研究团队已成功完成原理样片的制备，并向智能光芯片工业化方向推进，正在包括无人机、无人车、机器人等在内的端侧智能系统进行应用部署。

技术从学术界到工业界的转化是一个复杂的过程。尽管太极-II 芯片本身的能耗极低，但片上光源、片上光存储和外围电子设备的工程化进程方面，仍存在相关挑战。

为进一步探索和发展这些技术，研究团队正积极推进与产业界和科研机构的密切合作，以推动光芯片系统的进一步集成优化。

“我们希望在 2-3 年内实现光芯片与外围设备完整封装的产品级光计算系统，逐步实现芯片从专用应用到通用应用的转变。”薛智威表示。

太极-II 的面世，为智能光计算在大规模训练带来新的曙光，它与太极-I 将共同打造光算力新底座，为 AI 大模型训练和推理提供算力发展的新方案。

“太极”不仅是智能光计算系列芯片，更是光计算系统特有对偶属性的辩证协作架构。太极 I 和 II 正如其命名一样，既是分立的两个部分，又组合成一个统一的整体。

该团队通过太极-I 和太极-II 的相关研究，完成了对 AI 推理和训练的探索，它们又共同构成了大规模智能计算的整个生命周期，开启了智能光计算的新时代。

本次研究的完成，让人们看到了包括太极系列光芯片在内的智能光算力船帆的不懈航行，智能光计算平台终以解决 AI 算力所面临的问题为目标，将驶向 AI 算力的星辰大海，为绿色环保和高效率的 AI 大模型、通用人工智能等提供一种全新方案。