清华突破AI光训练芯片！成果登Nature

芯东西 · 公众号 · · 2024-08-10 13:36

正文

请到「今天看啥」查看全文

为AI大模型高效计算开辟新路径。

编辑 | Panken

芯东西8月10日报道，清华大学电子工程系方璐教授课题组、自动化系戴琼海院士课题组首创 全前向智能光计算训练架构 ，研制了 “太极-II”光训练芯片 ，摆脱了对离线训练的依赖，实现了光计算系统大规模神经网络的高效精准训练。相关研究成果本周发表于国际顶级学术期刊Nature。

清华大学电子工程系发文介绍说， 太极-II的面世，填补了智能光计算在大规模训练这一核心拼图的空白 。

光计算具有高算力低功耗的特性，是加速智能计算的一大前沿方向。Nature审稿人在审稿评述中提到“本文中提出的想法非常新颖， 此类光学神经网络（ONN）的训练过程是前所未有的 。所提出的方法不仅有效，而且容易实现。因此，它 有望成为训练光学神经网络和其他光学计算系统的广泛采用的工具 ” 。

清华大学电子系为论文第一单位，方璐教授、戴琼海教授为论文的通讯作者，清华大学电子系博士生薛智威、博士后周天贶为共同一作，电子系博士生徐智昊、之江实验室虞绍良博士参与了本项工作。本课题受到国家科技部、国家自然科学基金委、北京信息科学与技术国家研究中心、清华大学-之江实验室联合研究中心的支持。

论文链接：

https://www.nature.com/articles/s41586-024-07687-4

大模型时代最火AI芯片峰会来啦！！ 9月6-7日，由芯东西联合主办的2024全球AI芯片峰会将在北京举行。峰会设有数据中心AI芯片、智算集群等7大板块。目前，AMD人工智能事业部高级总监王宏强，云天励飞副总裁、芯片业务线总经理李爱军，清华大学交叉信息研究院助理教授、北极雄芯创始人马恺声等30+嘉宾/企业已确认出席并演讲。扫码报名～

01 .

巧用对称，助力光计算摆脱GPU依赖

光计算有望提高机器学习应用的速度和能源效率。然而，目前有效训练这些模型的方法受到计算机仿真的限制。

发表于国际学术顶刊Science的通用智能光计算芯片“太极” ，首次将光计算从原理验证推向了大规模实验应用，以 160TOPS/W 的系统级能效为复杂智能任务的推理带来曙光，但却未能释放光计算的“训练之能”。

相较于推理而言，模型训练更训练大量算力。电训练架构要求前向-反向传播模型高度匹配，对光计算物理系统的精准对齐提出了苛刻的要求，致使梯度计算难、离线建模慢、映射误差大，禁锢了光训练的规模与效率。

而方璐、戴琼海课题组找到了“ 光子传播对称性 ”这把钥匙，用 全前向光训练 突破了电训练架构对物理光计算的掣肘。

▲全前向智能光计算训练架构（图源：清华大学）

据论文第一作者、电子系博士生薛智威介绍，在太极-II架构下，梯度下降中的反向传播化为了光学系统的前向传播，光学神经网络的训练利用数据-误差两次前向传播即可实现。两次前向传播具备天然的对齐特性，保障了物理梯度的精确计算。如此实现的训练精度高，便能够支撑大规模的网络训练。

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

由于不需要进行反向传播，太极-II架构不再依赖电计算进行离线的建模与训练，大规模神经网络的精准高效光训练得以实现。

02 .

训练有数百万参数的光网络，

速度提升1个数量级

以光为计算媒介，以光的可控传播构建计算模型，光计算天然具有高速和低功耗的特性，利用光的全前向传播实现训练能够极大提升光网络训练的速度与能效。

系统实测结果表明，太极-II能够对多种不同光学系统进行训练，并在各种任务下均表现出了卓越的性能。

1、大规模学习领域： 突破了计算精度与效率的矛盾，将 数百万参数 的光网络训练速度提升了 1个数量级 ，代表性智能分类任务的准确率提升 40% 。

2、复杂场景智能成像： 弱光环境（每像素光强度仅为亚光子）下实现了能量效率为 5.40×10^6 TOPS/W 的全光处理，系统级能效提升 6个数量级 。在非视域等复杂场景成像应用中，实现了千赫兹帧率的智能成像，效率提升 2个数量级 。

3、拓扑光子学领域： 在不依赖任何模型先验下可自动搜索非厄米奇异点，为高效精准解析复杂拓扑系统提供了新思路。

▲通用智能光训练赋能复杂系统（图源：清华大学）

03 .

推动应用和理论进步，

为AI大模型提供算力新动力

太极-II在拓扑光子学领域亦展露出应用潜力，在不依赖任何模型先验下可自动搜索非厄米奇异点，为高效精准解析复杂拓扑系统提供了新思路。

如两仪分立，太极I和II分别实现了高能效AI推理与训练；又如两仪调和，太极I和II共同构成了大规模智能计算的完整生命周期。

“「定两仪太极之道，合正反乾坤之法」，我们这样形容太极系列这一组辩证协作架构，我们相信，它们将合力为未来AI大模型注入算力发展的新动力，构建光算力的新基座。”方璐说。

在原理样片的基础上，研究团队正积极地向智能光芯片产业化迈进，在多种端侧智能系统上进行了应用部署。

两代太极芯片相继揭示了智能光计算的巨大潜力。经过包括太极系列在内的光计算领域的不懈努力，智能光计算平台将有望以更低的资源消耗和更小的边际成本，为AI大模型、通用人工智能、复杂智能系统的高速高能效计算开辟新路径。

来源： 清华大学，Nature