[论文解读] Training neural networks with end-to-end optical backpropagation

James C. Spall, Xianxin Guo+1 more|arXiv (Cornell University)|Aug 9, 2023

Neural Networks and Reservoir Computing被引用 4

一句话总结

本文提出了首个完整的端到端光学反向传播框架，利用铷原子气室中的可饱和吸收体实现神经网络的训练。通过采用泵浦-探测配置，其中前向信号作为泵浦光，反向信号作为探测光，该方法实现了通过非线性激活层的全光学梯度计算，仅需极少的硬件开销，性能优于传统的软件内训练。

ABSTRACT

Optics is an exciting route for the next generation of computing hardware for machine learning, promising several orders of magnitude enhancement in both computational speed and energy efficiency. However, to reach the full capacity of an optical neural network it is necessary that the computing not only for the inference, but also for the training be implemented optically. The primary algorithm for training a neural network is backpropagation, in which the calculation is performed in the order opposite to the information flow for inference. While straightforward in a digital computer, optical implementation of backpropagation has so far remained elusive, particularly because of the conflicting requirements for the optical element that implements the nonlinear activation function. In this work, we address this challenge for the first time with a surprisingly simple and generic scheme. Saturable absorbers are employed for the role of the activation units, and the required properties are achieved through a pump-probe process, in which the forward propagating signal acts as the pump and backward as the probe. Our approach is adaptable to various analog platforms, materials, and network structures, and it demonstrates the possibility of constructing neural networks entirely reliant on analog optical processes for both training and inference tasks.

研究动机与目标

• 为解决在光学域内完全实现反向传播，特别是通过非线性激活函数的长期挑战。
• 消除在模拟光学神经网络训练过程中对数字计算或光电转换的依赖。
• 实现仅使用光学元件和过程的实时、原位光学神经网络训练。
• 证明光学反向传播可因减少“现实差距”效应而优于传统软件内训练。
• 开发一种通用且可适应的框架，适用于多种光学平台、材料和网络架构。

提出的方法

• 在铷原子气室中采用泵浦-探测过程，其中前向信号（泵浦光）使原子能级发生饱和，反向信号（探测光）测量激活函数的导数。
• 使用可饱和吸收体作为非线性激活单元，通过探测光束的透射率受泵浦光强度调制，以模拟激活函数的导数。
• 利用空间光调制器（SLMs）和柱面透镜实现光学矩阵-向量乘法（MVM），通过“汇聚”和“展开”操作实现线性层。
• 通过相同的光学装置反向传播光束实现反向传播，其中DMDs和SLMs用于实现转置权重矩阵和误差反向传播。
• 在每个训练迭代中进行三次测量（仅泵浦光、仅探测光、两者同时），以从荧光和未吸收的探测光中减去背景偏移，确保梯度估计的准确性。
• 采用相干检测与高速相机测量激活向量和误差向量，用于实时计算权重梯度并更新SLMs。

Figure 1: Illustration of optical training. (a) Network architecture of the ONN used in this work, which consists of two fully-connected linear layers and a hidden layer. (b) Simplified experimental schematic of the ONN. Each linear layer performs optical MVM with a cylindrical lens and a spatial li

实验结果

研究问题

• RQ1 是否可以在无需数字干预的情况下，完全在光学域内实现通过非线性激活函数的反向传播？
• RQ2 单一光学系统是否能够通过反向信号传播同时实现推理与训练？
• RQ3 全光学反向传播在准确性和对硬件非理想性的鲁棒性方面是否优于传统软件内训练？
• RQ4 在可饱和吸收体中采用的泵浦-探测配置是否能在反向传播过程中准确模拟非线性激活函数的导数？
• RQ5 所提出的方法是否可推广至不同的光学平台、材料和神经网络架构？

主要发现

• 光学反向传播框架成功利用铷原子气室中的可饱和吸收效应，实现了通过非线性激活层的梯度计算。
• 通过三次测量校准减去背景信号，实现了精确的梯度估计，有效补偿了荧光和残余探测光透射。
• 通过全光学反向传播训练的网络性能优于采用传统软件内方法训练的网络，表现出更小的“现实差距”引起的误差。
• 系统仅需极少的硬件开销，使用标准组件如SLMs、DMDs和单个气室即可完成前向与反向传播。
• 该方法具有通用性，可适配多种光学平台、材料和网络结构，包括基于干涉仪或其他线性层实现的结构。
• 实验结果证实，端到端光学训练是可行且有效的，标志着迈向全模拟光学机器学习的关键一步。

Figure 2: Multi-layer ONN characterisation. (a) Scatter plots of measured against theory results for MVM-1 (first layer forwards), MVM-2a (second layer forwards) and MVM-2b (second layer backwards). All three MVM results are taken simultaneously. Histograms of the signal and noise error for each MVM