近年间,具有高算力低功耗特性的智能光计算逐步登上了算力发展的舞台。此前清华大学通用智能光计算芯片“太极”首次将光计算从原理验证推向了大规模实验应用,拥有 160TOPS / W 的系统级能效,但现有的光神经网络训练严重依赖 GPU 进行离线建模并且要求物理系统精准对齐。
8月8日,据清华大学发布消息,该校电子工程系方璐教授课题组与自动化系戴琼海院士课题组再次取得了重大进展。他们首创了全前向智能光计算训练架构,并成功研制出“太极-Ⅱ”光芯片,这一成果标志着光计算系统在大规模神经网络的高效精准训练方面迈出了重要一步。
该研究成果已于2024年8月7日在线发表于国际学术期刊《自然》(Nature)。审稿人对“太极-Ⅱ”光芯片的研究成果给予了高度评价,认为“本文中提出的想法非常新颖,此类光学神经网络(ONN)的训练过程是前所未有的。所提出的方法不仅有效,而且容易实现。因此,它有望成为训练光学神经网络和其他光学计算系统的广泛采用的工具。”
清华大学电子系为论文第一单位,方璐教授、戴琼海教授为论文的通讯作者,清华大学电子系博士生薛智威、博士后周天贶为共同第一作者。本课题受到国家科技部、国家自然科学基金委、北京信息科学与技术国家研究中心、清华大学-之江实验室联合研究中心的支持。
创新之处在于全前向架构
“太极-Ⅱ”光芯片的创新之处在于其全前向智能光计算训练架构,这一架构利用空间对称性和洛伦兹互易性,消除了反向传播的需求,通过光子传播对称性,将神经网络训练中的前向与反向传播都等效为光的前向传播,从而实现了数据-误差两次前向传播即可完成训练的高效过程。原理如下:
- 将光学系统参数化,映射到神经网络上。
- 利用空间对称性和洛伦兹互易性,实现数据传播和误差传播的精准对齐。
- 基于测量数据/误差传播的输出光场计算梯度,并进行参数更新。
图:太极-II 智能光计算训练架构。a, 传统上,光学AI 系统是通过离线建模和优化来设计的,这导致其设计效率和系统性能受限。b, 一般的光学系统,包括自由空间系统和集成光子系统,都包含了调制区域(深绿色)和传播区域(浅绿色),在这些区域中,折射率分别是可调和固定的。c, 光学系统中的这些区域可以映射到神经网络表示中的权重和神经元连接,这使得可以在输入和输出之间构建一个可微分的神经网络表征(左上)。利用空间对称的互易性,数据和误差计算共享前向物理传播和测量,并在设计区域内计算在线梯度以更新折射率(右上和左下)。通过在线梯度下降,光学系统逐渐收敛(右下)。
据论文第一作者、电子系博士生薛智威介绍,在太极-II 架构下,梯度下降中的反向传播化为了光学系统的前向传播,光学神经网络的训练利用数据-误差两次前向传播即可实现。两次前向传播具备天然的对齐特性,保障了物理梯度的精确计算。由于不需要进行反向传播,太极-II 架构不再依赖电计算进行离线的建模与训练,大规模神经网络的精准高效光训练终于得以实现。
性能如何?
论文研究表明,太极-II 能够对多种不同光学系统进行训练,并在各种任务下均表现出了卓越的性能:
- 大规模学习领域:突破了计算精度与效率的矛盾,将数百万参数的光网络训练速度提升了 1 个数量级,代表性智能分类任务的准确率提升 40%。
- 复杂场景智能成像:弱光环境下(每像素光强度仅为亚光子)实现了能量效率为 5.40×10^6 TOPS / W 的全光处理,系统级能效提升 6 个数量级。在非视域场景下实现了千赫兹帧率的智能成像,效率提升 2 个数量级。
- 拓扑光子学领域:在不依赖任何模型先验下可自动搜索非厄米奇异点,为高效精准解析复杂拓扑系统提供了新思路。
此外,在非视域场景下实现了千赫兹帧率的智能成像,效率提升了2个数量级。在拓扑光子学领域,太极-II能够自动搜索非厄米奇异点,为解析复杂拓扑系统提供了新思路。
效果怎么样?
太极-II 架构还在多个领域均展现出巨大的应用潜力。
1. 光学神经网络(ONN)
- 深层 ONN 构建:太极-II 架构能够实现具有数百万参数的深层 ONN,并达到与理想模型相当的准确率。这对于光学计算至关重要,因为深层网络可以处理更复杂的任务并实现更高的性能。
- 非线性 ONN:太极-II 架构支持非线性光学神经网络的训练,通过在数据传播中引入非线性激活函数,并利用记录的函数输入输出来计算梯度,实现高效精准的训练过程。
图 :深度 ONN 的并行梯度下降训练。
2. 复杂场景智能成像
- 穿散射成像:太极-II 架构能够透过散射介质实现接近衍射极限的聚焦成像,达到更高的分辨率。这对于显微成像和宏观成像等领域具有重要意义。
- 非视域场景成像:太极-II 架构能够实现毫秒级的并行成像,并实现对非视域的目标进行全光处理。这对于非视域场景下的动态目标成像和识别具有重要意义。
图:利用 太极-II 架构,透过散射介质达到衍射极限分辨率。
3. 集成光子系统(PIC)
- 自设计 PIC 网络:太极-II 架构能够实现集成光子系统的自设计,并达到与理论相当的准确率。这对于构建高性能的光子计算系统具有重要意义。
- 非厄米系统解析:太极-II 架构能够自动搜索非厄米系统的奇异点,并无需物理模型。这对于研究和解析复杂拓扑系统具有重要意义。
图:基于太极-II 架构的集成光子系统在线训练。
产业化前景
今年4月发布的太极I具备879 T MACS/mm2的面积效率与160 TOPS/W的能量效率,首次赋能光计算实现自然场景千类对象识别、跨模态内容生成等人工智能复杂任务。太极-II的面世是继太极I芯片之后的一大突破,进一步揭示了智能光计算的巨大潜力。如两仪分立,太极I和II分别实现了高能效AI推理与训练;又如两仪调和,太极I和II共同构成了大规模智能计算的完整生命周期。
方璐表示:“‘定两仪太极之道,合正反乾坤之法’,我们这样形容太极系列这一组辩证协作架构,我们相信,它们将合力为未来AI大模型注入算力发展的新动力,构建光算力的新基座”。
当前,研究团队正积极推进智能光芯片的产业化进程,在多种端侧智能系统上进行应用部署,为人工智能大模型、通用人工智能、复杂智能系统的高速高能效计算开辟了新路径。
可以预见,智能光计算平台将以更低的资源消耗和更小的边际成本,为未来AI的发展注入新的活力。
在原理样片的基础上,研究团队正积极推进智能光芯片产业化,在多种端侧智能系统上进行了应用部署。预计智能光计算平台将以更低的资源消耗和更小的边际成本,为人工智能大模型、通用人工智能、复杂智能系统的高速高能效计算开辟新路径。