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

Abstract:
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.

heading: 1 Introduction
text: 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.

heading: 2 Related Work
text: 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.

heading: 3.1 Understanding the ONN On-Chip Learning Problem
text: 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.

heading: 3.2 Identity Calibration (IC): Variation-Agnostic Circuit State Preparation
text: 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.

heading: 3.3 Parallel Mapping (PM): Alternate Projection-based Model Deployment
text: 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.

heading: 3.4 Subspace Learning: Hardware-Aware Multi-Level Sparse Training
text: Besides mapping from an offline-trained model, L2ight also supports in-situ self-learning fully on chip. We name this feature as subspace learning. To make L2ight hardware-aware, we trade expensive full-space trainability for efficient subspace gradient evaluation, i.e., ∂L∂Σ which coincides with the general frequency-domain ONNs [19, 20] and subspace NN design concept [43]. Since this learning stage involves stochasticity, it turns out to be the efficiency bottleneck, especially the backward pass. Hence, we explore multi-level sparsity for efficient in-situ gradient approximation.
3.4.1 In-situ Subspace Gradient Acquisition via Reciprocity in Optics
The conventional way to compute first-order gradients w.r.t. Σ is ∂L∂Σ = diag ( U∗ ∂L∂W V ) . However, ∂L∂W = ∂L ∂yx
T requires arbitrary matrix multiplication, which is not implementable by weightstationary PTCs. Hence, we remap it as,
∂L ∂Σ = ∂L ∂yΣ
⊙ yV = ( ĨU∗
∂L ∂y
) ⊙ (ĨV ∗x). (5)
By shining in coherent WDM beams carrying the inputs and upstream gradients forward and backward through the reciprocal PTCs, respectively, as shown in Figure 6, the weight gradients can be efficiently obtained with lightweight elementwise multiplication ⊙, which can be offloaded to
electrical units. Ĩ naturally cancels out by the Hadamard product with no impacts on gradient fidelity.

