[論文レビュー] 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.

研究の動機と目的

• 非線形活性化関数を含む、完全に光的空間で誤差逆伝播を実装するという長年の課題を克服すること。
• アナログ光的ニューラルネットワークのトレーニング中に、デジタル計算や光電変換を一切不要とする。
• 光的部品とプロセスのみを用いて、リアルタイムかつイン・サイトでの光的ニューラルネットワークのトレーニングを可能とすること。
• 光的誤差逆伝播が、『リアリティギャップ』の影響を低減することで、従来のソフトウェア内トレーニングを上回る性能を示すことを実証すること。
• さまざまな光的プラットフォーム、材料、ネットワークアーキテクチャに適用可能な汎用的かつ柔軟なフレームワークを開発すること。

提案手法

• 前方信号（ポンプ）が原子遷移を飽和させ、後方信号（プローブ）が活性化関数の微分を測定するポンプ・プローブプロセスを、ルビジウム蒸気セルで実装する。
• 非線形活性化ユニットとして吸収率の飽和特性を利用する。プローブビームの透過率がポンプ強度によって変調され、活性化関数の微分を模倣する。
• 空間光変調器（SLM）と円筒レンズを用いて「ファンイン」と「ファンアウト」操作を実現し、光的行列・ベクトル乗算（MVM）により線形層を実装する。
• 同じ光学セットアップを逆方向に光ビームを伝搬させることで逆誤差伝播を実行し、DMDとSLMを用いて転置された重み行列と誤差の逆伝播を実装する。
• 1トレーニングイテレーションあたり3回の測定（ポンプのみ、プローブのみ、両方）を実施し、蛍光と未吸収プローブからのバックグラウンドオフセットを差し引くことで、正確な勾配推定を実現する。
• 高精度なカメラを用いたコherent検出により、活性化ベクトルと誤差ベクトルを測定し、リアルタイムで重み勾配を計算し、SLMを即座に更新する。

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 1つの光学系が、逆方向信号伝搬を用いて推論とトレーニングの両方を実行できるか？
• RQ3 すべての光的誤差逆伝播は、従来のソフトウェア内トレーニングに比べ、精度とハードウェアの不確実要因に対するロバスト性において優れているか？
• RQ4 吸収率の飽和特性を用いたポンプ・プローブ構成が、逆伝播時に非線形活性化関数の微分を正確に模倣できるか？
• RQ5 提案手法は、さまざまな光的プラットフォーム、材料、ニューラルネットワークアーキテクチャに一般化可能か？

主な発見

• 光的誤差逆伝播フレームワークは、ルビジウム蒸気セルにおける吸収率の飽和特性を用いて、非線形活性化層を介した勾配計算に成功した。
• 3回の測定によるキャリブレーションによりバックグラウンド信号（蛍光および残存プローブ透過）を差し引くことで、正確な勾配推定が実現した。
• すべての光的誤差逆伝播でトレーニングされたネットワークは、従来のソフトウェア内トレーニングで得られたネットワークを上回り、『リアリティギャップ』に起因する誤差が低減された。
• 標準部品（SLM、DMD、1つの蒸気セル）を用いて、前方および後方伝搬の両方で最小限のハードウェアコストで動作した。
• 本手法は、干渉型またはその他の線形層実装に基づくさまざまな光的プラットフォーム、材料、ネットワーク構造に一般化可能で、汎用的かつ柔軟なアプローチである。
• 実験結果により、エンドツーエンドの光的トレーニングが実現可能かつ有効であることが確認され、完全にアナログな光的機械学習への重要な一歩を示した。

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