Diffractive tensorized unit for million-TOPS general-purpose computing

本文发表于 Nature Photonics,文章由清华大学戴琼海院士、方璐教授团队完成,提出并实验验证了一种名为衍射张量单元(DTU)的全新光子处理器架构,旨在解决衍射光计算长期存在的不可重构和可扩展性受限的核心难题,首次在芯片上实现了支持通用计算的高性能衍射光子处理器。核心创新点可以分为近核调制和张量化处理两个部分。

架构创新点:

在近核调制问题中,不同于直接调控衍射核心内部的数百万个静态的、固化的衍射神经元,采用动态衍射核心(DDTC)和静态衍射核心(SDTC)相结合的方式,通过局部调控实现实际传输矩阵的重构。如图所示为通过DDTC实现DTU重构的具体过程。


Extended Data Fig. 1


DTU的传输过程可以被解释为:

其中包含了信号输入和参数输入是传输系统,与的维度匹配,这个操作进一步细化为:

对于,假设等效的传递矩阵是,DDTC的设计的主要挑战是,是否存在使得对于任意给定的,至少存在一个解使得。也就是说,是否可以通过设计的静态系统,在动态参数的辅助下,使得对于输入信号,可以实现的输出,也就是传输矩阵是

通过的引入使得,也就是改变系统的功能,实现矩阵的重构。此时:

,则:

问题变成了是否可以找到对应的使得这个线性方程组有解。充分必要条件是参数通道数P不小于输出通道数O。且矩阵的秩等于通道数O。

在实现了单个DTU的基础上,通过张良话架构可以实现大规模数据的处理。DTU采用张量分解方法。将庞大的计算任务分解成多个小型的张量核心,然后映射到由多个DTC组成的计算集群中进行并行处理,赋予了DTU卓越的可扩展性。


Fig. 1


性能验证结果:

Ø  基础通用计算能力:DTU能够以10-6的均方误差完成任意1024*1024规模的矩阵乘法(基于32*32DTC的32次时域复用)。

Ø  高级AI任务实现:DTU实现了自然语言生成,跨模态识别,图像分类与视频生成这些复杂的AI应用。在实验中完成了MNIST、Fashion-MNIST图像分类和自然语言生成任务的推理实验,准确率分别达到97.7%、85.4%和58.6%,与仿真结果高度吻合,证明了架构的可行性。


Fig. 3


在自然语言生成(NLG)任务中,(a)(b)图描述了如何利用DTU实现单词预测。输入的词序列被编码成光信号,在DTU阵列中进行循环计算和调制,最终生成预测的下一个词。(c)图中的训练收敛曲线证明了DTU与Bi-LSTM/Transformer相当,具有通用性。(e)图展示了词向量随着预测的演化过程。(f)示例了句子的生成。(g)(h)展示了循环次数和输入长度对准确度的影响。(f)为中文诗歌、对联的结果。


Fig. 4


在跨模态的识别问题中,将视频帧序列化和张量化,利用循环DTU集群处理视频内容并生成文字描述。在MSVD和MSR-VTT等标准数据集上验证了DTU的性能。


Fig. 5


在真实制造的DTU芯片上实现了图像分类和NLG两种有代表性的任务。在MNIST手写数字分类任务上,芯片的总体准确率达到了97.7%。在NLG任务上,芯片的实际输出与仿真结果的偏差仅为1.9%,一系列实验的平均准确率达到了58.6%,非常接近仿真结果显示的60.5%。

原文链接:Diffractive tensorized unit for million-TOPS general-purpose computing | Nature Photonics