heading: 3.4.2 Multi-Level Sparse Subspace Learning
text: Inspired by sparse backpropagation methods [44, 37, 49, 36, 39], we propose multi-level sparse subspace learning to cut down both energy cost and total time steps in on-chip gradient evaluation.
Balanced Feedback Sampling. To improve the efficiency of the error feedback process, i.e., W T ∂L∂y , as shown in Figure 7, we sample the feedback matrix W T ∈ RN×M with a structured sparse mask PW = cW (SW ⊗ 1) generated by the Kronecker product between a boolean mask SW ∈ {0, 1}Q×P with sparsity αW and an all-ones matrix 1, where the scaling factor cW is set to 1αW = PQ Tr(STWSW )
for unbiased estimation, proven in Appendix D. The efficiency benefits come from two aspects: (1) the structurally masked PTCs are entirely idle, directly saving energy, and (2) the product accumulation depth/step is reduced by a factor of αW , effectively trimming time steps.
However, two major downsides exist on traditional uniform and layer-wise topk sampling [39].
First, on a backward path, multiple feedback sampling operators will be cascaded, such that importance-unaware uniform sampling can lead to an exponentially large variance [37]. Second, topk sampling is overly greedy and tends to break the load balance as the feedback latency can be bottlenecked by the longest partial product accumulation path, shown in Figure 7. To tackle this, we propose a balanced top-K sampling (btopk) to draw SW from a guided distribution that locally prefers blocks with large Frobenius norm, which can be efficiently evaluated by ∥Wpq∥2F = Tr(|Σpq|2). It strikes a balance between gradient variance and bias by fine-grained row-wise top-K sampling and eliminates load-imbalance by guaranteeing the same sparsity for different rows of W T , i.e., ∑ p SW (1, :) = ∑ p SW (2, :) = · · · = ∑ p SW (Q, :). Figure 8(a), 8(b) shows the gradient approximation fidelity in terms of average angular similarity [5] and normalized matrix
distance. Our btopk-sampled weight gradients align well with the true gradients. With the unbiased (exp) normalization factor αW , btopk shows the best gradient angular similarity and inference accuracy compared with others.
Information-Preserving Column Sampling. Input feature sparsification can also effectively cut down the gradient evaluation cost [39, 37], especially for costly CONV layers. However, with traditional spatial sampling (SS) [39, 37], the input feature map x barely maintains its sparsity regularity after being transformed to flattened patches X via im2col if the kernel size is larger than 1, shown in Figure 9. Hence, we propose a novel column sampling (CS) as a better solution. We sample X using
a mask SC{0, 1}H ′W ′ with a uniform sparsity αC , which is shared across batches with negligible overhead. This leads to both information preservation and efficiency improvement. First, in Figure 9, a pixel appears in multiple columns, such that partial information can be maintained after column sampling. Second, this highly-structured column dropping directly translates to less PTC forward energy and fewer partial gradient accumulation steps. In contrast, with a spatial mask SS and spatial sparsity αS , the masked pixel will be completely dropped with poor regularity after im2col, at the cost of large variance due to information loss and almost no runtime improvement on this dense linear projection engines. Note that for CONV1×1, CS turns out to be equivalent to SS, which can simultaneously save memory and runtime. Figures 8(c), 8(d) show that our proposed
CS can obtain better gradient approximation fidelity than prior SS. Different normalization has small effects on model accuracy since feature sampling only happens locally within each layer, without any variance cascade effect. Note that simultaneous scaling by αW and αC tends to generate overly-
confident gradient approximation, which empirically leads to harmful gradient variance. Hence, we will adopt αC=1 in all experiments.
Data Sampling. After parallel mapping, the ONN is initialized fairly close to the target pre-trained model. It is reasonable and intuitive to calibrate it with a representative calibration set instead of the entire training set. Inspired by the mini-batch dropping (SMD) technique [49], we integrate this SMD technique into our framework to further explore data-level sparsity. Within one training epoch, we randomly skip each iteration with probability αD, directly translating to training time reduction.
3.5 Complexity Analysis of Three Stages in L2ight
We assume the total step in IC, PM, and SL is T1, T2, and T3, respectively. The ONN has L layers, each including an N ×N weight matrix partitioned into multiple k × k blocks. Identity Calibration and Parallel Mapping. Each block optimizes k(k − 1) phases using ZOO. All LN2/k2 blocks are optimized in parallel. The total step is 2k(k − 1)T1 for IC and 2LN2(k − 1)T2/k + 3 for PM. The total PTC call is around 2LN2T1 or 2LN2T2 for IC and PM, respectively.
Subspace Learning. We assume the feature map size is H ×W with a batch size of B. The detailed complexity analysis is given in Appendix G. The total step is approximately T3LNBHW/k.
According to our training cost profiler, IC and PM in total is 3-order-of-magnitude cheaper than the SL stage, since the batched parallel regression is deterministic and data-independent.

heading: 4 Results

heading: 4.1 Experiment Setup
text: Datasets. We evaluate L2ight on Vowel [10], MNIST [29], FashionMNIST [53], CIFAR-10, CIFAR100 [27], and TinyImagenet [7]. On CIFAR-10/100 and TinyImagenet, we adopt random crop, flip, color jittering for augmentation.
Models. All models are implemented with our open-source PyTorch-centric ONN library torchonn. We evaluate on a customized MLP (8-16-16-4) [18] on Vowel, CNN-S (CONV8K3S2-CONV6K3S2FC10) [18] on MNIST, a CNN-L ({CONV64K3}×3-Pool5-FC10) on FashionMNIST, and VGG8 [8]/ResNet-18 1 [24] on CIFAR-10/100. CNN-L/FashionMNIST is used for ablation studies. VGG-8/ResNet-18 on CIFAR-10/100 are used for accuracy and efficiency comparison. Training details can be found in Appendix E.
Efficiency Evaluation. We assume fully parallel 9× 9-blocking matrix multiplication in photonic tensor cores and sequential partial product accumulation in electronics. All experiments and performance measurements are based on software simulation with various noise modeling. Our simulator counts the total number of PTC calls as the normalized energy indicator and the longest accumulation path as the normalized latency/runtime indicator. Details of profiling can be found in Appendix G.

