L2ight: Enabling On-Chip Learning for Optical Neural Networks via Efficient in-situ Subspace Optimization

Silicon-photonics-based optical neural network (ONN) is a promising hardware platform that could represent a paradigm shift in efficient AI with its CMOScompatibility, flexibility, ultra-low execution latency, and high energy efficiency. In-situ training on the online programmable photonic chips is appealing but still encounters challenging issues in on-chip implementability, scalability, and efficiency. In this work, we propose a closed-loop ONN on-chip learning framework L2ight to enable scalable ONN mapping and efficient in-situ learning. L2ight adopts a three-stage learning flow that first calibrates the complicated photonic circuit states under challenging physical constraints, then performs photonic core mapping via combined analytical solving and zeroth-order optimization. A subspace learning procedure with multi-level sparsity is integrated into L2ight to enable in-situ gradient evaluation and fast adaptation, unleashing the power of optics for real on-chip intelligence. Extensive experiments demonstrate our proposed L2ight outperforms prior ONN training protocols with 3-order-of-magnitude higher scalability and over 30× better efficiency, when benchmarked on various models and learning tasks. This synergistic framework is the first scalable on-chip learning solution that pushes this emerging field from intractable to scalable and further to efficient for next-generation self-learnable photonic neural chips. From a co-design perspective, L2ight also provides essential insights for hardware-restricted unitary subspace optimization and efficient sparse training. We open-source our framework at link.

[1]  Huaxi Gu,et al.  Countering Variations and Thermal Effects for Accurate Optical Neural Networks , 2020, 2020 IEEE/ACM International Conference On Computer Aided Design (ICCAD).

[2]  David A. B. Miller,et al.  Analyzing and generating multimode optical fields using self-configuring networks , 2020, Optica.

[3]  Jiayu Zhou,et al.  Subspace Network: Deep Multi-Task Censored Regression for Modeling Neurodegenerative Diseases , 2018, KDD.

[4]  Joan Bruna,et al.  Exploiting Linear Structure Within Convolutional Networks for Efficient Evaluation , 2014, NIPS.

[5]  Eduard A. Gorbunov,et al.  Stochastic Three Points Method for Unconstrained Smooth Minimization , 2019, SIAM J. Optim..

[6]  Y Fainman,et al.  Towards 100 channel dense wavelength division multiplexing with 100GHz spacing on silicon. , 2014, Optics express.

[7]  Dirk Englund,et al.  Deep learning with coherent nanophotonic circuits , 2017, 2017 Fifth Berkeley Symposium on Energy Efficient Electronic Systems & Steep Transistors Workshop (E3S).

[8]  A. Boes,et al.  11 TOPS photonic convolutional accelerator for optical neural networks , 2021, Nature.

[9]  Shanhui Fan,et al.  Training of Photonic Neural Networks through In Situ Backpropagation , 2018, 2019 Conference on Lasers and Electro-Optics (CLEO).

[10]  Kamyar Azizzadenesheli,et al.  signSGD: compressed optimisation for non-convex problems , 2018, ICML.

[11]  Benjamin G. Lee,et al.  Design and Fabrication of Low-Insertion-Loss and Low-Crosstalk Broadband $2\times 2$ Mach–Zehnder Silicon Photonic Switches , 2015, Journal of Lightwave Technology.

[12]  Xiang Zhou,et al.  Ultra-High-Capacity DWDM transmission system for 100G and beyond , 2010, IEEE Communications Magazine.

[13]  Andrea Melloni,et al.  Canceling Thermal Cross-Talk Effects in Photonic Integrated Circuits , 2019, Journal of Lightwave Technology.

[14]  Michael R. Watts,et al.  AIM Process Design Kit (AIMPDKv2.0): Silicon Photonics Passive and Active Component Libraries on a 300mm Wafer , 2018, 2018 Optical Fiber Communications Conference and Exposition (OFC).

[15]  Ray T. Chen,et al.  Toward Hardware-Efficient Optical Neural Networks: Beyond FFT Architecture via Joint Learnability , 2021, IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems.

[16]  Xiaogang Wang,et al.  Convolutional neural networks with low-rank regularization , 2015, ICLR.

[17]  Yiran Chen,et al.  Learning Structured Sparsity in Deep Neural Networks , 2016, NIPS.

[18]  Gordon Wetzstein,et al.  Inference in artificial intelligence with deep optics and photonics , 2020, Nature.

[19]  Saeed Ghadimi,et al.  Stochastic First- and Zeroth-Order Methods for Nonconvex Stochastic Programming , 2013, SIAM J. Optim..

[20]  Song Han,et al.  Deep Compression: Compressing Deep Neural Network with Pruning, Trained Quantization and Huffman Coding , 2015, ICLR.

[21]  Roland Vollgraf,et al.  Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms , 2017, ArXiv.

[22]  Philippe M. Bardet,et al.  Massively parallel amplitude-only Fourier neural network , 2020, Optica.

[23]  Cho-Jui Hsieh,et al.  A Comprehensive Linear Speedup Analysis for Asynchronous Stochastic Parallel Optimization from Zeroth-Order to First-Order , 2016, NIPS.

[24]  Reck,et al.  Experimental realization of any discrete unitary operator. , 1994, Physical review letters.

