Online training and pruning of multi-wavelength photonic neural networks

2.1.1 Offline training

While PNNs can operate with picosecond latency for real-time applications [40], [41], [42], they rely on slower digital computers (i.e. CPU or GPU) for training, a process called offline training [40], [43]. As shown in Figure 1a, in this approach, NN parameters (e.g., neuron weights and biases) are optimized in software using backpropagation (BP) based gradient descent algorithm:

where w ^k denotes the weight matrix in the kth training epoch, α is the learning rate, L is the loss function. The parameters are then mapped onto the physical hardware (PNN chips) using highly accurate calibration models (also called look-up table (LUT)). For the MRR in the ith row, the jth column of an M × N MRR weight bank, it writes as

Figure 1:

Comparison of offline training with online training and pruning approach. (a) Conventional offline training method supporting only small-scale PNN chips. In this approach, NN parameters are calculated in software and mapped onto the PNN chips. The resonance variations of MRRs necessitate complicated look-up tables, and lead to higher power consumption for MRR control. (b) Online training and pruning method compatible for large-scale PNN chips. The training of PNNs occurs on the same chip used for inference, accounting for any chip-to-chip variations and environmental fluctuations. Our approach simultaneously optimizes the PNN accuracy and the total power consumption for tuning all the MRRs.

Here w _ij denotes the weight to be executed on the i, jth MRR, and I _ij is the required tuning current. The LUT, denoted as a mapping function f _ij , is determined by the relative position of the jth laser wavelength and the MRR resonant state. Ideally, the MRRs operating at the same laser wavelength are expected to exhibit identical resonant wavelengths, simplifying the LUT to f _j .

However, in practice, the MRRs in the same column may exhibit significantly different resonant wavelengths due to static and dynamic variations. As derived in Supplementary Note 1, an LUT for the i, jth MRR – accounting for FPVs, ambient temperature variations, input optical power, and the self-heating effect – is approximately given by

where σ _ij denotes the deviation of resonant wavelength due to FPVs, T is the time-varying ambient temperature, P_opt is the input optical power. Based on Eq. (5), it remains complicated to generate an LUT that accounts for all the non-idealities across large-scale MRR weight banks. Any inaccuracies in LUTs will directly introduce arbitrary deviations in the mapped weight parameters, leading to significant performance degradation across NN system benchmarks such as classification tasks.

In an MRR-based weight bank using thermo-optic tuning to counteract these non-idealities, the required electrical power for actively programming weights of the i, jth MRR can be written as [24]:

where R is the resistance of the metal heater. The power breaks down into a static weight locking power that locks the MRR resonance to the desired wavelength, and a configuration power to program the weight value [24]. Therefore, the total power for the weight configuration consumed by an M × N MRR weight bank is

For simplicity, we denote P^weight as P in the following sections except Section 2.3.

2.1.2 Online training

Online training, also known as “in situ” or “chip-in-the-loop” training, was originally proposed to alleviate the intensive use required of CPU/GPUs for training non Von-Neumann architecture based computing hardware [44], where the training occurs on the same hardware used for inference. It was also recognized that online training can mitigate hardware nonidealities, and promise for iterative weight updates in real-time [45]. The state-of-the-art online training algorithms for PNNs [6], [45], [46], [47], [48] can be classified into two major categories: gradient-based and gradient-free algorithms. Gradient-free algorithms, such as genetic algorithms [47], are straightforward to implement in practice but often face inherent challenges with convergence and scalability. In contrast, gradient-based algorithms are typically more efficient and, therefore, more widely adopted. The most commonly used gradient-based training algorithm for software-based NNs, backpropagation, analytically computes gradients by back-propagating errors using the chain rule [37]. While significant progress has been made in experimentally realizing BP on photonic hardware [48], [49], [50], the process typically requires global optical power monitoring and evaluation of nonlinear activation function gradients in software. This approach introduces additional system complexity and latency overhead. Other gradient-based algorithms, such as direct feedback alignment [45], replace the chain rule in back-propagation with a random weight matrix, but their validation has been limited to small, shallow NNs.

In our approach, the online training of our MRR-based PNNs is implemented by a perturbation-based gradient descent algorithm, which estimates gradients only based on forward inference running on PNN chips. Here, the NN parameters to be optimized are set to be MRR tuning currents instead of neuron weights:

The gradient of loss function with respect to MRR tuning currents (I ^k ) is approximately given by perturbation measurement:

where ΔI ^k is the perturbation rate of the MRR tuning currents in the kth training epoch, and Δ L is the measured change of loss function induced by the perturbation. In practice, ΔI ^k should be sufficiently small to preserve the validity of approximation in Eq. (9), while remaining large enough to produce measurable Δ L above experimental noise.

2.1.3 Pruning