heading: 4.2 Main Results
text: Scalability Comparison with Prior ONN Learning Protocols. Figure 10 compares L2ight with
two SOTA ONN on-chip training protocols, FLOPS [21] and MixedTrn [18]. For ZO methods, i.e., FLOPS and MixedTrn, we count the energy and latency of forward PTC query in Appendix G. Prior protocols can only handle toy models (∼2,000 params) given their algorithmic inefficiency and instability, while our L2ight shows >1,000× higher scalability to handle large ONNs (∼10 M) on challenging
tasks with comparable accuracy to full-space pre-trained models. Though MixedTrn achieves com-
1https://github.com/kuangliu/pytorch-cifar
parable accuracy to L2ight on small benchmarks, we are still 1.7× faster and 6.9× more energy efficient.
The superiority of L2ight provides three important insights: (1) decoupling ZOO from stochasticity and partitioning a large-scale regression problem into a batch of sub-tasks can greatly mitigate the curse of dimensionality both in convergence and efficiency. (2) mapping before learning can fully leverage the pre-trained model to reduce the learning cost. Prior methods have to learn from a severely corrupted solution under variations, while L2ight recovers most accuracy via mapping, leaving a very light workload for subspace learning. (3) Restricted subspace learning provides adequate degree of freedom for training from scratch and task transfer. Also, its compatibility with first-order methods significantly boosts the trainability and breaks the scalability barrier for ONN training. We now validate the above insights by extensive experiments.
Training Efficiency Comparison with Prior Sparse Training Methods. In Figure 11, we show accuracy and efficiency comparison of 1) baseline L2ight-SL (BS), 2) L2ight-SL with spatial sampling (RAD), 3) L2ight-SL with weight and spatial sampling (SWAT-U), and 4) L2ight-SL with all three introduced sampling methods (feedback, column, and data sampling), and 5) our proposed full flow with IC, PM, and sparse SL (L2ight). To clarify, L2ight-SL performs subspace learning on-chip from scratch without using pre-trained weights, while L2ight includes the full flow, i.e., pre-training, mapping, and on-chip training. When we perform subspace learning from scratch, our proposed multi-level sampling strategies outperform previous RAD and SWAT-U by ∼3× in hardware cost with comparable accuracy. Though RAD can save the forward peak memory, it leaves the most expensive backward pass unoptimized, which does not fully exploit the sparsity in ONN training. SWAT-U tries to save forward cost by simultaneously sparsifying the forward and feedback weights with shared masks/patterns. However, in our experiment, the forward sparsification considerably degrades the performance, which dilates the efficiency benefits from it. Parallel mapping can fully leverage the pre-trained weights and help our full three-stage flow L2ight achieve the best accuracy with much faster convergence, leading to over 30× higher energy efficiency and fewer time steps.
Note that the energy efficiency and latency improvement is not just on the photonic part but a systematic performance boost. Our three-level sampling methods directly skip the pruned block, which means the corresponding cost of memory transaction, computation, control, and communication are removed together. Therefore, the sampling sparsity can be directly translated to the energy/latency improvement ratio regardless of whether the electrical part dominates the total cost.

heading: 4.3 Ablation Studies and Discussion

