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: This paper introduces L2ight, a framework for mapping a pre-trained machine learning model to an optical neural network implemented with MZIs. The framework also supports on-chip learning: this is done by fine-tuning the implemented weights to compensate for manufacturing imperfections in the chip, or to perform transfer learning from one domain to another.   To do so, L2ight first performs an identity calibration step, then map the pre-trained weights, and finally uses subspace learning to enable efficient model fine-tuning to a specific task or to a specific chip. The paper provides extensive empirical evidence motivating this approach, focusing on compute vision applications. It is able to deal with architectures orders of magnitudes larger than the current state-of-the-art.