[25]  Chao Wang,et al.  CirCNN: Accelerating and Compressing Deep Neural Networks Using Block-Circulant Weight Matrices , 2017, 2017 50th Annual IEEE/ACM International Symposium on Microarchitecture (MICRO).

[26]  Yann LeCun,et al.  The mnist database of handwritten digits , 2005 .

[27]  Alex Krizhevsky,et al.  Learning Multiple Layers of Features from Tiny Images , 2009 .

[28]  Jiaqi Gu,et al.  ROQ: A Noise-Aware Quantization Scheme Towards Robust Optical Neural Networks with Low-bit Controls , 2020, 2020 Design, Automation & Test in Europe Conference & Exhibition (DATE).

[29]  Xuan Li,et al.  Parallel convolutional processing using an integrated photonic tensor core , 2021, Nature.

[30]  Xu Sun,et al.  meProp: Sparsified Back Propagation for Accelerated Deep Learning with Reduced Overfitting , 2017, ICML.

[31]  G. Lo,et al.  An optical neural chip for implementing complex-valued neural network , 2021, Nature Communications.

[32]  Arild Nøkland,et al.  Direct Feedback Alignment Provides Learning in Deep Neural Networks , 2016, NIPS.

[33]  Volker J. Sorger Massively parallel amplitude-only Fourier neural network , 2021 .

[34]  Ray T. Chen,et al.  Efficient On-Chip Learning for Optical Neural Networks Through Power-Aware Sparse Zeroth-Order Optimization , 2020, AAAI.

[35]  Pengfei Xu,et al.  E-Train: Training State-of-the-art CNNs with Over 80% Less Energy , 2019 .

[36]  Xu Wang,et al.  Chip-Scale Optical Matrix Computation for PageRank Algorithm , 2020, IEEE Journal of Selected Topics in Quantum Electronics.

[37]  Ivan V. Oseledets,et al.  Speeding-up Convolutional Neural Networks Using Fine-tuned CP-Decomposition , 2014, ICLR.

[38]  Natalia Gimelshein,et al.  PyTorch: An Imperative Style, High-Performance Deep Learning Library , 2019, NeurIPS.

[39]  Xu Han,et al.  Efficient training and design of photonic neural network through neuroevolution , 2019, Optics express.

[40]  Elad Hoffer,et al.  Scalable Methods for 8-bit Training of Neural Networks , 2018, NeurIPS.

[41]  Jiaqi Gu,et al.  Towards Area-Efficient Optical Neural Networks: An FFT-based Architecture , 2020, 2020 25th Asia and South Pacific Design Automation Conference (ASP-DAC).

[42]  Yuanzhou Yang,et al.  Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes , 2018, ArXiv.

[43]  Bernard Ghanem,et al.  A Stochastic Derivative-Free Optimization Method with Importance Sampling: Theory and Learning to Control , 2019, AAAI.

[44]  Zhenzhi Wu,et al.  GXNOR-Net: Training deep neural networks with ternary weights and activations without full-precision memory under a unified discretization framework , 2017, Neural Networks.

[45]  Tor M. Aamodt,et al.  Sparse Weight Activation Training , 2020, NeurIPS.

[46]  Yiran Chen,et al.  ReBoc: Accelerating Block-Circulant Neural Networks in ReRAM , 2020, 2020 Design, Automation & Test in Europe Conference & Exhibition (DATE).

[47]  C. Wright,et al.  Photonics for artificial intelligence and neuromorphic computing , 2020, ArXiv.

[48]  Yun Liang,et al.  Efficient Recurrent Neural Networks using Structured Matrices in FPGAs , 2018, ICLR.

[49]  Jiaqi Gu,et al.  Design Technology for Scalable and Robust Photonic Integrated Circuits: Invited Paper , 2019, 2019 IEEE/ACM International Conference on Computer-Aided Design (ICCAD).

[50]  Marc Sorel,et al.  Non-invasive monitoring and control in silicon photonics using CMOS integrated electronics , 2014, 1405.5794.

[51]  Wojciech Samek,et al.  Dithered backprop: A sparse and quantized backpropagation algorithm for more efficient deep neural network training , 2020, 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW).

[52]  Menachem Adelman,et al.  Faster Neural Network Training with Approximate Tensor Operations , 2018, NeurIPS.

[53]  Jiaqi Gu,et al.  FLOPS: EFficient On-Chip Learning for OPtical Neural Networks Through Stochastic Zeroth-Order Optimization , 2020, 2020 57th ACM/IEEE Design Automation Conference (DAC).

[54]  Puneet Gupta,et al.  Million-Channel Parallelism Fourier-Optic Convolutional Filter and Neural Network Processor , 2020, 2020 Conference on Lasers and Electro-Optics (CLEO).

[55]  Nan Hua,et al.  Universal Sentence Encoder , 2018, ArXiv.

[56]  Ryan P. Adams,et al.  Randomized Automatic Differentiation , 2020, ICLR.

[57]  Li Fei-Fei,et al.  ImageNet: A large-scale hierarchical image database , 2009, CVPR.

[58]  Jian Sun,et al.  Deep Residual Learning for Image Recognition , 2015, 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).

[59]  Luca P. Carloni,et al.  Silicon Photonics Codesign for Deep Learning , 2020, Proceedings of the IEEE.