In software-based NNs, pruning is a model compression technique that removes redundant weight parameters (by setting their values to zero) whilst maintaining accuracy. While the implementation of pruning has also been investigated in PNNs, prior approaches either lack robustness against nonidealities [51], [52] or require extensive offline training [26]. As shown in Figure 1b, our proposed approach accounts for the total MRR tuning power in the “chip-in-the-loop” process, by incorporating a power-aware “pruning” term into the conventional loss function:

where L ̃ is the modified loss function given by the sum of the conventional loss function L and the power-aware pruning term γP, γ is an empirically determined hyperparameter that defines pruning strength. According to Eqs. (6) and (7), L ̃ can also be expressed as

The selection of the hyperparameter γ is critical in this framework, as it explicitly parameterizes the trade-off between the NN accuracy and energy efficiency. Depending on the demand of the end users, minimizing this modified loss function can enable concurrent optimization of both the NN accuracy and the total power consumption. For example, in a power-constrained environment, the training problem becomes

with a relatively larger value of γ, where P_max denotes the maximum power available for weight configurations. Oppositely, in scenarios prioritizing high accuracies, it becomes

with a relatively smaller value of γ, where L max denotes the upper bound of the conventional loss function permitted under a given accuracy constraint.

2.2.1 Setup

Recently, integrated PNNs have been extensively investigated and experimentally validated across various on-chip scales, achieving low latencies of hundreds of picoseconds at diverse applications [6], [41], [53]. Notable demonstrations include a 3-layer PNN with 9 neurons for image classification [53], a 3-layer 6 × 6 PNN for vowel classification task [6], and a single-layer 4 × 2 PNN for fiber nonlinearity compensation [40]. In our proof-of-concept experiment, a 3 × 2 MRR weight bank is used to demonstrate our proposed online training and pruning approach and the associated energy savings.

Our experimental setup is illustrated in Figure 2a. First, three channels of input data (denoted as x₁, x₂, x₃) are generated from a high sampling rate signal generator (Keysight M8196A), and modulated onto laser 1, 2, and 3 (Pure-Photonics, PPCL500) respectively via Mach–Zehnder modulators (MZMs). The lights from these three lasers, each at a distinct wavelength, are combined using a wavelength-division multiplexer (MUX) and then split equally between two MRR weight banks. Each MRR weight bank consists of three MRRs, which are designed with slightly different radii (8 μm, 8.012 μm, 8.024 μm). The MRRs at the same column (i.e., MRR1 and MRR4, MRR2 and MRR5, MRR3 and MRR6) are expected to share the same corresponding resonance wavelengths, aligned with 200 GHz spaced ITU grids (1,546.92 nm, 1,548.51 nm, 1,550.12 nm, respectively). However, as shown in Figure 2b, the measured weight bank spectrum indicates that the resonance wavelengths of all six MRRs deviate from the designed values due to FPVs and temperature changes. To counteract FPVs, the wavelengths of three lasers are further tuned (1,547.24 nm, 1,548.60 nm, and 1,550.40 nm) to align with the deviated resonance wavelengths. Moreover, all MRRs are thermally tuned via embedded N-doped heaters to actively program and configure weights (denoted as w₁₁, w₁₂, w₁₃, w₂₁, w₂₂, w₂₃), allowing for individual weighting of input analog data in the three wavelength channels. The thermal tuning characteristics of six MRRs at 30 °C, are illustrated in Figure 2c.

Figure 2:

System setup of online training and pruning. (a) Schematic of our experimental setup. MZM, Mach–Zehnder modulator. MUX, wavelength multiplexer. PIC, photonic integrated circuit. MRR, microring resonator. BPD, balanced photodetector. ADC, analog-to-digital converter. CPU, central processing unit. σ(z) denotes the nonlinear activation function performed on software. The inset shows the micrograph of the two MRR weight banks and BPDs. (b) Normalized weight spectrum of two MRR weight banks at 30 °C. The MRRs at the same colors (i.e., MRR1 and MRR4, MRR2 and MRR5, MRR3 and MRR6) are designed to share the same resonance wavelengths, aligned with 200 GHz spaced ITU grids (1,546.92 nm, 1,548.51 nm, 1,550.12 nm). (c) Tuning characteristics of six MRRs at 30 °C. (d) Scatterplot of Iris flower dataset. (e) Schematic of a simple 3 × 2 neural network for Iris classification. (f) Online training and pruning procedure. At each iteration, the PIC performs matrix-vector multiplication without and with perturbations, while the CPU evaluates the gradients of power-aware loss function, and updates the MRR tuning currents for the next iteration.

Our proof-of-concept experiment is validated on a standard Iris dataset, where three species of Iris flowers (Setosa, Versicolour, Virginica) are classified using only three out of four input features (sepal length, petal length, and petal width), as shown in Figure 2d. This three-class classification problem is converted into a two-step binary classification (Figure 2e) mapped to the on-chip PNN, which utilizes eight trainable parameters. The parameters include the tuning currents of six MRRs – I₁₁, I₁₂, I₁₃, I₂₁, I₂₂, I₂₃ (corresponding to six weights) – and two biases, b₁ and b₂. The 150 samples are split into 120 for training and 30 for testing.