heading: 4.3.1 Multi-Level Sparsity in Efficient Training
text: Feedback Sparsity. To investigate the impact of feedback sampling strategies, we visualize the gradient approximation fidelity and accuracy curves in Figure 12(a). uniform sampling shows varying performance under different sparsity values due to large gradient variances. topk shows worse performance after sufficient steps due to its biased gradient estimation from overly greedy block selection. In contrast, our proposed load-balancing btopk strikes a balance between variance and bias via block-wise sampling and also leads to less runtime as it forces load balance among massively parallel PTCs. In Table 2, feedback sampling saves 50-60% time steps on the most costly error feedback∇xL, leading to 1.5-1.8× overall time step reduction with minimum accuracy drop. Feature Sparsity. Figure 12(b) compares the accuracy and weight gradient computation time steps on two feature sampling techniques. Though spatial sampling (ss) can save peak storage by dropping a subset of activations during the forward pass, it shows no gradient computation step reduction. Our hardware-friendly column sampling (cs) directly leads to energy and runtime reduction due to its
structured sparsity. In Table 2, when column sampling is further added, we observe ∼50% PTC energy saving on weight gradient computation∇ΣL at the cost of ∼1% accuracy drop. Data Sparsity. In the data level, we also demonstrate how SMD with sparsity αD impacts the training efficiency in Figure 12(c). With the best selected αW and αC , data sparsity directly reduces training time by skipping iterations [49]. The data sampling selects a uniform subset of the training set to represent the original data distribution, leading to less data processing with comparable generalization in the extracted features. Another explanation is that the variance increased by partial replacement serves as a regularization mechanism to improve the model performance [49]. For relatively easy tasks, aggressive sparsity (αD=0.8) is a sweet point, while for larger datasets shown in Table 2, a medium sparsity (0.5) can be a good setting to balance both the training cost and accuracy. With all three sampling methods integrated, our L2ight-SL shows competitive accuracy and ∼3× higher efficiency than RAD and SWAT-U. More advanced dataset sampling methods are left for future exploration.

heading: 4.3.2 Learnability of Restricted Subspace ONNs
text: Impacts of Calibration/Mapping Quality. Table 2 shows that with IC and PM, the full L2ight flow
achieves the highest accuracy with 32-35× efficiency boost over baselines. We further evaluate the impact of different mapping accuracy and the calibration quality on subspace learning in Figure 13. First, parallel mapping or pre-training is not a must. Our subspace learning supports first-order optimization on-chip from random initialization. Second, the optimality on subspace bases influences the final accuracy as it determines the upper bound of accuracy that can be recovered by subspace learning. With roughly optimized space bases, i.e., U ,V ∗, subspace learning can efficiently train basis coefficients, i.e., Σ, achieving 5-6% higher accuracy and 9.9× less energy and steps compared with random unitaries (train from scratch). Third,
subspace optimization shows low sensitivity on mapping quality and is able to compensate for the suboptimality in singular vectors within a reasonable range. Even with 60% mapped accuracy,
singular value optimization has enough capability to recover the accuracy to ∼90%. Fourth, our subspace learning is robust to gradient noises caused by non-ideal Ĩ (MSEU≈MSEV≈0.013), which shows that L2ight can tolerate reasonable suboptimality in the calibration and mapping stages.
In-situ Transferability in the Restricted Subspace. Another important question to answer is the
transferability of subspace learning. After mapping, we fix the inherited unitaries and adapt to different tasks by only training the singular values. Figure 14 shows that the inherited bases span a good design space with enough transferability. The insitu subspace transfer learning shows 1-2% higher final accuracy. Also, it uses 3∼5× fewer steps to obtain the same accuracy as training from scratch. Hence, our proposed L2ight finds a highly trainable design point while the learnability is still mostly maintained.

heading: 5 Conclusion
text: In this work, we propose the first scalable and efficient on-chip learning framework L2ight for emerging optical neural networks. Our proposed three-stage flow synergistically enables on-chip self-learning via automatic circuit state calibration, parallel model mapping, and efficient subspace learning. To further improve the learning efficiency, we explore multi-level sparsity, including balanced feedback sampling, information-preserving column feature sampling, and runtime-reduced data sampling. Extensive ablation studies and comparison experiments show 3-order-of-magnitude scalability improvement over prior on-chip training protocols and 30× efficiency boost compared with previous sparse training methods. We open-source a PyTorch-centric ONN library torchonn, based on which we release our on-chip learning framework L2ight at link. In the future, we will go beyond current software simulation and experimentally validate the effectiveness of L2ight on real photonic neural chips.
Acknowledgments The authors acknowledge the Multidisciplinary University Research Initiative (MURI) program through the Air Force Office of Scientific Research (AFOSR), contract No. FA 9550-17-1-0071, monitored by Dr. Gernot S. Pomrenke.

Review summaries:
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).    

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.