[论文解读] Optical Transformers

Maxwell G. Anderson, Shi-Yuan Ma+3 more|arXiv (Cornell University)|Feb 20, 2023

Neural Networks and Reservoir Computing被引用 6

一句话总结

本文提出光学变换器——一种利用光学矩阵-向量乘法实现的神经网络加速器，相比数字系统可实现显著更高的能效。通过在原型光学硬件系统上展示准确的推理，并对大规模模型进行仿真，作者表明光学系统具有 1/d 的每 MAC 能量优势，对于万亿参数模型，其能效最高可达最先进的数字处理器的 8,000 倍。

ABSTRACT

The rapidly increasing size of deep-learning models has caused renewed and growing interest in alternatives to digital computers to dramatically reduce the energy cost of running state-of-the-art neural networks. Optical matrix-vector multipliers are best suited to performing computations with very large operands, which suggests that large Transformer models could be a good target for optical computing. To test this idea, we performed small-scale optical experiments with a prototype accelerator to demonstrate that Transformer operations can run on optical hardware despite noise and errors. Using simulations, validated by our experiments, we then explored the energy efficiency of optical implementations of Transformers and identified scaling laws for model performance with respect to optical energy usage. We found that the optical energy per multiply-accumulate (MAC) scales as $\frac{1}{d}$ where $d$ is the Transformer width, an asymptotic advantage over digital systems. We conclude that with well-engineered, large-scale optical hardware, it may be possible to achieve a $100 imes$ energy-efficiency advantage for running some of the largest current Transformer models, and that if both the models and the optical hardware are scaled to the quadrillion-parameter regime, optical computers could have a $>8,000 imes$ energy-efficiency advantage over state-of-the-art digital-electronic processors that achieve 300 fJ/MAC. We analyzed how these results motivate and inform the construction of future optical accelerators along with optics-amenable deep-learning approaches. With assumptions about future improvements to electronics and Transformer quantization techniques (5$ imes$ cheaper memory access, double the digital--analog conversion efficiency, and 4-bit precision), we estimated that optical computers' advantage against current 300-fJ/MAC digital processors could grow to $>100,000 imes$.

研究动机与目标

• 解决训练和推理大规模深度学习模型（尤其是变换器）日益增长的能耗问题。
• 探索光学计算是否能为大规模变换器模型提供一种可扩展、高能效的替代方案，以替代数字电子处理器。
• 证明尽管存在噪声和误差，光学硬件仍能准确执行变换器运算。
• 建立光学能效与模型性能在不同模型规模和光学能效预算下的缩放规律。
• 基于实验和仿真结果，为未来光学加速器及适用于光学的深度学习架构设计提供指导。

提出的方法

• 使用基于空间光调制器（SLM）的原型系统，在小规模光学实验中执行代表变换器推理的矩阵-向量乘法操作。
• 从光学硬件中收集真实世界的噪声、误差和非理想性数据，用于校准完整光学变换器的高保真仿真。
• 通过系统性地引入实验测量所得的误差、噪声以及权重/输入非理想性，对光学变换器推理进行仿真。
• 对光学神经网络（ONN）加速器的总能耗进行建模，包括每 MAC 的光学能耗、内存访问能耗以及数模转换开销。
• 评估能效的缩放规律，作为模型宽度（d）和总光学能耗的函数，表明光学每 MAC 能耗遵循 1/d 缩放规律。
• 在不同假设下（如电子器件性能提升、内存访问优化、4 位量化）预测性能和能效优势。

Figure 1: General scheme of an optical neural network (ONN) accelerator. Data is encoded and fed into the network, and the output is subject to shot noise. There are many experimental realizations of ONN accelerators such as Mach-Zehnder Interferometer meshes (Shen et al., 2017 ; Bogaerts et al., 20

实验结果

研究问题

• RQ1 尽管存在固有噪声和误差，光学硬件能否准确执行变换器所需的线性运算？
• RQ2 光学变换器的能效如何随模型规模和光学能耗的变化而缩放？
• RQ3 在大规模变换器推理中，光学处理器相比数字处理器在理论和实际能效方面具有多大的优势？
• RQ4 训练和量化方案如何影响光学变换器中的光子使用量和性能表现？
• RQ5 面向万亿参数模型的未来光学加速器，其关键设计权衡与缩放极限是什么？

主要发现

• 在真实硬件上，即使存在噪声和误差，光学矩阵-向量乘法仍能以高精度完成，该结论通过实验和仿真得到验证。
• 光学每 MAC 操作的能耗与模型宽度 d 呈 1/d 缩放关系，为每 MAC 能耗恒定的数字系统提供了渐近能效优势。
• 对于当前大规模变换器（如 1750 亿参数模型），光学系统相比最先进的数字处理器（300 fJ/MAC）可实现 100 倍的能效优势。
• 当模型和光学硬件均扩展至千万亿参数规模时，光学系统对数字处理器的能效优势可超过 8,000 倍。
• 在乐观假设下（内存访问成本降低 5 倍、数模转换效率翻倍、采用 4 位量化），能效优势可能超过 100,000 倍。
• 为维持 8 位数字模型性能所需的光子使用量随模型规模呈亚线性增长，但能效对权重和激活统计特性高度敏感，而这些特性又受量化和训练方案的影响。

Figure 2: Optical Transformer evaluation: prototype hardware; simulator model; Transformer architecture. Bottom: typical Transformer architecture, but with ReLU6 activation. Top Left: experimental spatial light modulator (SLM)-based accelerator setup. From some layers—marked with a laser icon—we sam