The analog optical signals from the drop and through port are then captured and differentiated by two balanced photodetectors (BPDs), which gives the electrical output of the weighted summations, Σ _i w_1ix _i and Σ _i w_2ix _i . The output is further read by an oscilloscope and demodulated in a CPU, which evaluates the power-aware loss function (Eq. (10)) and calculates gradients based on perturbation (Eq. (9)). Finally, the CPU updates the MRR tuning currents for the next training epoch (Eq. (8)) and commands the MRR driver, which is equipped on a customized printed circuit board (PCB) (Supplementary Note 2). As shown in Figure 2f, the training of our PNN occurs with the photonic chip in the loop, iteratively optimizing the loss function and reducing the total power consumption.

2.2.2 Demonstration on 3 × 2 PNN

PNN training is conducted under three conditions: conventional offline training, online training without pruning, and online training with pruning. Each training configuration is repeated 10 times to ensure consistency. The perturbation rate of MRR tuning currents (ΔI ^k ) is set to 0.05 mA. To evaluate the gradients with respect to the MRR tuning currents ∇ I k L , each training epoch batches the training data seven times through the PNN: six times perturbing each tuning current I i j k + Δ I i j k and one time to measure L ̃ with non-perturbed currents I i j k . The plots in Figure 3a show the average of both categorical cross-entropy loss ( L ) and power-aware loss ( L ̃ ) versus training epoch. It is observed that the PNN can quickly converge to the optimal weights within 25 epochs using perturbation-based online training.

Figure 3:

Online training results without, or with pruning method. (a) Experimental results of online training losses without, or with the pruning method. (b) Simulated result indicating the tradeoff between the prediction accuracy and power efficiency. (c–e) Confusion matrices for the 150 samples, obtained by the conventional offline training, online training without, or with the pruning method, respectively.

We also simulate the trade-off between the NN accuracy and the associated energy savings (Figure 3b). The simulation indicates that the trade-off is optimized at γ_opt = 0.01 mW⁻¹, where significant reduction of overall tuning power (exceeding 70 %) can be achieved before the accuracy drops off. As predicted in Section 2.1.3, in scenarios prioritizing high accuracies where γ < γ_opt, the loss function L ̃ is dominated by the conventional loss function L , such that the PNN maintains a high accuracy of >93 %. Contrarily, in scenarios constraining overall power consumption where γ > γ_opt, the power-aware loss function L ̃ is dominated by the pruning term γP, leading to a larger amount of energy savings but may degrade the PNN accuracy. As shown by the confusion matrix in Figure 3c, the offline training experiment only produces 82 % accuracy on the 150 samples. In the accuracy-prioritized online training experiment (without pruning, γ set to 0) (Figure 3d), the overall classification of the Iris flower task is improved to 96 %, with a total MRR tuning power of 9.54 mW. In the online training experiment constraining power consumption (with pruning, γ set to 0.0075 mW⁻¹) (Figure 3e), we observe a 44.7 % reduction (from 9.54 mW to 5.28 mW) of total tuning power while the PNN maintains a classification accuracy of 95.33 %.

Furthermore, to demonstrate the adaptability of our online training method to temperature drifts, additional experiments are conducted at multiple different temperatures ranging from 26 °C to 40 °C. Before online training, the MRR tuning currents are randomly initialized within the tuning range of 0–1.5 mA (only weights w₁₁ and w₂₁ are plotted, Figure 4a), producing untrained classification outputs with only 33.33 % accuracy. After online training for 100 epochs (Figure 4b), the PNN can be trained to the optimal weights, regardless of the changes of MRR tuning characteristics due to temperature drifts. As a result, high classification accuracies is consistently achieved at both 30 °C (99.33 %, upper row of Figure 4a and b) and 34 °C (96.67 %, lower row of Figure 4a and b). As shown in Figure 4c, the classification accuracies maintain above 90 % even if the external temperature changes from 26 °C to 40 °C. We also evaluate the resilience of our PNN training to intentional thermal disturbances. In Figure 4d, it is shown that the PNN can quickly recover its classification performance within about 25 epochs even if the external temperature is disturbed from 27 °C to 32 °C (at the 100th epoch), and from 32 °C to 30 °C (at the 200th epoch). These proof-of-concept experimental results further confirm that our online training and pruning method can handle any PNN hardware non-idealities, including chip-to-chip FPVs, temperature disturbances, and even nonlinearities induced by input optical power.

Figure 4:

Adaptive PNN training at different temperatures. (a, b) Illustrate two of the six MRR tuning currents (weights I₁₁ and I₂₁) before and after online training, respectively, at 30 °C (the upper plots) and 34 °C (the lower plots). (c) Experimental obtained Iris classification accuracies at different temperatures. (d) Online training experiment performed with thermal disturbances. The temperature is changed from 27 °C to 32 °C at the 100th epoch, and from 32 °C to 30 °C at the 200th epoch.