基于相干光的并行光学矩阵-矩阵乘法实现高效张量处理新范式

《Nature Photonics》：Direct tensor processing with coherent light

本刊推荐：研究人员针对现有光学计算方法难以高效处理张量运算的瓶颈，开展了并行光学矩阵-矩阵乘法（POMMM）研究。该技术通过单次相干光传播实现全并行张量处理，在实数和复数域均与GPU运算高度一致，成功应用于卷积神经网络（CNN）和视觉Transformer（ViT），为下一代光学计算提供了兼具可扩展性与高效率的通用加速框架。

　　

在人工智能技术飞速发展的今天，神经网络模型的训练和推理过程对计算资源提出了巨大需求。作为核心运算的张量处理通常依赖图形处理器（GPU）张量核心进行矩阵-矩阵乘法（MMM），但这种方法面临内存带宽需求大、功耗高以及张量核心资源利用不足等挑战。与之相比，光学计算凭借其大带宽、高并行性和低能耗的固有特性，成为加速神经网络计算的理想平台。然而，现有光学计算方法多针对标量运算优化，难以高效处理基于张量的任务，这限制了其在神经网络等复杂应用中的实用性。

为突破这一瓶颈，上海交通大学电子信息与电气工程学院Guo Xuhan团队在《Nature Photonics》发表了题为"Direct tensor processing with coherent light"的研究论文。该研究开发出并行光学矩阵-矩阵乘法（POMMM）新范式，通过单次相干光传播实现全并行张量处理，显著提升了光学计算的理论计算能力和效率。

研究人员通过巧妙利用傅里叶变换的时移和频移特性，构建了POMMM的核心原理。首先将矩阵A的元素编码到空间光场的振幅和位置上，每行采用特定梯度的线性相位编码；接着对携带复振幅信号的光场进行列方向傅里叶变换，由于时移特性，光场的每行代表矩阵A所有行的叠加；然后通过振幅调制执行该"混合矩阵"与矩阵B^T的Hadamard积，实现所有行-列Hadamard积的并行计算；最后进行行方向傅里叶变换完成求和。得益于频移特性和行维度中不同线性相位调制的存在，矩阵A不同行的贡献自然分离到最终计算结果的不同空间频率位置。

实验验证方面，研究团队搭建了光学原理验证原型系统。使用振幅空间光调制器（ASLM）编码矩阵A，通过4f光学系统成像到相位空间光调制器（PSLM）施加线性相位调制，随后光束通过柱面透镜组件实现行方向成像和列方向傅里叶变换，在第二个ASLM上进行矩阵B^T的振幅调制，最后通过另一组柱面透镜完成行方向傅里叶变换和列方向成像，输出矩阵（AB）^T由定量互补金属氧化物半导体（qCMOS）相机捕获。整个张量操作完全并行，单次曝光即可同时生成所有数值。

POMMM的演示与评估

实验结果表明，POMMM在不同场景下与基于GPU的MMM保持高度一致，包括不同大小的非负矩阵（如对称矩阵和上三角矩阵）以及包含共轭矩阵对的实值矩阵。大样本定量分析显示，平均绝对误差小于0.15，归一化均方根误差小于0.1，证实了POMMM框架的准确性和可靠性。

POMMM在GPU兼容张量处理中的应用

凭借POMMM与GPU张量核心间的高度一致性，该技术理论上支持直接部署标准GPU神经网络架构。研究团队在POMMM仿真和原型上使用MNIST和Fashion-MNIST（FMNIST）数据集进行了CNN和ViT网络的直接推理实验。这些架构包含现代神经网络中常见的三种张量处理步骤：多通道卷积、多头自注意力和多样本全连接层。

实验结果显示，所有平台上的推理输出高度一致，表明POMMM支持广泛的张量处理操作，并能直接部署GPU训练的权重。此外，研究团队使用两个连续的仿真POMMM单元直接执行无约束GPU训练权重的所有线性运算，实现了与GPU结果高度一致的推理精度。通过基于U-Net模型的图像风格迁移任务进一步验证了POMMM的可扩展性，其中最大MMM规模达到[256,9,216] × [9,216,256]。

多波长扩展前景与性能分析

由于POMMM范式仅依赖相干光的振幅和相位调制，原则上支持波长复用，从而通过单次传播过程实现张量-矩阵运算。这使得能够对[L,N,M]张量与[M,N]矩阵进行并行光学处理（四阶并行）。因为光场的空间频率分布取决于波长，经历相同线性相位调制的不同波长分量将被映射到空间频率域的不同位置。

研究团队通过两个[5,5]复数值矩阵的实部和虚部（分别调制到540纳米和550纳米波长）验证了波长复用POMMM。结果显示，POMMM能够通过单次光学传播在不同尺度和数据类型上稳健执行MMM，其与多波长扩展的内在兼容性凸显了其向并行张量-矩阵乘法扩展的强大潜力。

研究结论与讨论

与现有通用光学计算范式相比，POMMM在单波长和多波长扩展下均展现出显著的理论计算优势。实验证明的计算规模证实了POMMM在实际应用中的可扩展性。尽管原型系统的实际能效仅为2.62 GOPJ^-1，但由于POMMM仅需被动相位调制（排除数据输入和输出），原则上与各种自由空间光学计算设备兼容。其极高的理论计算并行性使得与高速、大规模专用光子硬件集成时能够显著提升有效性能，使POMMM成为下一代光学计算平台的理想计算范式。

该研究通过理论仿真和物理光学原型验证了POMMM范式，证明了其与标准GPU基于MMM在各种输入矩阵尺度上的高度一致性。基于基本计算单元，研究人员开发了GPU兼容的光学神经网络（ONN）框架，并展示了不同GPU神经网络架构（如CNN和ViT）的直接光学部署，融入了多通道卷积、多头自注意力和多样本全连接层等各种张量操作。研究还探索了POMMM的可扩展性，包括其数据类型、复杂神经网络和任务的大规模计算能力，以及用于高阶张量处理的多波长复用。通过全面比较现有光学计算范式，凸显了POMMM范式作为变革性方法的优越性能，以及在张量处理中理论上提高的效率和多功能性。

研究结果表明，POMMM有潜力实现更复杂和更高阶的通用并行光学计算，满足未来计算需求，为推进下一代光学计算提供了可扩展、高效率的基础框架。这一创新技术不仅解决了现有光学计算方法在张量处理方面的关键限制，还为人工智能和神经网络应用提供了新的加速路径，标志着光学计算领域的重要进展。