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

1 Introduction . The escalating scales of deep learning models and datasets have brought increased demand for computing capacities in electronic processors . Stringent performance and efficiency constraints in practical applications raise a surging need to develop more efficient computing solutions . As a promising substitute for conventional electronics , optical neural networks ( ONNs ) have attracted extensive research interests owing to their sub-nanosecond latency and attojoule/multiply-accumulate operation ( MAC ) energy efficiency [ 42 , 6 , 51 , 41 , 54 , 14 , 15 ] , shown in Figure 1 ( a ) . However , robustness and trainability are still critical issues for photonic AI engines [ 58 , 22 , 60 ] . Due to the analog computing nature of ONNs , the photonic DNN model inevitably suffer from performance degradation or even complete malfunction [ 58 , 60 ] with the existence of manufacturing errors , non-ideal device controls , and undesired circuit noises , shown in Figure 1 ( b ) . Though non-ideal effects can be simulated and considered during software training [ 58 , 22 ] to improve noise tolerance , the variation simulation is physically inaccurate ( especially with unknown process variations ) and prohibitively expensive , shown in Figure 1 ( c ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Recently , on-device training has become an appealing trend towards adaptable and self-learning ONNs . However , training on photonic neural chips is non-trivial and much less explored than on conventional platforms . Prior work [ 57 , 25 , 21 , 18 ] only demonstrated small prototypes , and their scalability and efficiency are rather limited . To push the limits of DNNs in optics , we propose an efficient three-stage learning framework L2ight that consists of variation-agnostic identity calibration , alternate projection-based parallel mapping , and multi-level sparse subspace learning . The main contributions of this work are four-fold , • Scalability . For the first time , an ONN learning protocol can scale up to million-level parameters under practical circuit non-ideality , over 3-order-of-magnitude more scalable than prior arts . • Efficiency . We explore multi-level sparsity in in-situ gradient evaluation to trim down unnecessary on-chip training energy and runtime cost . • Learnability . By trading redundant representability , our restricted subspace optimization can provide ONNs with enough adaptability for on-device self-learning and task transfer . • Robustness . Various practical device noises and process variations are considered in situ to facilitate noise-resilient photonic AI engines . • To our best knowledge , this is the first framework that supports on-chip training on millionparameter ONNs , over 1000× more scalable and 30× more efficient than prior art . We open-source a PyTorch-centric [ 38 ] ONN library torchonn and release our on-chip training framework at link . 2 Related Work . Optical Neural Network and Training Methods . One of recent ONN architectures adopts singular value decomposition ( SVD ) to implement matrix multiplication [ 42 ] , i.e. , y = Wx = UΣV ∗x . Cascaded 2-by-2 optical devices , i.e. , Mach-Zehnder interferometers ( MZIs ) , are used to construct unitary matrices as the product of a series of 2-dimensional unitary rotators , U ( n ) = D ∏2 i=n ∏i−1 j=1 Rij ( ϕij ) . A detailed introduction to ONNs can be found in Appendix A . Beyond offline training [ 22 ] , ONN on-chip training methods are proposed to offload the process back onto photonics [ 25 , 21 , 18 ] , shown in Table 1 . Brute-force device tuning ( BFT ) [ 42 , 59 ] and evolutionary algorithms [ 57 ] are applied to search MZI settings . An adjoint variable method ( AVM ) [ 25 ] is proposed to directly evaluate gradients using in-situ light field monitoring . Stochastic zeroth-order optimization ( ZOO ) [ 21 , 18 ] is later applied to improve the training efficiency . However , prior methods are hard to scale to larger ONNs either due to algorithmic inefficiency or unrealistic hardware complexity . Efficient NN Training Methods . Extensive work has been devoted to accelerating DNN training , among which an important branch is sparse backpropagation . Previous methods mainly focus on approximating matrix multiplication by sparsifying the pre-activation gradients [ 44 ] , forward and feedback matrices [ 1 , 39 ] , and input feature maps [ 37 ] . Quantization to the pre-activation gradients is adopted in [ 52 ] to induce sparsity by trading off quantization steps and performance . Other methods also exist , e.g. , distributed and low-precision training [ 2 , 3 , 26 ] . However , they are not readily applicable to analog photonic engines , thus not in the scope of our discussion . Subspace Neural Networks . Subspace neural networks are special DNN models with restricted parameter space but demonstrate comparable representability to classical NNs . Sparse NNs [ 23 , 50 ] , low-rank NNs [ 9 , 28 , 45 ] , structured NNs [ 11 , 30 , 48 ] , Fourier-domain NNs [ 19 , 35 , 34 ] , and general frequency-domain NNs [ 20 ] were introduced to trim down the redundancy in DNNs by restricting the NN structure , matrix rank , numerical resolution , etc . In this work , we deeply explore the trade-off between ONN learnability , trainability , and efficiency in the restricted unitary subspace . Challenges of ONN On-Chip Training . As a unique hardware-restricted optimization problem , ONN in-situ learning encounters fundamental challenges causing scalability issues in prior methods : • Lack of full-observability for in-situ light field . Tracking physical optical field on every waveguide in U and V ∗ is not scalable or practical when ONNs scale up . Per device light field monitoring and calibration [ 17 , 25 ] involves intractable hardware complexity . In practice , only Σ can be precisely monitored and efficiently tuned . • Limited input/output observability . In photonic tensor cores , for efficiency consideration , only the final output signals after UΣV ∗ can be coherently detected . Intermediate signals of a single unitary projection can not be easily read out without extra hardware support . • Inaccessible gradients for most control variables . Due to the above two limitations , it is challenging to obtain true derivatives w.r.t . the MZI rotation phases in U and V ∗ [ 42 , 21 , 18 ] , casting fundamental in-situ optimization difficulty as ONN scales up . To enable in-situ self-learning for ONNs , the proposed synergistic framework L2ight provides a scalable , efficient , and on-chip-implementable solution that overcomes those hardware restrictions . 3 Synergistic ONN On-Chip Learning Framework L2ight In this section , we give a formal description of the ONN on-chip training problem and detailed demonstration of our proposed three-stage learning flow L2ight , shown in Figure . 2 . 3.1 Understanding the ONN On-Chip Learning Problem . The ONN that supports on-chip learning is shown in Figure 3 , constructed by local storage , control units , interconnects , and photonic tensor cores with coherent I/O [ 33 , 56 ] and wavelength-division multiplexing ( WDM ) [ 55 , 46 ] for parallel processing . The target is to optimize MZI phases Φ directly on chip under variations . Formally the hardware-restricted learning problem is , Φ∗ = argmin Φ L ( W ( ΩΓQ ( Φ ) +Φb ) ; Dtrn ) , s.t . W ( Φ ) = { Wpq ( Φpq ) } p=P−1 , q=Q−1 p=0 , q=0 , Wpq ( Φpq ) = Upq ( Φ U pq ) Σpq ( Φ S pq ) V ∗ pq ( Φ V pq ) , Upq ( Φ U pq ) = D U pq 2∏ i=k i−1∏ j=1 Rpqij ( ϕ U pqij ) , V ∗ pq ( Φ V pq ) = D V pq 2∏ i=k i−1∏ j=1 Rpqij ( ϕ V pqij ) , Σpq ( Φ S pq ) = max ( |Σpq| ) diag ( · · · , cosϕSpq , i , · · · ) , Φb ∼ U ( 0 , 2π ) , Γ ∼ N ( γ , σ2γ ) . ( 1 ) The linear projection in an ONN adopts blocking matrix multiplication , where the M ×N weight matrix is partitioned into P×Q blocks of size k×k . During the optimization of Φ , we jointly consider control resolution limits Q ( · ) [ 22 , 42 ] , device process variations Γ [ 22 , 21 , 18 ] , thermal crosstalk among adjacent devices Ω [ 21 , 60 ] , and unknown phase bias due to manufacturing error Φb for in-situ robustness-aware training . A detailed non-ideality analysis is in Appendix A.3 . For practicality , robustness , and convergence consideration , we select k=9 , which is explained in Appendix F . 3.2 Identity Calibration ( IC ) : Variation-Agnostic Circuit State Preparation . After manufacturing , unknown process variations in waveguides make the initial state of PTCs unpredictable [ 47 , 60 ] . A primary task is to prepare U and V ∗ to be identity matrices . However , the calibration problem , i.e. , minΦU , ΦV ∑ p , q ( ∥Upq ( ΦUpq ) − I∥22 + ∥V ∗pq ( ΦVpq ) − I∥22 ) , is not solvable given the observability and controllability constraints on U and V ∗ . The closest auxiliary problem that we can solve is the one with absolute operations on unitaries , i.e. , minΦU , ΦV ∑ p , q ( ∥|Upq ( ΦUpq ) | − I∥22 + ∥|V ∗pq ( ΦVpq ) | − I∥22 ) . We denote those two mean square errors as MSEU and MSEV . We rewrite it as a surrogate minimization of LIC that can lead to the same solution , min Φ ∑ p , q ∥Upq ( ΦUpq ) ΣpqV ∗pq ( ΦVpq ) Σ−1pq−I∥ . ( 2 ) The optimal solution for this auxiliary problem is U = V ∗ = Ĩ , where Ĩ is not guaranteed to be an identity matrix but a more general sign-flipping matrix with arbitrary and unobservable sign flips on the same columns in U and rows in V ∗ , shown in Figure 4 ( a ) . We adopt zeroth-order optimization ( ZOO ) on ΦU and ΦV to calibrate U and V ∗ to approach Ĩ , shown in Figure 4 ( b ) . We show the converged solution of Eq . ( 2 ) with unobservable sign flips and suboptimality only has marginal impacts on the following training procedure in later sections . 3.3 Parallel Mapping ( PM ) : Alternate Projection-based Model Deployment . The target is to map the pre-trained weights W onto photonic MZI meshes W̃ ( Φ ) with high fidelity . We formulate the parallel mapping as a batched k × k-blockwise regression problem , min Φ ∑ p , q ∥W̃pq ( Φpq ) −Wpq∥22 . ( 3 ) As analyzed before , ∂W ∂ΦU and ∂W ∂ΦV are too expensive to compute in situ . We propose a parallel mapping flow with alternate zeroth-order optimization on ΦU and ΦV . After convergence , we will perform analytical op- timal singular-value projection ( OSP ) to minimize the regression error given fixed singular vectors . We show why OSP gives the optimal solution under sign flips and how to perform it on the PTC . Claim 1 . Optimal singular-value projection ( OSP ) : the optimal singular value problem , i.e. , Σopt = argminΣ ∥UΣV ∗ −W ∥ , can be analytically solved on-chip with arbitrary and unknown sign flip . Proof . Σopt =diag ( U−1W ( V ∗ ) −1 ) = diag ( U∗WV ) = diag ( ( Ĩ∗V ∗W ∗UĨ ) ∗ ) . ( 4 ) OSP can be directly achieved using the limited operation set , i.e. , { U , U∗ , V , V ∗ } , supported by the reciprocal PTC itself . Specifically , we configure V ∗ = Ĩ and Σ = I , and shine in a coherent WDM light beam that carries W from right ports . Since the coherent photonic circuit is reciprocal [ 33 ] , we can read ĨU∗W on the left ports . Then we configure U = Ĩ and Σ = I , and shine in its adjoint field from left , i.e. , W ∗UĨ∗ . We can directly read out the projected optimal diagonal on the right because the sign flips in the unitary matrices naturally cancel out on the diagonal . Figure 5 compares different ZO optimizers on this task . Coordinate-wise optimizers ( ZCD [ 31 ] and ZTP [ 13 ] ) outperform the gradient-based ZGD [ 16 ] with higher accuracy and convergence speed . This procedure is highly parallel and efficient since the mapping involves no stochasticity and only happens locally within each PTC . We can also observe that OSP effectively reduces the normalized matrix distance ( ∥W − W̃ ∥22/∥W ∥22 ) and boosts the accuracy by 2-5 % almost for free .

Summary: Optical Neural Networks (ONN) with their attojoule per Multiply Accumulate energy efficiency and sub-nanosecond latency are becoming useful implementation hardware for large scale deep learning models and datasets. Despite their efficiency and speed, ONNs suffer significant loss in accuracy due to manufacturing defects/non-ideal device controls and circuit noises. Simulating these non-idealities during software training is not scalable.   Training ONN on-chip is a hard problem that is exacerbated due to above factors. This paper tackles the problem by proposing a three-stage learning framework L2ight. The three stages are variation-agnostic calibration, alternate-projection based parallel mapping and multi-level sparse subspace learning.   L2ight enables scaling up on-chip ONN training, for the first time to models with 10M params (1000x scalable than baseline). This scalability is built on top of an efficient training algorithm utilizing multi-level sparsity and a restricted subspace optimization step which enables sufficient adaptability for on-device learning and task transfer. L2ight also achieves 30x higher efficiency than prior art. Lastly, the subspace learning step is agnostic to device/process variation and mapping noise enabling robust implementation of deep neural networks (DNN).