清华大学电子工程系方璐教授课题组
自动化系戴琼海院士课题组
另辟蹊径
首创了全前向智能光计算训练架构
研制了“太极-II”光训练芯片
实现了光计算系统
大规模神经网络的高效精准训练
该研究成果
以“光神经网络全前向训练”为题
于北京时间8月7日晚
在线发表于《自然》期刊
清华大学电子系为论文第一单位,方璐教授、戴琼海教授为论文的通讯作者,清华大学电子系博士生薛智威、博士后周天贶为共同一作,电子系博士生徐智昊、之江实验室虞绍良博士参与了本项工作。本课题受到国家科技部、国家自然科学基金委、北京信息科学与技术国家研究中心、清华大学-之江实验室联合研究中心的支持。
Nature审稿人在审稿评述中指出“本文中提出的想法非常新颖,此类光学神经网络(ONN)的训练过程是前所未有的。所提出的方法不仅有效,而且容易实现。因此,它有望成为训练光学神经网络和其他光学计算系统的广泛采用的工具。”
巧用对称,助力光计算摆脱GPU依赖
近年间,具有高算力低功耗特性的智能光计算逐步登上了算力发展的舞台。通用智能光计算芯片“太极”的问世便是其中的一个缩影,它首次将光计算从原理验证推向了大规模实验应用,以160TOPS/W的系统级能效为大规模复杂任务的“推理”带来了曙光,但未能够释放智能光计算的“训练之能”。
相较于模型推理而言,模型训练更需要大规模算力。然而,现有的光神经网络训练严重依赖GPU进行离线建模并且要求物理系统精准对齐。正因如此,光学训练的规模受到了极大的限制,光高性能计算的优势仿佛被禁锢在无形的枷锁之中。
在这个时候,方璐、戴琼海课题组找到了“光子传播对称性”这把钥匙,将神经网络训练中的前向与反向传播都等效为光的前向传播。
据论文第一作者、电子系博士生薛智威介绍,在太极-II架构下,梯度下降中的反向传播化为了光学系统的前向传播,光学神经网络的训练利用数据-误差两次前向传播即可实现。两次前向传播具备天然的对齐特性,保障了物理梯度的精确计算。如此实现的训练精度高,便能够支撑大规模的网络训练。
由于不需要进行反向传播,太极-II架构不再依赖电计算进行离线的建模与训练,大规模神经网络的精准高效光训练终于得以实现。
全前向智能光计算训练架构
高效精准,智能光训练事事可为
以光为计算媒介,以光的可控传播构建计算模型,光计算天然具有高速和低功耗的特性,利用光的全前向传播实现训练能够极大的提升光网络训练的速度与能效。
论文研究表明,太极-II能够对多种不同光学系统进行训练,并在各种任务下均表现出了卓越的性能。
大规模学习领域:突破了计算精度与效率的矛盾,将数百万参数的光网络训练速度提升了1个数量级,代表性智能分类任务的准确率提升40%。
复杂场景智能成像:弱光环境下(每像素光强度仅为亚光子)实现了能量效率为5.40×10^6 TOPS/W的全光处理,系统级能效提升6个数量级。在非视域场景下实现了千赫兹帧率的智能成像,效率提升2个数量级。
拓扑光子学领域:在不依赖任何模型先验下可自动搜索非厄米奇异点,为高效精准解析复杂拓扑系统提供了新思路。
通用智能光训练赋能复杂系统
携手太极,推动AI光算力扬帆远航
太极-II的面世,继太极I芯片之后进一步揭示了智能光计算的巨大潜力。
如两仪分立,太极I和II分别实现了高能效AI推理与训练;
又如两仪调和,太极I和II共同构成了大规模智能计算的完整生命周期。
方璐表示:“‘定两仪太极之道,合正反乾坤之法’,我们这样形容太极系列这一组辩证协作架构,我们相信,它们将合力为未来AI大模型注入算力发展的新动力,构建光算力的新基座”。
在原理样片的基础上,研究团队正积极地向智能光芯片产业化迈进,在多种端侧智能系统上进行了应用部署。
可以预见,经过太极系列在内的光计算领域的不懈努力,智能光计算平台将有望以更低的资源消耗和更小的边际成本,为人工智能大模型、通用人工智能、复杂智能系统的高速高能效计算开辟新路径。
来源|电子系
排版|胡文悦
责编|苑洁
审核|许亮
清华大学版权所有