SqueezeLight: A Multi-Operand Ring-Based Optical Neural Network With Cross-Layer Scalability
Jiaqi Gu, Chenghao Feng, Hanqing Zhu, Zheng Zhao, Zhoufeng Ying, Mingjie Liu, Ray T. Chen, David Z. Pan
Abstract
Optical neural networks (ONNs) are promising hardware platforms for next-generation artificial intelligence acceleration with ultrafast speed and low-energy consumption. However, previous ONN designs are bounded by one multiply–accumulate operation per device, showing unsatisfying scalability. In this work, we propose a scalable ONN architecture, dubbed <monospace xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">SqueezeLight</monospace> . We propose a nonlinear optical neuron based on multioperand ring resonators (MORRs) to squeeze vector dot-product into a single device with low wavelength usage and built-in nonlinearity. A block-level squeezing technique with structured sparsity is exploited to support higher scalability. We adopt a robustness-aware training algorithm to guarantee variation tolerance. To enable a truly scalable ONN architecture, we extend <monospace xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">SqueezeLight</monospace> to a separable optical CNN architecture that further squeezes in the layer level. Two orthogonal convolutional layers are mapped to one MORR array, leading to order-of-magnitude higher software training scalability. We further explore augmented representability for <monospace xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">SqueezeLight</monospace> by introducing parametric MORR neurons with trainable nonlinearity, together with a nonlinearity-aware initialization method to stabilize convergence. Experimental results show that <monospace xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">SqueezeLight</monospace> achieves one-order-of-magnitude better compactness and efficiency than previous designs with high fidelity, trainability, and robustness. Our open-source codes are available at <uri xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">https://github.com/JeremieMelo/SqueezeLight</uri> .