提高集成光子电路的计算性能，清华提出了一种衍射图神经网络框架

　　编辑 | 萝卜皮
　　光子神经网络使用光子而不是电子执行受大脑启发的计算，以实现显著提高的计算性能。然而，现有架构只能处理具有规则结构的数据，而无法泛化到欧几里得空间之外的图结构数据。
　　清华大学的研究人员提出了衍射图神经网络 (DGNN)，这是一种基于衍射光子计算单元 (DPU) 和片上光学器件的全光学图表示学习架构，以解决这一限制。具体来说，图节点属性被编码到带状光波导中，由 DPU 转换，并由光耦合器聚合以提取其特征表示。
　　DGNN 在光速光学消息通过图结构期间，捕获节点邻域之间的复杂依赖关系。他们通过基准数据库演示了 DGNN 在节点和图级分类任务中的应用，并获得了卓越的性能。该研究团队的工作为使用深度学习设计用于高效处理大规模图形数据结构的专用集成光子电路开辟了一个新方向。
　　该研究以「 All-optical graph representation learning using integrated diffractive photonic computing units 」为题，于 2022 年 6 月 15 日发布在《 Science Advances 》。
　　深度学习技术在广泛的人工智能（AI）应用中取得了巨大进步，包括计算机视觉、语音识别、自然语言处理、自动驾驶汽车、生物医学科学等。其核心是利用多层神经网络从大数据中学习层次和复杂的抽象，这是由集成电子计算平台（如中央处理器、图形处理单元（GPU）、张量处理单元和现场可编程门阵列）的不断发展推动的。
　　然而，电子计算性能已接近物理极限，难以跟上人工智能发展需求的增长，这是需要大规模深度神经模型的广泛应用中的普遍困境。
　　近年来，光子计算的研究越来越受到关注，以光子为计算介质，利用其高并行性、最低功耗和光速信号处理的先进特性来构建光子神经网络。已经提出了许多光子神经网络架构来促进复杂的神经启发计算，例如衍射神经网络、光学干涉神经网络、光子尖峰神经网络和光子储存计算。
　　现有架构在处理具有矢量或网格状图像形式的规则结构的数据方面最为成功。然而，各种科学领域分析的数据超出了这种基本的欧几里得领域。作为典型代表，编码复杂系统中实体（即节点）之间丰富关系（即边）的图结构数据在现实世界中无处不在，小到化学分子，大到大脑网络。
　　为了处理图结构数据，图神经网络（GNN）已经发展成为一类广泛的新方法，能够集成局部节点特征和图拓扑结构进行表示学习。
　　在这些模型中，基于消息传递的 GNN 具有灵活性和效率的主要优势，它通过在图节点处生成神经消息并将边缘传递给其邻居进行特征更新。它已成功应用于许多基于图的应用，包括分子属性预测、药物发现、基于骨架的人体动作识别、时空预测等。然而，如何有效地利用光子计算来使基于图的深度学习受益仍然很大程度上尚未探索。
　　在这里，研究人员提出了衍射 GNN (DGNN)，这是一种新颖的光子 GNN 架构，可以在图形结构数据上执行光学消息传递。DGNN 建立在用于生成光学节点特征的集成衍射光子计算单元 (DPU) 的基础之上。每个 DPU 包括用金属线实现的连续衍射层，以将节点属性转换为光学神经信息，其中部署条形光波导以对输入节点属性进行编码并输出转换结果。从节点邻域发送的光学神经信息使用光耦合器聚合。
　　在 DGNN 架构中，DPU 可以水平级联以扩大感受野，从而从任意大小的相邻节点中捕获复杂的依赖关系。此外，DPU 还可以垂直堆叠以提取更高维的光学节点特征，以提高其学习能力，这受到许多现代深度学习模型（例如 Transformer 和图注意力网络）中使用的多头策略的启发。
　　在这种可扩展的光学消息传递方案的基础上，研究人员首先演示了半监督节点分类任务，其中 DGNN 提取的光学节点特征被输入光学或电子输出分类器以确定节点类别。结果表明，该光学 DGNN 在合成图模型和三个真实世界图基准数据集上的电子 GNN 方面实现了具有竞争力甚至更出色的分类性能。
　　此外，DGNN 还支持图级分类，其中额外的 DPU 用于将全光节点特征聚合为图级表示进行分类。基于骨架的人类动作识别的结果证明了该架构对图分类任务的有效性。
　　图示：光学 DGNN 的架构。（来源：论文）
　　稀缺的训练标签
　　研究人员分析了 DGNN 在训练标签大小有限的情况下的有效性，这是半监督学习中的常见情况。在相同的架构设置下，他们比较了 DGNN 在不同尺寸训练标签下电子模型基线的性能，包括每类 1、5、10、15、20 和 25 个标签。
　　通过对每个尺寸的训练标签进行 10 次评估，绘制了带有误差条的测试精度条形图。对衍射调制层进行二值化有助于克服网络训练期间的局部最小问题，并提高分类精度。DGNN 架构优于所有标签稀缺设置的所有基线，特别是在较小的训练集大小时，例如每个类只有一个标签，这表明相对于其他电子计算方法具有更高的泛化能力。
　　图示：具有稀缺训练标签的 Amazon Photo 上的分类。（来源：论文）
　　带有锥形输出波导的 DPU
　　锥形波导用于将较大区域的输出光场耦合到集成 DPU 的输出端口。锥形输出波导的改进耦合效率使光电探测器能够接收更多的光功率并提高光电转换过程中的信噪比 (SNR)。更高的信噪比为分类器提供了更高质量的输入信号，保证了分类任务的稳定性。
　　锥形和单模波导的输出能量分布和模型性能的定量评估，研究人员使用 FDTD 来评估合成 SBM 图测试节点上光学特征的功率分布和训练好的 DGNN-E 模型。锥形输出波导的起始纤芯宽度经过优化，设置为 2 μm，而不是单模输出波导中使用的 500 nm。对于每个测试图节点，通过计算两个端口的输出功率相对于输入光源功率的比例得到DPU的功率传输率，从而得到所有图节点上传输率的频率直方图。具有锥形输出波导的 DPU 的平均功率传输率为 2.01%，比单模输出波导的 0.36% 高约 5.6 倍。
　　利用 DPU 的估计功率传输率，研究人员评估了片上光电探测器的光电流 SNR，公式在材料和方法中详述，在不同的输入光源功率下。同时，进一步评估了 DGNN-E 模型在 SNR 方面的测试精度，在 top-k 相邻节点设置为 16 下，通过将光电探测器噪声包括到节点特征中并重新训练电子分类器（图 S6D）。增加输入光源功率和 DPU 的功率传输率提高了光电流 SNR，并在合成 SBM 图上实现了更稳定的模型性能。
　　图示：三个基准图数据库上的半监督节点分类。（来源：论文）
　　在这项工作中，采用单轮消息传递的 PPRGo 模型直接捕获高阶邻域信息。DGNN 的计算能量效率是在输入光源功率为 10 mW 的基础上计算的，在锥形和单模输出波导中，DGNN 的光电流信噪比分别达到 34.6 和 20.2 dB，相应的模型测试精度分别为 94.4%和 92.3%。
　　DPU的计算精度
　　量化比特决定了 DPU 的计算精度，可以从光电流 SNR 中推断出来。在数字信号处理中，在模数转换器的量化过程中引入了量化误差。假设信号具有覆盖所有量化级别的均匀分布，则信量化噪声比可以表示为 SQNR = 20log10(2Q) = 6.02Q dB，其中 Q 表示量化位数。因此，使用具有 10 mW 输入光源功率的锥形输出波导的 34.6 dB 光电流 SNR 对应于约 6 个量化位。
　　计算密度和能源效率
　　值得注意的是，一旦 DGNN 架构设计经过优化和物理制造，用于计算节点和图形表示的片上光学器件以及推理过程中的光学输出分类器都是无源的。这种基于图的 AI 任务的推理过程以光速处理，仅受输入数据调制和输出检测率的限制，与电子 GNN 相比消耗的能量很少。
　　具体而言，假设 DGNN 使用 MSG（·）将每个节点的 n 维属性转换为 m 维光学神经信息，使用 AGG（·）聚合 k 个节点的光学特征，并为 C 类分类任务堆叠 P 个头部。因此，每个节点的 MSG（·）模块包含每个节点的 n m 权重矩阵，每个头部的 AGG（·）模块包含 m 维向量的 k 个节点之和，分类器包含 mP C 权重矩阵。
　　因此，DGNN 的每个推理周期包含用于特征提取的 (2nmk + mk)P 操作 (OP) 和用于分类的 2mPC 操作，即具有 (2nk + k + 2C)mP 的总操作。
　　图示：DGNN 在动作识别任务上的图分类。（来源：论文）
　　考虑到基于现有硅光子代工厂的 30GHz 数据调制和光电探测速率，DGNN 的计算速度为 (6nk + 3k + 6C)mP   10^10 OPs/s。假设典型光源功率为 10 mW，DGNN 的能量效率为 (6nk + 3k + 6C)mP   10^12 OPs/J。
　　对于n = 20，m = 2，k = 8，P = 4，C = 8的节点分类设置，计算速度为 82.6 TOP/s（Tera-Operations/s），能效为 8.26 POP/s（ Peta-Operations/s) 每瓦。对于下图中计算区域大小为 61.5 μm x 45 μm 的 DPU 模块，使用 3 2 权重矩阵执行 MSG(·) 函数，计算密度为每平方毫米 130 TOP/s。
　　假设每个 MZI 的大小为 100 μm x 100 μm，则使用片上 MZI 光子器件的相同 3 2 权重矩阵的相应实现将需要 300 μm x 200 μm 的计算区域大小，大约是 21.7 倍 。
　　请注意，最先进的 GPU Tesla V100 的能效和计算密度分别为每瓦 100 GOP/s 和每平方毫米 37 GOP/s（Giga-Operations/s）。DGNN 架构在能效上实现了四个数量级以上的提升，在计算速度上实现了三个数量级以上的提升。
　　图示：合成图上的半监督节点分类。（来源：论文）
　　架构的可扩展性
　　所提出的 DGNN 架构只执行一次 AGG(·) 以直接考虑高阶节点特征，避免了在提取远程邻域信息时的指数邻域扩展问题，并有利于学习更大图的可扩展性。原则上，架构的头数可以缩放到任意大小，并且基本DPU模块（例如图1C中的）可以水平堆叠并与Y耦合器和条形波导互连，以聚合来自任意大小邻域的光学神经消息。
　　此外，该架构具有通过进一步堆叠 DPU 模块扩展多轮光学消息传递的灵活性。DPU 模块可以通过增加每层金属层和元原子的数量来放大，DPU 的输入和输出数量可以通过额外的光调制器和波导交叉来放大。
　　图示：放大 DGNN 节点的神经消息维度。（来源：论文）
　　该架构的工作波长可以从单波长扩展到多波长，以进一步提高计算吞吐量。通过重新训练输出分类器可以减轻系统误差的累积。
　　图示：使用二进制调制训练 DGNN-E。（来源：论文）
　　此外，原位训练方法还可以通过开发具有可编程调制系数的片上 DPU 模块（例如，使用一维氧化铟锡进行调制）来解决系统错误并提高训练效率。
　　限制和未来的工作
　　在本研究中，DGNN 中的光学特征聚合是使用 2 1 光学 Y 耦合器实现的，其组合比为 50:50，不支持为不同的相邻节点分配不同的权重，即加权和。尽管平均特征聚合在节点和图级分类任务中已经取得了显著的性能，但使用加权和的消息传递可以进一步提高模型容量，并且可以使用片上调幅器（例如相变材料）来实现。
　　另一个限制是所提出的 DGNN 架构使用线性模型进行光学消息传递。尽管现有的工作已经证明了实现光学非线性激活函数的可能性，但非线性操作在 GNN 中并不重要，正如之前的工作所研究的那样。这可以通过 DGNN 在现实世界的基准数据集中实现的卓越模型性能来证明。
　　例如，在大型稀缺训练标签下，DGNN 在 Amazon Photo 上几乎达到了最先进的性能，并在稀缺标签设置下显著优于电子 GNN。因此，在 DGNN 中包含非线性激活函数留待未来工作，作为进一步增强模型学习能力的潜力。
　　综上所述，研究人员希望该工作将激发高级光学深度学习架构的未来发展，该架构具有超越欧几里得域的集成光子电路，用于高效的图形表示学习。
　　论文链接：  https://www.science.org/doi/10.1126/sciadv.abn7630