Physical neural networks using sharpness-aware training

Domain: physics
Published: 2026-01-19
Editorial outcome: major_revision

Abstract
Recent advances in AI are pushing the limits of traditional hardware, making physical neural networks (PNNs) a promising alternative. However, training PNNs remains challenging: in silico training suffers from model-reality mismatch, while in situ training produces device-specific models that do not transfer across fabrication variations. Both approaches are further compromised by post-deployment perturbations, such as thermal drift or misalignment, which invalidate trained models and require retraining. We address these challenges through sharpness-aware training (SAT), inspired by sharpness-aware minimization, which links loss landscape geometry to generalization. We establish a connection between loss landscape sharpness and robustness in physical systems and leverage it to improve PNN training. SAT is compatible with both in silico and in situ settings: it mitigates model-reality gaps, enables cross-device transfer, and provides strong resilience to post-deployment perturbations without retraining. We demonstrate SAT across three PNN platforms and multiple tasks, including classification, compression, reconstruction, and generation, showing its broad applicability.

Introduction
The tremendous success of modern artificial intelligence (AI) is driven by a synergic operation of hardware and algorithms1. Powerful hardware, particularly graphics processing units, has revolutionized neural network (NN) computing by drastically reducing the time and resources needed to process large datasets. Central to this progress is backpropagation (BP), the foundational algorithm for training NNs, which efficiently dictates how NNs learn and evolve by minimizing errors2.

Recent AI models have begun to strain the capabilities of traditional digital hardware. The cost of training and inferring AI models doubles every 2 months, outpacing Moore’s Law. Consequently, AI systems are increasingly facing challenges caused by hardware efficiency. To address hardware limitations, a different computing paradigm, neuromorphic computing, has gained traction. Many neuromorphic computers are physical systems leveraging photonics3–5, analog electronics6, and other wave-based physics7,8, as shown in Fig. 1a, b. The physical models of these systems exhibit mathematical isomorphism with NNs, enabling their computing capabilities. Their physical configurations can alleviate the data movement bottleneck found in traditional digital hardware, substantially improving both computing speed and energy efficiency. Large-scale9–13, high-speed14–17, and energy-efficient18–22 physical neural networks (PNNs) have been widely reported23–25.

However, the development of suitable training methods that are synergistic with these computing hardware has lagged behind. The training process remains cumbersome and unreliable for PNNs. Traditional gradient-based algorithms, such as BP, cannot be directly applied to PNNs due to the difficulty of obtaining gradients from physical systems. Moreover, even with appropriate training methods, PNNs struggle to maintain computing accuracy during deployment, as physical systems are inherently highly susceptible to disturbances compared to digital hardware, leading to errors3,26,27.

In reviewing the existing training methods, they can be divided into two main categories: in silico training and in situ training (Fig. 1c). To train a PNN in silico, the process begins by extracting a digital model of the physical system. The training is then conducted on a digital computer, typically using BP. Once the model is trained, the parameters are applied to the PNN. However, in silico training requires highly accurate digital modeling of the physical system, which is challenging due to unavoidable fabrication-induced errors and intrinsic noise present in the physical systems. As a result of such a model-reality gap, the inference accuracy is often lower than the ideal value, as illustrated in Fig. 1c-i. In contrast, in situ training alleviates the need for precise digital modeling by looping the physical system directly into the training process. However, in situ training faces the problem of obtaining the gradients of the physical system. Current methods for obtaining gradients include: (1) Approximating the gradients using a digital twin model of the physical system obtained through data-driven methods4,28–31; (2) Estimating the gradients through finite-difference methods by applying perturbations and measuring difference32–34; (3) Directly measuring the gradients physically35–40. Each of these approaches introduces overhead, either through excessive measurement requirements or the computational cost of creating a digital twin model. To mitigate this overhead, various BP-free or gradient-free approaches have been proposed3,41,42. However, their training efficiency and effectiveness in complex applications have yet to be fully demonstrated. Moreover, a drawback of all in situ learning approaches is that, because the specific physical system is looped into the training process, the trained results are only applicable to that particular system and cannot be transferred to other systems, even those with the same design. Our simulation results show that the in situ-trained model is even more vulnerable to imperfections than the in silico-trained model when transferring the model to other systems. A critical issue across both in silico and in situ training methods is that perturbations after deployment, such as thermal drift or alignment errors, will make the trained parameters ineffective, leading to accuracy loss and requiring retraining3.

This work proposes and demonstrates a training technique, sharpness-aware training (SAT), that overcomes the challenges of both in silico and in situ training. SAT can be implemented in both in silico and in situ training settings. It closes the model-reality gap in in silico training by enabling accurate training of PNNs using efficient BP algorithms, without the need for an exact digital model or looping the actual hardware during the process. The systems trained in silico by SAT even outperform those trained in situ, even in the presence of modeling and fabrication errors. Additionally, SAT addresses a key limitation of current in situ training methods. Parameters trained on one device can be reliably transferred to other devices for inference without accuracy loss, even if they are different due to fabrication variances. Equally important, the technique exhibits strong resilience to post-deployment perturbations, a common problem in both in silico and in situ training. This robustness allows PNNs to continue operating accurately under perturbations without the need for retraining.

The key to these benefits is that SAT leverages the geometry of the loss landscape in a physical system: it minimizes not only the loss value but also the loss sharpness, which indicates the robust region of the physical system, as shown in Fig. 1d. By identifying parameters in regions of uniformly low loss, SAT ensures that the loss remains minimal despite challenges common in analog computing systems, such as modeling inaccuracies, fabrication variances, and external perturbations. SAT is inspired by the machine learning technique sharpness-aware minimization (SAM)43, which enhances model generalization by connecting loss sharpness to model generalization. In contrast, SAT introduces a different connection between loss sharpness and the robustness of physical systems. Furthermore, SAT develops methodologies for automatically locating flat minima in physical systems and enables direct training in physical systems.

SAT is universally applicable to various PNNs, regardless of whether their models are explicitly known or unknown. We demonstrate the generality of SAT through three typical photonic NNs: integrated microring resonator (MRR) weight banks, Mach-Zehnder interferometer (MZI) meshes, and diffractive-optics-based NNs. These demonstrations highlight the effectiveness and accuracy of SAT during both training and deployment stages, even under error sources inherent in analog computing. This includes imprecise modeling prior training, thermal and electrical disturbances that heavily affect sensitive devices such as resonators, fabrication variances in integrated optics, and deployment errors commonly seen in free-space optical systems. We further expand the validation across different tasks, including image classification, compression, reconstruction, and generation. Overall, SAT offers a practical, effective, and computationally efficient solution for training and deploying PNNs in real-world applications.

Results
Sharpness-aware training: seeking loss minimum and loss sharpness minimum in a physical system
A deep neural network (DNN) is a multi-layer system, where each layer consists of synapses and nonlinear nodes, as depicted in Fig. 1a. The mathematical model of the lth layer of DNN is given by,

Xl+1=Al(WlXl+bl) (1)

Here, Xl is the input of the lth layer, Wl and bl denote the trainable weights and bias, and Al(⋅) represent the lth layer nonlinear activation layer. A PNN emulates this process through a physical system controlled by a set of tunable physical parameters Θ, such as electrical currents. These parameters determine the system’s transfer function, which can be expressed by y=F(x; Θ) with x and y representing the system input and output and Θ representing the tunable physical parameters. Tuning these control parameters affects the system’s transfer function, making them equivalent to changing weights in a DNN. Therefore, training a PNN is to find the optimal control parameters that minimize the loss function for a given task.

To train such a system accurately, obtaining the exact mathematical expression for y=F(x; Θ) is essential for conventional methods. However, this is nearly impossible for real-world physical systems, which are governed by highly complex physics and are prone to parameter drift, device crosstalk, and environmental disturbances. Discrepancies between the ideal model and the actual system lead to training errors. In situ training reduces the reliance on an exact mathematical expression. It obtains the true loss function by experimental measurement. With the true loss function, even if there are slight discrepancies between the hardware and the ideal model, the loss can still be minimized through iterative optimization4. However, in situ training is inefficient because it is specific to the particular system being trained and cannot be easily transferred to other devices, even those fabricated with the same design. This limitation is unfavorable compared to the typical functioning of modern AI hardware, where training is conducted on a single high-end computer or a cloud-based computer cluster, and the trained parameters are then deployed across many edge devices for inference26. Moreover, a critical issue across both in silico and in situ training methods is that perturbations after deployment, such as thermal drift or installation errors, can render the trained parameters invalid.

We address all the challenges of both in silico and in situ training with our SAT method, which leverages the geometry of the loss landscape in a physical system to modify the training objective accordingly. The loss landscape, illustrated in Fig. 1e, represents how the loss changes as parameters are adjusted. Traditional training methods like BP only focus on minimizing the loss function, which typically ends at the “High sharpness minimum” as depicted in Fig. 1e. The “High sharpness minimum” is characterized by sharp curvatures, where small changes in parameters lead to rapid increases in loss. In machine learning, such suboptimal minima lead to poor model generalization. In this work, we associate these sharp minima in the loss landscape with the poor robustness of physical systems. At such sharp minima, slight parameter changes due to modeling discrepancies, fabrication variances, or environmental disturbances can cause rapid increases in loss and errors, affecting both the training and deployment stages when using in silico and in situ training methods.

Our SAT method addresses these problems by automatically searching for flat minima in the loss landscape, as the “Low sharpness minimum” illustrated in Fig. 1e. At these flat minima, the loss remains minimal even with parameter changes, providing high resilience and robustness against errors and imperfections in physical systems. In this work, we derive a universal framework for finding flat minima across different PNN systems.

To find the “Low sharpness minimum,” the goal is not only to minimize the loss but also to ensure parameters lie in the neighborhoods with uniformly low loss value. The optimization process of SAT can be mathematically expressed as finding Θ minimizing the loss function L, as shown in Eq. (2):

L(Θ) = L(y, ytarget; Θ) + α || ∂L(y, ytarget; Θ) / ∂Θ ||^2 (2)

The first term, L(y, ytarget; Θ), in the loss function seeks to find the loss minima, while the second term, || ∂L(y, ytarget; Θ) / ∂Θ ||^2 reflects the sensitivity of the objective function to changes in the parameters. This term aims to identify parameters that lie in the neighborhoods with uniformly low loss values. As shown in Fig. 1d, this additional term can be expressed as ∂L(y, ytarget; Θ) / ∂Θ = ∂L(y, ytarget; Θ) / ∂W * ∂W / ∂Θ, where ∂W can optimize the stability of individual components and ∂L(y, ytarget; Θ) / ∂W optimizes the stability of the whole system. Here, the component stability refers to how the weight of a single component remains stable under small control parameter perturbations, while the system stability reflects how the final loss remains stable under small weight perturbations. Both component and system stability are improved during optimization. The regularization term α determines the penalty of sharpness for loss during the actual training process, that is, the trade-off between sharpness and loss.

However, directly calculating the second term requires computing second-order derivatives, the Hessian matrix, which leads to the computational complexity increasing from O(n) to O(n^2), where n denotes the matrix size. This additional cost makes Eq. (2) challenging to implement in practice. To address this, Foret et al. propose approximating the gradient for minimizing the loss function using a first-order gradient of L at Θ + ΔΘ43, the gradient is calculated through Eq. (3),

ΔΘ = α * (∂L / ∂Θ) / || ∂L / ∂Θ ||^2
∂L_SAT(Θ) = ∂L(Θ + ΔΘ) / ∂Θ (3)

Here, ΔΘ represents the parameter perturbation that induces the largest change in the loss within the neighborhood of Θ. The previously defined hyperparameter α also indicates how large the perturbation ΔΘ should be given to Θ. And ∂L/∂Θ depicts the loss gradient with respect to parameter Θ. The specific value of α is chosen based on whether it can maximize the model’s robustness without decreasing the accuracy. We give a sensitivity analysis over the hyperparameter α, and the optimum is around 0.1 to 1. The detailed sensitivity analysis is depicted in the Supplementary Note 1.2. According to the result, we choose the value of α to be 0.1.

Minimizing the loss function can thus be achieved through a two-step process of automatic differentiation, as illustrated in Fig. 1e. The first step involves identifying the point Θ + ΔΘ that results in the maximum loss within the neighborhood (denoted as “Midpoint” in Fig. 1e) using Eq. (3). Once these parameters are identified, the gradient is recalculated at “Midpoint.” This second gradient calculation with Eq. (3) reflects the direction in which the maximum loss in the neighborhood can be reduced, i.e., directing to the “flat minimum,” as shown at the point labeled “Low sharpness minimum” in Fig. 1e. Consequently, this two-step BP not only finds the minimum loss but also ensures that the loss remains low in the vicinity of the minimum, achieving a “flat minimum.”

The loss minimization process highlights the difference from our prior work, optical pruning27. Optical pruning focuses solely on improving the stability of individual devices inside the PNNs. In contrast, SAT not only improves the stability of individual components but also enhances the system’s stability by ensuring that the loss remains low in the vicinity of the minimum. Furthermore, optical pruning is limited to relatively simple systems, whereas SAT is broadly applicable to different physical systems because it can automatically identify robust minima without requiring explicit knowledge of the underlying physical dynamics. SAT can also outperform noise-aware training (NAT)44–47, a widely used method to improve physical system robustness by injecting synthetic noise during training. The reason is that NAT typically injects simple additive Gaussian noise into the model, but real-world imperfections, such as modeling inaccuracies, pixel-level misalignments, fabrication errors, and temperature-induced shifts, often deviate from this assumption. Consequently, NAT struggles to fully capture the complex, system-specific non-Gaussian error characteristics, limiting its effectiveness in enhancing robustness. In contrast, SAT makes no assumptions about the nature of system imperfections. Instead, by connecting robustness to the intrinsic geometry of the loss landscape, SAT formulates a general optimization objective and naturally accounts for diverse and non-Gaussian imperfections without requiring explicit modeling. This allows SAT to optimize PNNs with greater accuracy and resilience, even when the specific forms of system noise are unknown and non-Gaussian.

As shown in Fig. 1f, SAT is applicable to both mainstream in silico and in situ training workflows. In in silico training, the training process is conducted on the mathematical model y = F(x; Θ) of a PNN to directly determine the controlling parameters. Even when the mathematical model deviates from the actual system, the trained parameters remain valid and maintain high accuracy when applied to the physical system. In in situ training, the loss function is calculated directly from measurements, and the two-step differentiations are obtained using the gradient estimation methods mentioned earlier. We further demonstrate the generality of SAT through three typical photonic NNs: integrated microring resonator (MRR) weight banks, Mach-Zehnder interferometer (MZI) meshes, and diffractive-optics-based NNs. These demonstrations highlight the effectiveness and accuracy of SAT during both training and deployment stages, even under error sources inherent in analog computing. This includes imprecise modeling prior training, thermal and electrical disturbances that heavily affect sensitive devices such as resonators, fabrication variances in integrated optics, and deployment errors commonly seen in free-space optical systems. We further expand the validation across different tasks, including image classification, compression, reconstruction, and generation. Overall, SAT offers a practical, effective, and computationally efficient solution for training and deploying PNNs in real-world applications.

In the following section, we demonstrate SAT on three distinct PNNs. These demonstrations not only highlight the broad applicability of SAT, but also illustrate its advantages over existing training methods from three perspectives: (1) When SAT is implemented in in silico training setting, it closes the model-reality gap, enabling accurate BP training under imprecise models (section “In silico SAT: enabling accurate training with an inaccurate digital model”); (2) When SAT is implemented in in situ training setting, SAT addresses the challenge of limited transferability caused by fabrication variances. Parameters trained on one device using SAT can be reliably deployed to other devices without accuracy degradation, even when hardware discrepancies exist (section “In situ SAT: facilitating transferable in situ training”); (3) In both in silico and in situ training settings, PNNs trained by SAT can continuously operate accurately after training under perturbations without retraining (sections “In silico SAT: enabling accurate training with an inaccurate digital model”, “In situ SAT: facilitating transferable in situ training,” and “Extending SAT to PNNs without explicitly known models”); (4) SAT is broadly applicable to physical systems regardless of whether an accurate physical model is explicitly available (sections “In silico SAT: enabling accurate training with an inaccurate digital model”, “In situ SAT: facilitating transferable in situ training,” and “Extending SAT to PNNs without explicitly known models”). The advantages of SAT over existing in silico and in situ training methods, including those designed to improve robustness, such as NAT and optical pruning, are further demonstrated in the following sections, and summarized in Table 1.

In silico SAT: enabling accurate training with an inaccurate digital model
This section demonstrates our approach’s effectiveness in the in silico training scheme. Using MRR-based PNNs as our first demonstration platform, we demonstrate two advantages of SAT compared to traditional training techniques. Firstly, SAT closes the model-reality gap: it does not require highly accurate digital modeling of the physical system or considering side effects or system imperfections, yet it still outperforms the method requiring them. Secondly, PNN trained by SAT is more robust against dynamic noise and post-deployment perturbations, particularly temperature drift.

In the MRR-based PNNs, weight synapses are realized using MRR weight banks, as illustrated in Fig. 2a. The input matrix is represented by an array of analog signals, each modulated onto a laser at a distinct wavelength. The multiplication operation is performed by an array of MRR weight banks, shown in Fig. 2b. Each MRR in the weight bank has a slightly offset radius, resulting in a unique resonance wavelength that aligns with the corresponding input laser’s wavelength. By tuning the resonance wavelength of the MRR, the distribution of light between the Drop and Through ports can be precisely controlled, allowing for tunable weights, as shown in Fig. 2c. Thermal tuning is the commonly employed mechanism for controlling MRRs due to its ease of implementation on silicon photonics chips. Once the input signals are weighted, they are detected by a balanced photodetector (PD)17,27. In this experiment, we use a single-end PD to detect the light from the Through port. MRR-based PNNs offer several advantages, including power-efficient tuning, straightforward weight assignment, and high computational density due to the compact device footprint24. However, MRRs are highly sensitive to thermal changes, such as thermal drift, thermal crosstalk, and ambient disturbances. In our previous work27, we observed that a resonance shift of just 12pm (corresponding to a temperature change of 0.2°C) caused the PNN’s accuracy to drop from an ideal 99.0% to 67.0% when classifying the Modified National Institute of Standards and Technology (MNIST) dataset48.

Here, we conduct our experiment using a 4×4 MRR weight bank, as shown in Fig. 2b (refer to the Supplementary Notes 4.2 and 4.3 for detailed device design and experimental process). To demonstrate the effectiveness of SAT to address challenging problems, we have several experimental settings. (1) When constructing the mathematical model y=F(x; Θ), we treat all MRRs as identical. In this model, Θ indicates an array of currents used to control each MRR. We experimentally measure the current-weight relationship on only one specific MRR and apply this relationship uniformly across all MRRs in the model. However, as shown in Fig. 2d, fabrication variations lead to different current-weight characteristics for each MRR, resulting in modeling errors. (2) We intentionally ignore all side effects that could introduce modeling errors in our training model. These side effects include thermal crosstalk between devices, noise, and imperfections in peripheral circuits such as wavelength-dependent gain in erbium-doped fiber amplifier and non-ideal modulations. Ignoring these factors further causes the mathematical model to deviate from the actual system. (3) After training, the deployed system experiences thermal drift, further impacting the system’s performance.

We first compare the training performances of our SAT method with the traditional training method using standard BP on the MNIST database. As shown in Fig. 2e, both methods converge after 20 training epochs with a training accuracy of 99.0%. (see Supplementary Note 4.1 for training details). The current distributions obtained using the two methods are shown in Fig. 2f. The effectiveness of SAT is highlighted by the loss landscape illustrated in Fig. 2g. The loss landscape clearly shows that the system trained with SAT exhibits a flatter profile, indicating enhanced stability and robustness against variations in control parameters compared to traditional training methods49. To emphasize the largest curvatures, the x and y axes represent the two principal directions of the Hessian matrix of the loss function with respect to the control currents, while the z-axis corresponds to the loss function. This robustness is also reflected in the maximum eigenvalue of the Hessian matrix, λmax. A smaller eigenvalue suggests a flatter loss landscape. With SAT, λmax reduces dramatically from 746.8 with standard BP to just 1.2, highlighting the effectiveness of SAT in identifying stable regions within the physical system.

After training, we deploy the trained currents to the MRR-based PNN to evaluate the inference accuracy. As mentioned in the experimental setup, there is a deviation between the model used for training and the actual system due to static errors and dynamic noise. Consequently, when using standard BP, the accuracy for MNIST classification drops to 80.0% from the ideal accuracy of 98.0%. In comparison, the accuracy achieved with SAT remains at 97.0%, as shown in Fig. 2h. This result is obtained at a temperature of 22°C, the same temperature used to characterize the MRR current-weight relationship. To further simulate the real-world conditions under temperature drift, we vary the chip temperature using a thermoelectric controller (TEC) from 21 to 23°C and test the inference accuracies, with results shown in Fig. 2h. The system trained with SAT demonstrates remarkable robustness to temperature changes, maintaining its computational accuracy, while the accuracy of the system trained with standard BP degrades from 80.0 to 17.0% as the temperature varies from 21 to 23°C.

Next, we directly remove the TEC and measure the accuracy under the room temperature of 20°C. Note that the training model is still characterized at 22°C with TEC. In this circumstance, despite the temperature change being as large as 2°C, SAT still maintains high accuracy at 91.0%, in contrast to 52.0% for standard BP.

We further compare SAT with NAT. NAT achieves an accuracy of 73.0%, indicating that it is less resilient to such a temperature variation compared to SAT, which achieves 91.0%. This difference is also evident in the sharpness metric λmax: while NAT reduces λmax compared to standard BP, it still reaches 414.5, which remains substantially higher than 1.2 achieved by SAT. The reason is that NAT typically only accounts for Gaussian errors but struggles to fully capture the complex, system-specific non-Gaussian error characteristics, limiting its effectiveness in enhancing robustness. In contrast, SAT naturally accounts for diverse and non-Gaussian, allowing SAT to optimize PNNs with greater accuracy and resilience, even when the specific forms of system noise are unknown and non-Gaussian. We also compare the SAT method with our previous work on optical pruning27. The core idea in that work is to train the parameters of the MRR-based PNN to enhance the noise robustness of each MRR. Optical pruning focuses solely on improving the stability of individual devices. Under large model deviations and temperature fluctuations, optical pruning achieves an accuracy of only 73.0%. In contrast, SAT not only improves the stability of individual components but also enhances the overall system’s stability. Consequently, SAT achieves the highest accuracy at 91.0%.

Additionally, we extend the application of our approach to much more challenging tasks, including image classification, compression, reconstruction, and generation, as shown in Fig. 3a–m. Our results show that PNNs trained with SAT maintain high accuracy in all these tasks under temperature variations within ±0.25°C, a thermal fluctuation level in current co-packaged optics50. In contrast, PNNs trained using standard methods suffer performance degradation and fail under the same conditions. In the image classification task, we apply the ResNet18 model in the MRR-based PNN to classify the CIFAR-10 dataset. The model’s robustness is evaluated by varying the operating temperature. As shown in Fig. 3b–d, SAT maintains stable accuracy across ±0.25°C temperature variations, whereas standard BP fails. Specifically, under a −0.25°C shift, BP’s accuracy drops to 10.6%, while SAT maintains 83.4%, nearly matching its performance under nominal conditions. Image compression and reconstruction tasks use an autoencoder framework. We first compress CIFAR-10 images into a latent space with an MRR-based encoder, followed by digital decoding. As shown in Fig. 3f–h, for the SAT-trained model, even under ±0.25°C shifts, the reconstruction quality, assessed using mean squared error (MSE), remains at around 0.023, comparable to its nominal performance (0.013). In contrast, the BP-trained model fails. Under a −0.25°C shift, the reconstruction collapses with an MSE of 0.529. Next, we implement a generative adversarial network where the MRR-based PNN serves as the generator. The quality of the generated images is evaluated using Fréchet Inception Distance (FID)46. As shown in Fig. 3j, SAT ensures stable FID values under temperature shifts, whereas BP suffers from severe degradation. Figure 3k, l further highlights the robustness of SAT. The images generated by SAT remain clear under a −0.25°C shift, while BP-generated images degrade. Additionally, the t-SNE visualization (Fig. 3m) further confirms that SAT-generated images maintain distributional consistency with the original dataset, unlike those generated by BP. Details are clarified in Supplementary Note 4.4.

In situ SAT: facilitating transferable in situ training
This section addresses the transferable learning challenge of in situ training: the trained results are specific to a particular system and cannot be easily transferred to other systems, even those fabricated with the same design, due to fabrication variances. This lack of transferability is a drawback for modern AI hardware, where training is typically performed on a single high-end computer or cloud-based cluster, and the trained parameters are then deployed across multiple edge devices for inference. Here, we demonstrate that SAT can be combined with in situ training approaches to overcome this major drawback and enable transferable in situ training.

We demonstrate SAT on the MZI meshes based NN through numerical simulations using the library provided by Ziyang et al.28. In MZI mesh-based NNs, the basic building block is a 2×2 MZI, which consists of two 50-50 beam splitters and two phase shifters. This block performs an arbitrary 2×2 unitary transformation by configuring two phase shifters. By cascading multiple MZI blocks, the MZI mesh can achieve arbitrary unitary transformations51. Programming such a system remains challenging, as the model is difficult to characterize due to unknown initial phases caused by phase errors in the phase shifters. Consequently, many in situ training approaches have been proposed to program such systems. However, wafer-level manufacturing errors make the trained result specific to only the trained system but not transferable to other systems. For example, as pointed out by ref. 52, variation of the beam splitter as small as 2%, which is a typical wafer-level variance, can degrade the NN accuracy by nearly 50%. We demonstrate in situ SAT on a simulated digital-optical hybrid NN, following the setup in ref. 28. The digital NN consists of two convolutional layers followed by two fully connected layers, with the functions of feature extraction and image down-sampling. The optical NN is a 64×64 MZI mesh followed by a square detection layer, serving as the output layer. The first 10 output ports are used to generate classification output on the MNIST dataset.

We evaluate six approaches under fabrication variances: (1) in silico training with standard BP on a digital model without considering fabrication variances; (2) in silico NAT, which injects synthetic noise during training; (3) physical-aware training (PAT), an in situ approach that leverages the actual PNN system for forward propagation and a proxy model for gradient estimation; (4) dual-adaptive training (DAT), an in situ approach modified from PAT, capable of deriving a more accurate proxy model for improved gradient estimation; (5) in silico SAT, trained in silico on a digital model and deliberately ignoring fabrication variances; (6) an in situ SAT, which adopts the PAT framework but adding searching for the sharpness minimum. Detailed setups for these approaches are clarified in Supplementary Note 5. The purpose is to verify: (1) In silico SAT can address model-reality gap and even outperforms in situ training methods; (2) SAT can integrate with in situ training, particularly PAT, and improve their performance in light of inaccurate modeling in PAT; (3) In situ SAT can enable transferable learning to other devices, overcoming the major drawback of current in situ training methods.

To simulate the fabrication errors, we generate phase error σps and splitting error σbs sampling from a Gaussian distribution with a mean of 0 and variance of 0.15, and add these errors to the MZI mesh model. As shown in Fig. 4c, when using in silico training with standard BP, these fabrication errors cause an accuracy reduction from 97.4 to 68.6%, corresponding to an error rate of 31.4%. This is because the fabrication variances create a substantial mismatch between the ideal model used for training and the actual system. The training time is 356.1s. Next, we apply the NAT training method. During training with NAT, the noise we add to the phase parameter follows a Gaussian distribution with a mean value of 0 and a standard deviation of 0.01 rad. The values are chosen to maximize robustness without decreasing the accuracy. Additionally, we add Gaussian noise to intentionally increase the trained model’s robustness, and the trained accuracy is only 79.8%, corresponding to an error rate of 20.2% caused by noise injection. The training time is 338.2s.

To address the model mismatch, we turn to in situ training methods, including PAT and DAT, which loop the device into the training process. In PAT, the forward pass uses the system with fabrication errors, while the backward model is an ideal MZI mesh without error. Despite using in situ training, under large fabrication errors, PAT achieves an accuracy of only 70.2% with a training time of 393.8s. The low accuracy is due to the approximated gradients from the ideal model deviating from the true gradients, which hinders the loss minimization. To obtain a more accurate proxy model for BP, DAT combines the ideal model with an additional digital model using the data-driven method. DAT brings the approximated gradients closer to the real gradients, therefore improving the accuracy to 94.6%. But this improvement is at the cost of adding six additional systematic error prediction networks (SEPNs), with the training time increasing to 1272.7s. The results are consistent with ref. 28.

We also find the model trained by DAT is highly sensitive to noise and perturbations. We use the maximum eigenvalue of the Hessian matrix λmax to evaluate the trained model, as a large λmax indicates high sensitivity to parameter changes. The λmax for DAT is 252.4, nearly three times that of the model trained in silico with standard BP, which has a λmax of 93.5. This result implies that while DAT provides accurate gradients and improves classification performance, this comes at the cost of reduced robustness. As a result, the trained system will be highly vulnerable to perturbations and noise after deployment, and the trained parameters cannot be transferred to other systems, even those fabricated with the same design.

Next, we evaluate in silico SAT and in situ SAT (which is adapted from the PAT framework but with additional training to search for the sharpness minimum). We first implement SAT trained in silico on an ideal model. The training time is 513.8s, which is longer than that of in silico training with standard BP due to the need for two steps of BP. However, the trained system is much more robust to the fabrication errors, achieving an accuracy of 92.3% even under large fabrication errors. Remarkably, this performance not only outperforms in silico training methods but also outperforms in situ training methods, including PAT and DAT. Moreover, the Hessian max eigenvalue λmax is only 0.3, which is lower than other methods under comparison. The result shows the trained model has high robustness compared to others.

To further improve accuracy, we perform in situ SAT by combining SAT with PAT, resulting in an accuracy of 97.1%. This outcome suggests that incorporating SAT effectively addresses the modeling inaccuracy in PAT: in situ SAT demonstrates the ability to achieve high training performance even with inaccurate approximating gradients. Moreover, in situ SAT also exhibits high robustness. The Hessian max eigenvalue λmax is only 0.6, indicating that SAT improves the robustness of the trained system.

This robustness allows the parameters obtained through in situ training to be transferred to other devices without a loss in accuracy. To evaluate this, we apply the trained model to other devices with different randomly generated phase errors and splitting ratio errors. We systematically increase the errors in both phases and the splitting ratio and assess classification accuracy, with results shown in Fig. 4d. SAT is capable of maintaining high accuracy at 95.2%, indicating that SAT can facilitate transferable in situ training. In contrast, the system trained by DAT decreases to 58.6% when the phase and splitting ratio errors reach 0.15.

In summary, SAT enhances both in silico and in situ training schemes. When trained in silico, SAT achieves high accuracy under large fabrication errors, surpassing state-of-the-art in situ methods. Moreover, integrating SAT with in situ training (such as PAT) further improves accuracy. Importantly, SAT enables the trained parameters to be transferable to other devices, addressing a key limitation of current in situ training methods.

Extending SAT to PNNs without explicitly known models
We further demonstrate our method’s effectiveness in diffractive optical NNs53. The training is performed in silico. Unlike the MRR-based PNN system or MZI-based PNN system, where the relationship between the weights and the control parameters is known (though not entirely precise), in diffractive optical NNs, the relationship between the misalignment factors, such as rotation angle and the controlling parameters, is not explicitly known.

In diffractive optical NNs, as shown in Fig. 5a, b, an NN layer is decomposed into multiple matrix multiplications and realized by an imaging system. In the experiment setup depicted in Fig. 5c, element-wise multiplication is achieved by first encoding the image to be processed on an organic light-emitting diode (OLED) serving as the light source. The optical image then passes through a spatial light modulator (SLM), which performs pixel-wise intensity modulation on the input image, thus realizing element-wise multiplication53. (Details are explained in the Supplementary Note 6.1). In this setup, misalignment between the OLED and SLM can severely degrade system performance. Our experiment shows that 1° rotation angle misalignment between the OLED and the SLM can cause a drop in accuracy, from 98.0 to 43.0% on the MNIST dataset. Other misalignments include x- and y-axis shifts between the OLED and the SLM, and the distance change between them (i.e., z-axis shift) that results in image scaling.

In such a system, there is no explicit model between the misalignment factors and controlling parameters, i.e., characterizing the model y = F(x; Θ) is difficult, where Θ represents the misalignment factors such as rotation angle and axis shifts between the OLED and the SLM. To address this, we propose a modified SAT method in which the gradient of the loss with respect to the misalignment factor is estimated by the finite difference method. The optimization flow is shown in Fig. 5d using rotation angle θ as an example. During the first automatic differentiation process, as noted by the yellow arrows, to estimate the gradient, a small perturbation δθ is added to the rotation angle, and the loss is calculated twice, once without this small perturbation and once with it. The gradient with respect to the rotation angle is then estimated using the finite difference method, as depicted in Fig. 5d. After obtaining this gradient dL/dθ, the rotation angle and weights are updated according to Eq. (4),

θadv = θ0 + α1 * (dL/dθ) / || dL/dθ ||^2
Wadv = W0 + α2 * (∇WL(W0)) / || ∇WL(W0) ||^2 (4)

Here, θ0 and W0 denote the initialized rotation angle and weights, θadv and Wadv represent the adversarial rotation angle and weights that reflect the maximum loss within the initial parameters’ neighborhood. The hyperparameters α1 and α2 determine the perturbations added to the rotation angle and weights, respectively. The functions of α1 and α2 are equivalent to α in Eq. (3). Because parameters θ and weights W are optimized separately, we use different subscripts to indicate the difference.

Next, we maintain the rotation angle at θadv during the second automatic differentiation process and conduct forward propagation, as noted by the blue arrows. Then the gradient with respect to weights is calculated again at Wadv. This calculated gradient is used to update the initial weights W0. Finally, the rotation angle is reset to the initial value θ0 as the rotation angle remains constant once calibrated in the real-world system. The training is performed on a computer in an in silico version without considering side effects like dark noises. (The detailed training process is illustrated in Supplementary Note 6.2).

We apply the same optimization approach to address the issues of shift and scaling. The training is performed in silico. After training, we deploy the trained model in our experimental setup to perform NN inference. During the SAT training, we assume perfect alignment between the OLED and SLM. This simplification is intentional, as SAT-trained systems are expected to be inherently robust to implementation errors, even when those errors are unknown. As a result, we expect the system trained by SAT to maintain high accuracy in the presence of practical misalignments such as rotation, shift, or scaling.

To quantitatively assess the tolerance of our method to such misalignment errors at the implementation stage, we follow the approach in ref. 28 and apply affine transformations to introduce controllable deviations (rotation, shift, and scale). We then evaluate the system’s accuracy under these conditions. This procedure enables us to systematically assess the resilience of the trained model to realistic imperfections in the physical setup. The experimental results shown in Fig. 5e illustrate that SAT improves the model’s robustness against the misalignment issues in the free-space computing system. Specifically, SAT maintains MNIST classification accuracy at 98.0% at 1° rotation angle misalignment without accuracy reduction, compared to 43.0% using the standard training method. In terms of pixel-wise shift, we intentionally shift the pixels of the SLM from 0 to 2 pixels in the y-axis direction. The results show that SAT achieves 97.0% accuracy at 1-pixel shift, while the standard training method attains only 36.0% accuracy. Finally, we scale the image size in the SLM to emulate the z-axis shift, with the scaling factor sweeping from 1.0 to 1.1. Our experimental results show that the standard training method achieves only 12.0% accuracy, but SAT improves the accuracy to 93.0% with the scaling factor at the level of 1.05. These results demonstrate that SAT is generally applicable to different computing systems, even for those systems without having an explicitly known model. The modified SAT can successfully address this issue by using the finite difference method to estimate the gradient of the loss with respect to those parameters.

Discussion
In summary, we propose a training method that addresses key limitations in both in silico and in situ training of PNNs.

Through demonstrations on three typical PNNs, an MRR-based PNN, an MZI mesh-based PNN, and a diffractive optical NN, we highlight several distinguishing features of SAT that offer clear advantages over current in silico and in situ training methods: (1) SAT enables accurate in silico training of a PNN, even with an imprecise digital model and under deployment errors; (2) SAT allows systems trained in silico to outperform those trained in situ, even in the presence of fabrication errors; (3) systems trained using SAT exhibit high robustness to environmental fluctuations and noise; (4) SAT can be integrated with in situ training methods, facilitating transferable in situ training across different devices; and (5) SAT is universally applicable to real-world challenging tasks and various PNNs, regardless of whether their models are explicitly known or unknown. Overall, SAT offers a practical, effective, and computationally efficient solution for training and deploying PNNs in real-world applications.

Training cost
While SAT has a slight cost in additional automatic differentiation during training, it remains far more efficient than methods requiring extensive experimental measurements or the training of a large digital twin model. However, the additional training steps can be further reduced using the approach proposed in ref. 54, which reduces the sharpness with almost zero additional computational cost by calculating the Kullback–Leibler divergence between the NN’s outputs using current weights and past weights, as a substitute for SAT’s sharpness measure. This approach will offer the benefits of reduced sharpness without extra overhead.

Joint parameter optimization
During the training of diffractive-optics-based NNs, we perform optimization for each parameter (rotation angle, shift, and scale factor) and evaluate its impact on inference performance. We can also jointly optimize these parameters. Here, we extend the verification of our method to jointly optimizing multiple control parameters. The simulation details and results are depicted in the Supplementary Note 6.3. We propose a sequential joint optimization strategy: optimizing the rotation angle in epoch 1, the pixel shift in epoch 2, and the scale factor in epoch 3. This approach enables joint consideration of all parameters while maintaining the same computational complexity as optimizing a single parameter at a time. Our simulation results show that joint optimization can simultaneously enhance the model’s robustness over different misalignments, while the robustness trained by joint parameter optimization is comparable to that of single parameter optimization.

Compatible with nonlinear systems
Although our experimental demonstration uses a linear system and assumes digital nonlinearity in simulations, SAT can naturally be extended to systems with physical nonlinearity. Because the gradient approximation process is the same for linear and nonlinear systems. To prove this, we also include an additional simulation study on a nonlinear deep diffractive neural network described in ref. 21. The simulation details and results are included in Supplementary Note 6.4. The simulated system consists of two programmable linear layers and two nonlinear layers. The first nonlinear layer is implemented using a saturable absorber, which introduces an intensity-dependent transmission while preserving the optical phase. And the second nonlinear component is the PD array, which introduces a quadratic nonlinearity during the optical-to-electrical conversion process. We train the physical network using both standard BP and our SAT. During training, we employ the finite difference method to approximate the loss gradient with respect to the rotation angle (which acts as a trainable control parameter), and we intentionally introduce perturbations in rotation to evaluate robustness. The simulation results show our method maintains high classification accuracy under rotation angle variations.

Methods
Dataset
The MNIST dataset contains 70,000 grayscale images of handwritten digits (0–9)48. These include 60,000 training examples and 10,000 examples, each formatted as a 28×28 pixel grid with corresponding digit labels. The CIFAR-10 dataset contains 60,000 color images categorized into 10 distinct classes (airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck), including 50,000 training examples and 10,000 test examples55. Each image is formatted as a 32×32 pixel grid (RGB color channels) with corresponding class labels.

Device fabrication
The device is fabricated in the commercial multi-project wafer process from Applied Nanotools. The silicon thickness of this process is 220nm, and the buffer oxide layer thickness is 2μm. The MRR weight bank contains 16 MRRs and is arranged in a 4×4 matrix. Each row sub-weight bank consists of four MRRs coupled with two bus waveguides in an add/drop configuration. The radii for each MRR are around 20μm, and we intentionally introduce a slight 10nm radius difference to avoid resonance collision. And the result show that resonance wavelengths are roughly spaced by 0.8nm. The gap between the MRR and the bus waveguide is 100nm. MRR Q factor is around 10^4. Circular metal heaters (TiW) are built on top of each MRR for thermal weight tuning. Metal (Al) vias and traces are deposited to connect the heater contacts of the MRR weight bank to the electrical metal pads. The heater resistance is around 375Ω. The tuning efficiency is around 0.53nm/mW.

MRR-based PNN experimental setup
During the experiment, we selected four MRRs with the labels “00,” “03,” “12,” and “15.” Before the experiment, we first actuate initialization currents to align the resonance wavelengths of “00” & “03” and “12” & “15.” After initialization, the resonance wavelengths of “00” & “03” are both 1542.7nm, and the resonance wavelengths of “12” & “15” are both 1544.8nm. A TEC module is set to stabilize or intentionally change the temperature. The laser (KG-TLS-C-13-50-P-FA, Conquer) wavelengths are set to 1543.0 and 1545.0nm, respectively. Different wavelength lights are independently modulated using electrical-optical modulation (KG-DDMZM-RF-B, Conquer). The modulation signals are set with a baud rate of 20 MBaud (AFG 31000 SERIES, Tektronix). The different wavelength lights are then combined through a 50:50 coupler. Before coupling into the chip, the combined light is amplified through an EDFA (FA-30, PriTel) to around 12dBm. The coupling loss of the chip is around 14dBm. The two-channel output lights are then detected by the PD (PDT0313-FC-A, HP) at the through port. During the matrix multiplication, the negative weights are mapped to the tuning range [0, 1].

Details about training the digital-optical hybrid network
The digital-optical hybrid network is built based on the open-source library “DATMPNN”28 and “Neurophox”56. The network is a modified version of the classic LeNet-5 architecture, designed for image classification tasks, with an optical MZI mesh layer for classification. Details are clarified in Supplementary Note 5.

Specifically, we test 6 different training methods: in silico standard BP, in silico NAT, in situ PAT, in situ DAT, in silico SAT, and in situ SAT. The training device is an Nvidia 3070 Ti, and the training time is measured by directly calculating the difference between the start time and end time of the program run. The in silico standard BP involves training the ideal model directly for 40 epochs. During training with NAT, the noise we add to the phase parameter follows a Gaussian distribution with a mean value of 0 and a standard deviation of 0.01rad. The aim is to maximally increase the robustness without decreasing the accuracy. For PAT, training involves performing inference on the physical system and BP on the ideal system. The physical system deviates from the ideal system due to static phase errors and splitting errors. The total number of training epochs is 40 for accuracy and loss to converge.

In the case of DAT, an additional digital network, SEPN, is introduced to compensate for the gap between the physical system and the ideal system. DAT requires more training epochs to converge, totaling 60 epochs. During the first 30 epochs, the training process is the same as the PAT, and SEPN is trained to reduce the output difference between the physical system and the ideal system, but it is not used for gradient calculations. In the subsequent 30 training epochs, SEPN training stops, and SEPN assists the ideal model in providing accurate gradients. To minimize the gap as much as possible, we use 6 complex mini-UNet networks connected in parallel to extract multiscale features. The mini-UNet first downsamples the input image from its original size (H, W) to (H/2, W/2). It then further reduces this to (H/4, W/4). Each of these three feature maps, at different spatial resolutions, is processed independently by a dedicated complex convolutional NN. Finally, the processed feature maps are fused and passed through a series of transposed convolutions to reconstruct an output image with the same size as the original input. The detailed network of mini-UNet is the same as that in ref. 28 and is illustrated in Supplementary Note 5 and Fig. S17.

For in silico SAT, standard BP is still used for differentiation, but differentiation is performed twice per epoch. The total number of training epochs remains 40. For in situ SAT, differentiation is provided by the ideal model, which is the same as PAT, and differentiation is performed twice per epoch. The total number of training epochs is 40. The initial learning rate for training the PNN is 0.1, which is decayed by 0.5 at half of the maximum training epoch number. The learning rate for training the SEPNs is set at 0.001 for constant. The cross-entropy loss function is selected to calculate the task loss, and the MSE loss function is selected for adaptive training of the SEPNs.

Free-space experimental setup
In the free-space experimental setup, a custom OLED display (Sony ECX335S) is used, encoding pictures as the light source. Only the green channel is utilized in the experiments. The display has a resolution of 1920×1080 and a refresh rate of 60Hz. We develop custom Python control codes to load images. The display features 256 distinguishable brightness levels, corresponding to an 8-bit resolution. We achieve the weights through light intensity modulation using an SLM (LCOS-SLM, HDSLM80R Plus, UPOlabs) with a polarizing beam splitter (PBS) and a half-wave plate (HW) in a double-pass configuration. A zoom-lens microscope (BYH0330, Inseinlifung) is used between the OLED and SLM to match the pixel sizes between the OLED (5μm) and SLM (8μm). The light is then projected onto a camera using a second microscope (SHL-0745C, Shhunhuali). A CMOS camera (MV-SUF89OGC/M, Mindvision) captures the modulated light field as an image. The image is separated into different regions, and the sum of the intensities from these regions represent the matrix multiplication result.

Data availability
The raw datasets are all publicly available. MNIST48 and CIFAR-1055 datasets are downloaded from previous works. Source data are provided via Zenodo at https://doi.org/10.5281/zenodo.1790544057.

Code availability
The source code supporting the findings of this study has been deposited in the GitHub repository at https://github.com/cuhkhuangslab/Sharpness-Aware-Training/tree/main. The repository is publicly accessible under the MIT License, permitting reuse and modification with appropriate citation. The software code is also available via Zenodo at https://doi.org/10.5281/zenodo.1790544057.

Peer review rounds
[{"round": 1, "reviewer_comments": {"reviewer_1": "Xu et al. propose to apply a technique from conventional machine learning, <<sharpness-aware training>> (SAT), to training physical neural networks (PNNs), and empirically observe that it works well when applied to some optical neural networks. While SAT was initially introduced to improve generalization, here the authors find that it also makes training more robust to noise/drifts in physical systems.\nI believe the idea is sound and the results show that SAT could be generally useful in PNNs. Because of the broad utility of the approach, I strongly recommend publication of the manuscript in Nature Communications subject to the authors addressing the following comments.\n1. The paper seemingly overclaims novelty: SAT was not invented by the authors, but rather was previously reported in references [42] and [50]. The current abstract says: \"...through a novel technique called Sharpness-Aware Training\". But the technique is not novel! (Indeed it was reported 4 years ago, in the most prominent machine-learning conference.) It would be more appropriate for the abstract (and other parts of the paper) to describe what they have done as apply the (established; not novel) technique SAT to the training of PNNs and observe that it results in improved robustness to model imperfections and PNN noise/drifts etc.\n2. The standard approach to training a physical (analog) neural network for which one has a fairly accurate model is to perform training in a digital computer but to inject noise during training (sometimes called noise-aware training: doi:10.1109/TED.2021.3089987 ; review: doi:10.1002/aisy.202200029). For Figure 2j and for Figure 5c, it would be useful if the authors could add a comparison to noise-aware training (NAT) or related standard noise-injection techniques. I also think it would be useful for physics-aware training (PAT) to be added as a point of comparison in Figure 2j.\n3. The authors have not provided the code to replicate their results. Especially for a methodology paper, this is unacceptable and against Springer Nature's stated policy on code availability. The authors should provide the code in a GitHub (or similar) repository. At the very least, this repository should contain everything needed to rerun the simulations reported in Figure 5.", "reviewer_2": "The article introduces a novel training approach, Sharpness-Aware Training (SAT), aimed at enhancing both the training process and deployment robustness of physical neural networks (PNNs). I appreciate the insight that the authors establish a connection between sharp minima in the loss landscape and the poor robustness observed in physical systems. I believe this work has the potential to become a foundational methodology for training PNNs—specifically, by incorporating a sensitivity term of the objective function into the loss formulation to promote the learning of more robust parameters. In addition, I have a few questions and suggestions that I hope the authors can clarify, which I believe would further improve the quality of the paper.\n1.When the influence of certain parameters on the physical process is not explicitly known, approximating the gradient via finite differences offers an effective way to bypass the need for explicit modeling. Figure 3(e) illustrates the separate optimization processes for rotation, shift, and scale parameters. Have the authors considered jointly optimizing these parameters instead of optimizing them independently?\n2.Lines 390–391 state: “During the experiment, we use the affine function to quantitatively adjust rotation, shift, and scaling parameters.” Is it truly feasible to achieve such precise control over angle and pixel-level shifts in a physical experimental setup? Figure 3e appears to show results from a real experiment, but in practice, these deviations—especially angular or spatial misalignments—are typically not precisely known or controlled.\n3.Line 297 states that “the system trained with standard BP degrades significantly from 80.0% to 7.0% as the temperature varies from 21°C to 23°C.” However, I could not clearly observe this drastic performance drop in Figure 2(h). Could the authors clarify this discrepancy or improve the clarity of the visualization?\n4.I am curious about the tolerance limits of the proposed method. Specifically, how large can the deviations (e.g., phase errors, temperature shifts) be before the performance of the SAT-trained model degrades significantly or fails altogether?\n5.The method effectively mitigates performance degradation caused by various noise sources, but how is the trade-off between sharpness and loss handled in practice? How should the hyperparameter \\alpha be chosen? A similar issue arises with the self-defined hyperparameters \\alpha_1, r, \\mu, and \\rho. It would be beneficial if the authors could provide a principled guideline or sensitivity analysis.\n6.Minor issues: In Figure 2(h), it would be more accurate to describe the performance shown as the ideal performance rather than the theoretical limit. In Equation (4) of the main text, the numerator appears to be missing the letter \"L\".\n7.I strongly encourage the authors to consider releasing the code, even if only for the MZI-based implementation, which would significantly enhance the impact and reproducibility of this work.", "reviewer_3": "This manuscript introduces the use of Sharpness-Aware Minimization (SAM), a technique originally developed in the machine learning domain, to improve the robustness of physical neural networks. By simultaneously minimizing the loss and its sharpness, the authors aim to eliminate the need for in-situ training, thereby enabling the network to tolerate environmental fluctuations and fabrication variance. The effectiveness of this approach is experimentally evaluated across three distinct hardware platforms.\nWhile the manuscript presents a technically sound implementation and provides extensive experimental validation, I have the following major concerns regarding its novelty and overall impact:\n1. Lack of originality in methodology: The training is entirely performed in silico, with only the deployment carried out on physical platforms. The core idea—SAM—is directly borrowed from machine learning literature and is already well-established.\n2. Limited conceptual advancement: While the experimental implementation differs from the authors’ earlier work, the underlying idea and overall conceptual direction remain highly similar. This raises concerns about a lack of novelty in terms of scientific framing and motivation. Such conceptual continuity may be appropriate for an incremental study, but may not meet the threshold for a high-impact general-interest journal.\n3. Absence of full neural network functionality: Despite being framed as work on “physical neural networks,” the experimental demonstrations are limited to simple linear functions. No physical implementation of nonlinear activation or full inference capability is presented.\n4. Presentation issues: The manuscript contains several typographical issues and instances of repeated text. A careful proofreading is necessary to ensure clarity and professionalism in presentation.\nThat said, I appreciate the authors’ effort in generating high-quality figures and conducting thorough, cross-platform experimental validation. These aspects reflect a strong technical execution.\nIn summary, while the paper demonstrates a useful application of an existing training technique to physical systems, the work currently does not meet Nature Communications' expectations for innovation and conceptual significance. Substantial improvements in both novelty and experimental depth would be required for this work to be considered for publication."}, "author_response": "Dear Reviewers,\nThank you very much for taking your time to review our manuscript (ID: NCOMMS-25-26432-T) titled\n“Perfecting Imperfect Physical Neural Networks using Sharpness-Aware Training”. We appreciate all your\ncomments and suggestions. We have carefully revised the manuscript and the Supplementary Notes accordingly.\nThe modifications to the manuscript and Supplementary Notes are highlighted in blue.\nDetailed point-by-point responses are provided below:\nResponse To Reviewer #1\nComment 1\nXu et al. propose to apply a technique from conventional machine learning, <<sharpness-aware training>>\n(SAT), to training physical neural networks (PNNs), and empirically observe that it works well when applied\nto some optical neural networks. While SAT was initially introduced to improve generalization, here the\nauthors find that it also makes training more robust to noise/drifts in physical systems.\nI believe the idea is sound and the results show that SAT could be generally useful in PNNs. Because of the\nbroad utility of the approach, I strongly recommend publication of the manuscript in Nature Communications\nsubject to the authors addressing the following comments.\nResponse 1\nWe sincerely appreciate the reviewer’s high recognition of our work and valuable suggestions.\nComment 2\n1. The paper seemingly overclaims novelty: SAT was not invented by the authors, but rather was previously\nreported in references [42] and [50]. The current abstract says: \"...through a novel technique called Sharpness-\nAware Training\". But the technique is not novel! (Indeed it was reported 4 years ago, in the most prominent\nmachine-learning conference.) It would be more appropriate for the abstract (and other parts of the paper) to\ndescribe what they have done as apply the (established; not novel) technique SAT to the training of PNNs and\nobserve that it results in improved robustness to model imperfections and PNN noise/drifts etc.\nResponse 2\nWe thank the reviewer for pointing out this issue. We realize that we did not clearly explain the relationship\nbetween Sharpness-Aware Minimization (SAM) and Sharpness-Aware Training (SAT). The wording in the\noriginal abstract was inappropriate.\nYou are correct that SAT is inspired by the machine learning SAM, which was proposed to enhance generalization\ncapabilities of AI models. Its core insight is the connection between the sharpness of the loss landscape and a\nmodel’s generalization performance and leverage this connection to enhance generalization capabilities of AI\nmodels. Our work is inspired by this conceptual framework, however the originality of our contribution is to\nestablish a new connection—between the sharpness of the loss landscape and the robustness of physical\nsystems. Guided by this new perspective, we propose SAT tailored specifically for physical systems, which\nenables robustness to hardware imperfections, environmental disturbance, and transferrable learning, for both\nin-silico and in-situ training. Such cross-disciplinary integration is the essence of interdisciplinary innovation.\nIn addition to introducing a new conceptual framework, our work also makes a significant methodological\ncontribution. The original SAM cannot be directly applied to physical systems. SAM is designed to optimize\nweights and biases in AI models that are explicitly defined by precise mathematical functions. However, as\nwe pointed out in Section 2.1 of the Principle section, the mathematical models of PNNs are often imprecise\nor even unknown. To address this, SAT extends from SAM in two ways: (1) it accommodates cases where the\nunderlying model of the PNN is not explicitly known, and (2) it reformulates the standard weight optimization\ninto a control-parameter optimization, allowing direct training of the physical system. These new\ndevelopments make SAT highly generalizable and suitable for a wide range of physical systems, whether or\nnot a precise model is available, and supports both in-silico and in-situ training scenarios.\nThank you for pointing out the inappropriate statement in the abstract. In the revised version, we have clarified\nthe relationship between SAM and SAT as requested.\nRevision 2\nWe clarify the relationship between SAM and SAT as requested in the Abstract (Line 21-25) and Manuscript\n(Line 91-95).\n“Here, we address the challenges with both in-silico and in-situ training through Sharpness-Aware Training\n(SAT), \\textcolor{blue}{inspired by a machine learning technique Sharpness-Aware Minimization\n(SAM)~\\cite{foret2021iclr} that links loss landscape geometry with model generalization to enhance model\ngeneralization. In this work, we establish a new link between loss landscape geometry with robustness of physical\nsystems, and then leverage this connection to tackle above challenges in training physical systems.”\n“Furthermore, SAT develops methodologies for automatically locating flat minima in physical systems and\nenables direct training in physical systems.”\nComment 3\n2. The standard approach to training a physical (analog) neural network for which one has a fairly accurate\nmodel is to perform training in a digital computer but to inject noise during training (sometimes called\nnoise-aware training: doi:10.1109/TED.2021.3089987 ; review: doi:10.1002/aisy.202200029). For Figure\n2j and for Figure 5c, it would be useful if the authors could add a comparison to noise-aware training (NAT)\nor related standard noise-injection techniques. I also think it would be useful for physics-aware training\n(PAT) to be added as a point of comparison in Figure 2j.\nResponse 3\nWe thank the reviewer for raising the questions regarding Noise-Aware Training (NAT) and Physical-Aware\nTraining (PAT).\nWe are aware that NAT is widely adopted to enhance model robustness against system imperfections [1-4]. In\nour original manuscript, we have compared our method with NAT in the Microring resonator (MRR)-based\noptical neural networks (ONNs), but presented the results in Supplementary Note 4.1 rather than in the main text.\nIn the revised version, we have now included this comparison directly in the main manuscript (Figure 2j). We\nalso add an additional comparison in the Mach-Zehnder Interferometer (MZI)-based ONNs (Figure 5c), where\nthe results also clearly demonstrate the superior performance of SAT over NAT.\nNAT enhances robustness by injecting synthetic noise into the model during training, typically assuming simple\nadditive Gaussian distributions. However, imperfections in physical systems often deviate significantly from such\nsimplified Gaussian assumptions. These imperfections may include modeling errors and implementation-related\nvariations such as pixel-level misalignment, fabrication errors, and temperature-induced shifts, which are non-\nGaussian. As a result, NAT cannot fully capture the complex, system-specific error characteristics and thus has\nlimitations in addressing robustness issues arising from these non-Gaussian imperfections.\nIn contrast, SAT makes no prior assumptions about the nature of system imperfections. Instead, it defines a clear\noptimization objective (i.e., sharpness minimization) that is generally applicable across different systems. By\nlinking system robustness to the fundamental geometry of the loss landscape, SAT inherently accounts for a wide\nrange of imperfections, including modeling inaccuracies, environmental drift, and hardware noise, without\nrequiring explicit modeling of each. This enables SAT to optimize PNNs more accurately and robustly, even\nwhen the exact form of imperfections is unknown.\nTo prove the advantage, we first train the MRR-based PNNs with NAT and SAT respectively. The training model\nis obtained at the room temperature of 22°C. When training with NAT, a weight noise, which follows a gaussian\ndistribution with mean value at 0, and standard deviation at 0.3, is injected. The reason for choosing this level\nnoise is to maximally increase the robustness while without decreasing the accuracy. At the inference stage, the\nTEC is removed and the room temperature becomes 20°C. Despite the temperature change being as large as 2°C,\nSAT maintains a high accuracy of 91.0% while the accuracy of NAT is reduced to 61.0%, as shown in Figure R1.\nThe result now is included in the directly in the Figure 2j of main manuscript. Another evidence that SAT\noutperforms NAT that the sharpness metric 𝜆 of SAT is only 1.2 while that of NAT is 414.5 (Figure R2). A\nlarger 𝜆 indicates higher sensitivity to parameter perturbations.\nFigure R1. Inference accuracy on microring circuits and confusion matrix when removing the thermoelectric\ncontroller (TEC).\nFigure R2. Loss landscape and sharpness metric 𝜆 of the neural network trained by different methods.\nSimilar comparsion is also conducted on MZI-based ONN. During training with NAT, the noise we add to the\nphase parameter follows a Gaussian distribution with mean value at 0, and standard deviation at 0.01 rad. The\naim is to maximally increase the robustness while without decreasing the accuracy. The learning rates are set at\n0.1 for both NAT and SAT, and the training epoch numbers are 40 for both NAT and SAT. After training, we\nevaluate the robustness by adding the phase error and beam splitter error simultaneously at 0.15 rad. And the\naccuracy for NAT is 79.8% compared to SAT at 97.1%. This additional result is added in Figure 5c in the revised\nmanuscript and shown in the response letter Figure R3.\nFigure R3. Training result performance comparison.\nPAT is an in-situ training method designed to enable backpropagation directly on physical systems. However, it\ndoes not address resilience issues during or after training. Like other existing in-situ training approaches, PAT\nsuffers from two key limitations: (1) The trained system is vulnerable to time-dependent degradation caused by\nfactors such as noise, environmental disturbances, and hardware variability. (2) The trained parameters are not\ntransferable to other systems due to device variations caused by fabrication errors.\nIn contrast, SAT addresses both of these limitations, and SAT can be implemented with in-situ training setting.\nAs shown in our revised manuscript of ‘Section 2.3 In-situ SAT: facilitating transferable in-situ training’ (‘Section\n2.4 Facilitate transferable in-situ training’ in the original manuscript), we compared SAT with PAT and evaluated\nthe sensitivity of the trained models on MZI-based ONNs using the maximum eigenvalue (𝜆 ) of the Hessian\nmatrix. A larger 𝜆 indicates higher sensitivity to parameter perturbations, and thus lower robustness. SAT\nexhibits significantly lower 𝜆 compared to PAT (PAT: 114.7, in-situ SAT: 0.6), indicating that models trained\nwith SAT are much more resilient to parameter variations and system disturbances.\nThis is further supported by the new inference results (Figure R1) in the MRR-based ONNs as requested by the\nreviewer, which demonstrate that PAT fails to maintain performance under ambient temperature fluctuations\nwhen the thermoelectric controller (TEC) is removed and without retraining. In contrast, SAT maintains\nconsistent robustness under such variations. To be more specific, a typical PAT training contains a real physical\nforward model and an approximate digital backward model. The forward physical model considers fabrication\nvariance, while during the backward period, the backward digital model does not contain the fabrication\nvariance. In this circumstance, the accuracy converges to 96.0% during training. But when removing the\nthermoelectric controller (TEC), the accuracy drops significantly to 59.0%. The result indicates that PAT fails\nto maintain performance under ambient temperature fluctuations.\nRevision 3\n1. We add the NAT trained neural network loss landscape to Figure 2g, remove the NAT training results from\nthe Supplementary Note 4.1 to Figure 2j, and explain the reason why SAT outperforms NAT (Line 181-291).\n2. We add the NAT training results in Mach-Zehnder interferometer (MZI) mesh to Figure 5c, and include\nthe simulation details into the manuscript (Line 336-341).\nComment 4\n3. The authors have not provided the code to replicate their results. Especially for a methodology paper, this is\nunacceptable and against Springer Nature's stated policy on code availability. The authors should provide the\ncode in a GitHub (or similar) repository. At the very least, this repository should contain everything needed to\nrerun the simulations reported in Figure 5.\nResponse 4\nWe thank the reviewer for pointing our the issue regarding code availability. Therefore, we have summarized all\nthe simulation codes of different computing systems and store them in the GitHub repository.\nhttps://github.com/cuhkhuangslab/Sharpness-Aware-Training/tree/main\nReferences\n[1] Kariyappa, S., Tsai, H., Spoon, K., Ambrogio, S., Narayanan, P., Mackin, C., ... & Burr, G. W. (2021). Noise-\nresilient DNN: Tolerating noise in PCM-based AI accelerators via noise-aware training. IEEE Transactions on\nElectron Devices, 68(9), 4356-4362.\n[2] Yang, X., Wu, C., Li, M., & Chen, Y. (2022). Tolerating noise effects in processing-in-memory systems for\nneural networks: a hardware–software codesign perspective. Advanced Intelligent Systems, 4(8), 2200029.\n[3] Wu, C., Yang, X., Yu, H., Peng, R., Takeuchi, I., Chen, Y., & Li, M. (2022). Harnessing optoelectronic noises\nin a photonic generative network. Science advances, 8(3), eabm2956.\n[4] Mourgias-Alexandris, G., Moralis-Pegios, M., Tsakyridis, A., Simos, S., Dabos, G., Totovic, A., ... & Pleros,\nN. (2022). Noise-resilient and high-speed deep learning with coherent silicon photonics. Nature communications,\n13(1), 5572.\nResponse To Reviewer #2\nComment 1\nThe article introduces a novel training approach, Sharpness-Aware Training (SAT), aimed at enhancing both\nthe training process and deployment robustness of physical neural networks (PNNs). I appreciate the insight\nthat the authors establish a connection between sharp minima in the loss landscape and the poor robustness\nobserved in physical systems. I believe this work has the potential to become a foundational methodology for\ntraining PNNs—specifically, by incorporating a sensitivity term of the objective function into the loss\nformulation to promote the learning of more robust parameters. In addition, I have a few questions and\nsuggestions that I hope the authors can clarify, which I believe would further improve the quality of the paper.\nResponse 1\nWe sincerely thank the reviewer for the high recognition of our work. Your comments will be invaluable for\nimproving our manuscript.\nComment 2\n1.When the influence of certain parameters on the physical process is not explicitly known, approximating the\ngradient via finite differences offers an effective way to bypass the need for explicit modeling. Figure 3(e)\nillustrates the separate optimization processes for rotation, shift, and scale parameters. Have the authors\nconsidered jointly optimizing these parameters instead of optimizing them independently?\nResponse 2\nWe thank the reviewer for raising the question regarding joint parameter optimization.\nIn the original manuscript, we performed optimization for each parameter independently and evaluated its\nindividual impact on performance. However, our method is also applicable to joint optimization of multiple\nparameters. To address the reviewer’s request, we propose a sequential joint optimization strategy: optimizing\nthe rotation angle in epoch 1, the pixel shift in epoch 2, and the scale factor in epoch 3. This approach enables\njoint consideration of all parameters while maintaining the same computational complexity as optimizing a single\nparameter at a time.\nNow, we have included the joint optimization result in the manuscript (Line 448-456) and the Supplementary\nNote 6.3.\nTo be more specific, we first take rotation angle, shift pixel and scale factor as the joint optimization parameters.\nAfter training, we evaluate the model’s robustness through a 2D sweep over rotation angle and shift pixel and\nkeep the scale factor at 1. As shown in Figure R4, our method maintains high robustness against both rotation\nangle and shift pixel changes after the joint optimization. Next, the robustness of the model trained with joint\nparameter optimization is compared with that of models trained with single-parameter optimization for rotation\nangle, shift pixel, and scale factor, respectively. As shown in Figure R5, joint optimization achieves comparable\nrobustness to that of single-parameter optimization. This result demonstrates that our method is applicable to\njoint optimization of multiple parameters without degrading the robustness of the trained model.\nFigure R4. 2D accuracy sweep with different rotation angle and shift pixel.\nFigure R5. Inference accuracy with different rotation angle, shift pixel and scaling number.\nRevision 2\nWe include the joint optimization result in the manuscript (Line 454-462) and the Supplementary Note 6.3.\nComment 3\n2.Lines 390–391 state: “During the experiment, we use the affine function to quantitatively adjust rotation,\nshift, and scaling parameters.” Is it truly feasible to achieve such precise control over angle and pixel-level\nshifts in a physical experimental setup? Figure 3e appears to show results from a real experiment, but in\npractice, these deviations—especially angular or spatial misalignments—are typically not precisely known or\ncontrolled.\nResponse 3\nYou are correct that such deviations are not controllable in real-world implementations. However, the purpose of\nadjusting these parameters is not to control them, but rather to quantitatively evaluate the tolerance of our method\nto these variations.\nDuring SAT training, we assume perfect alignment between the OLED and SLM. This simplification is\nintentional, as SAT trained systems are expected to be inherently robust to implementation errors, even when\nthose errors are unknown. As a result, we expect the system trained by SAT to maintain high accuracy in the\npresence of practical misalignments such as rotation, shift, or scaling.\nTo quantitatively assess the tolerance of our method to such misalignment errors at the implementation stage, we\nfollow the approach in [1] and apply affine transformations to introduce controllable deviations. We then evaluate\nthe system's accuracy under these conditions. This procedure enables us to systematically assess the resilience of\nthe trained model to realistic imperfections in the physical setup.\nRevision 3\nWe revise the manuscript to illustrate our experimental implementation process with more details (Line 416-\n425).\nComment 4\n3.Line 297 states that “the system trained with standard BP degrades significantly from 80.0% to 7.0% as the\ntemperature varies from 21°C to 23°C.” However, I could not clearly observe this drastic performance drop in\nFigure 2(h). Could the authors clarify this discrepancy or improve the clarity of the visualization?\nResponse 4\nWe thank the reviewer for pointing out this issue. The confusion was caused by a typo. The accuracy of standard\nBP degrades from 80.0% to 17.0%, not 7.0%. We apologize for the confusion.\nRevision 4\n1. Correct the typo 7.0% to 17.0%.\n2. We revise the Figure 2(h) and add the grid on the Figure 2(h) to improve the clarity of the visualization.\n[Original Figure 2h]\n[Revised Figure 2h]\nComment 5\n4.I am curious about the tolerance limits of the proposed method. Specifically, how large can the deviations\n(e.g., phase errors, temperature shifts) be before the performance of the SAT-trained model degrades\nsignificantly or fails altogether?\nResponse 5\nWe thank the reviewer for raising the question regarding the tolerance limits.\nThe tolerance limitation is presented in Figure R6. For temperature tolerance, our experimental results show that\nwhen the temperature variation is within ±1 ℃, the accuracy on the MNIST dataset can be maintained at around\n95.0% (idea performance is 98.0%). When the temperature variation further increases, the accuracy drops\nsignificantly. For more complex tasks, such as image reconstruction and generation, the system becomes more\nsensitive to imperfections. In these cases, performance metrics such as Mean square error (MSE) and Fréchet\nInception Distance (FID) remain stable at around 0.023 (Ideal MSE is 0.013) and 120.59 (Ideal FID is 100.31)\nwithin a temperature fluctuation of ±0.25 ℃. Beyond this threshold, for example, when the temperature variation\nreaches +0.5℃, the image quality metrics MSE and FID degrades to 0.547 and 426.69, respectively.\nIn our results, we select ±0.25 ℃ as the thermal fluctuation limit, as it represents the maximum expected\nfluctuation in current co-packaged optical modules [2].\nRegarding phase errors, we set both the standard deviation of phase errors (σ ) and the beam splitter errors (σ )\nps bs\nto 0.15 rad. These values correspond to typical wafer-scale variations observed in current standard manufacturing\nprocesses [3]. Our results show that under such conditions, SAT can maintain an accuracy of 95.2% (the ideal\nperformance is 97.1%), whereas without SAT, the accuracy drops to 58.6%. When the errors are further increased\nto 0.20 rad, even with SAT, the performance degrades to 61.3%.\nFigure R6. Inference accuracy with different levels of temperature change and phase errors.\nComment 6\n5.The method effectively mitigates performance degradation caused by various noise sources, but how is the\ntrade-off between sharpness and loss handled in practice? How should the hyperparameter \\alpha be chosen?\nA similar issue arises with the self-defined hyperparameters \\alpha_1, r, \\mu, and \\rho. It would be beneficial\nif the authors could provide a principled guideline or sensitivity analysis.\nResponse 6\nWe thank the reviewer for raising the question about the trade-off between sharpness and loss, as well as the\nprincipled guideline.\nAfter carefully reading the manuscript and Supplementary Notes, we find some defined hyperparameters are\nredundant. Now we have simplified all definitions, and we apologize for any possible confusions. The conclusion\nis, parameters defined in the initial manuscript and Supplementary Note 𝛼 , 𝑟, 𝜇, and 𝜌, are not necessary. The\nonly required hyperparameter is 𝛼. Moreover, The value of the hyperparameter 𝛼 determines the penalty of\nsharpness for loss during the actual training process, that is, the trade-off between sharpness and loss.\nIn the following response, we will first introduce how to simplify the redundant hyperparameters to only 𝛼, and\nintroduce our principled guideline for selecting the 𝛼 value through sensitivity analysis.\nThe hyperparameters 𝛼, 𝛼 , 𝑟 were noted in the section 2.1 of the original maintext, and the hyperparameters 𝜇,\nand 𝜌 were noted in section 2.3. Here, we show how to simplify these hyperparameters step by step.\nIn section 2.1, we initially defined the loss function ℒ as (Equation (2) in the original main\ntext),\nℒ = ℒ)𝑦,𝑦 ;Θ.+𝛼0\n)ℒ+,,,!\"#$%!;/0\n0 (𝑅1)\n%\"&'(%\n)/ 1\nThe hyperparameter 𝛼 indicates how large the penalty of the regularization term should be added to the initial\nloss function.\nTo simplify the computational complexity, we applied the method in [4] by approximating the Hessian matrix\nthrough first-order Taylor expansion, as shown in Equation (R2-R6) (The original Equation (S2-S6) in\nSupplementary Note 1.1).\n𝜕ℒ)𝑦,𝑦 ;Θ.\n%\"&'(%\n𝜕6 6\n𝜕Θ 𝜕1ℒ(Θ) 𝜕ℒ⁄𝜕Θ\n1 = ⋅ (𝑅2)\n𝜕Θ 𝜕Θ1 ‖𝜕ℒ⁄𝜕Θ‖\n1\n∂ℒ(Θ)\nℒ(Θ+ΔΘ) ≈ ℒ(Θ)+ ⋅ΔΘ (𝑅3)\n∂Θ\n𝜕ℒ(Θ+ΔΘ) 𝜕ℒ(Θ) 𝜕1ℒ(Θ)\n≈ + ⋅ΔΘ (𝑅4)\n𝜕Θ 𝜕Θ 𝜕Θ1\n𝜕ℒ)𝑦,𝑦 ;Θ.\n%\"&'(%\n𝜕6 6\n𝜕Θ\n1 𝜕ℒ(Θ+ΔΘ) 𝜕ℒ(Θ)\n1 ≈ @ − B (𝑅5)\n𝜕Θ 𝑟 𝜕Θ 𝜕Θ\n𝜕ℒ (𝛩) 𝜕ℒ(𝛩) 𝛼 𝜕ℒ(𝛩+𝛥𝛩) 𝜕ℒ(𝛩) 𝜕ℒ(𝛩) 𝜕ℒ(𝛩+𝛥𝛩) 𝜕ℒ(𝛩)\n! = + * − -| = +𝛼 * − -| (𝑅6)\n𝜕𝛩 𝜕𝛩 𝑟 𝜕𝛩 𝜕𝛩 \"#$% &ℒ⁄&# 𝜕𝛩 ! 𝜕𝛩 𝜕𝛩 \"#$% &ℒ⁄&#\n‖&ℒ⁄&#‖! ‖&ℒ⁄&#‖!\nDuring the above calculation, the new hyperparameters 𝑟\nand 𝛼 were introduced. The hyperparameter 𝑟\nindicates how large the perturbation ΔΘ should be given to Θ.\nWe directly set 𝑟 = 𝛼 during training all the ONNs and it did not influence the final training results. In this\ncircumstance, 𝛼 equals 1. Therefore, the above 3 hyperparameters 𝛼, 𝛼 , 𝑟 are simplified into only 𝛼. And\nEquation (R6) is simplified into Equation (R7),\n𝜕ℒ (Θ) 𝜕ℒ(Θ) 𝛼 𝜕ℒ(Θ+ΔΘ) 𝜕ℒ(Θ) 𝜕ℒ(Θ+ΔΘ)\n1 = + @ − B| = | (𝑅7)\n𝜕ℒ⁄𝜕Θ 𝜕ℒ⁄𝜕Θ\n𝜕Θ 𝜕Θ 𝛼 𝜕Θ 𝜕Θ ΔΘ=𝛼 𝜕Θ ΔΘ=𝛼\n‖𝜕ℒ⁄𝜕Θ‖\n2\n‖𝜕ℒ⁄𝜕Θ‖\nin Equation (R8).\n𝜕𝐿/𝜕Θ\nΔΘ = 𝛼\n‖𝜕𝐿/𝜕Θ‖\n1\n𝜕ℒ 𝜕ℒ(Θ+ΔΘ)\n$ = | (𝑅8)\n)ℒ⁄)/ \n𝜕Θ 𝜕Θ 4/56\n‖)ℒ⁄)/‖\n&\nIn section 2.3, we originally defined two new hyperparameters, 𝜇 and 𝜌. The meaning of these two\nhyperparameters is still how large the perturbation should be given to the corresponding parameters, rotation\nangle 𝜃 and weights 𝑊. Initially we used two hyperparameters, 𝜇 and 𝜌 to indicate that their values can be\nselected independently.\nNow we have changed the hyperparameters defined in Equation (5) and (6) (Now Equation (4) in the main\ntext) in the main text to 𝛼 and 𝛼 (originally was 𝜇 and 𝜌), as shown in Equation (R9).\n7ℒ⁄7:\n𝜃 = 𝜃 +𝛼\n\"78 9 $\n‖7ℒ⁄7:‖\n&\n∇ ℒ(W )\n; %\nW = W +𝛼 (𝑅9)\n\"78 % 1‖∇ ℒ(W )‖\n; % 1\nHere 𝛼 and 𝛼 determine the perturbations added to the rotation angle and weights, respectively. The function\nof 𝛼 and 𝛼 are equivalent to 𝛼 in Equation (R1). Because parameters Θ and weights 𝑊 are optimized\nseparately, we use different subscript to indicate the difference.\nAbove all, all previously defined hyperparameters can be simplified into a single hyperparameter 𝛼, which defines\nhow large the perturbation ΔΘ should be given.\nNext, we introduce our principled guideline for selecting the 𝛼 value through sensitivity analysis. To be more\nspecific, we perform sensitivity analysis on both the robustness and accuracy of the trained model. Here, we\nchoose the Microring resonator (MRR)-based computing system as an example demonstrating system. We use\nthe sensitivity 𝜆 to evaluate the trained model’s robustness. Small 𝜆 indicates high system robustness.\n!\"# !\"#\nThe sensitivity analysis result is shown in Figure R5, by gradually increasing the 𝛼 from 0 to 100, the model’s\nsensitivity first reduces and then increases. In contrast, the model’s accuracy remains stable when 𝛼 is smaller\nthan 1, and drops significantly when we further increase the value of 𝛼. This result indicates that small\nperturbations can effectively increase the model’s robustness while maintaining the accuracy, but larger\nperturbations (Large sharpness penalty) would make the training totally ineffective.\nBased on the sensitivity analysis result, we choose the optimum value of 𝛼 at 0.1, because of the highest model’s\nrobustness and accuracy are achieved simultaneously at this point.\nFigure R7. Sensitivity analysis\nRevision 6\nWe simplify the definition of hyperparameters in Equation 2-4 and give a detailed principled guideline of\nchoosing the hyperparameter in the manuscript (Line 154-155, Line 162-166) and Supplementary Note 1.1.\nComment 7\n6.Minor issues: In Figure 2(h), it would be more accurate to describe the performance shown as the ideal\nperformance rather than the theoretical limit. In Equation (4) of the main text, the numerator appears to be\nmissing the letter \"L\".\nResponse 7\nWe thank the reviewer for pointing out the inexact wording in our expressions and typo. We have revised the\nmanuscript and corrected the errors.\nRevision 7\n1. Revise ‘theoretical limit’ to ‘ideal performance’.\n2. Revise Equation (4) (Now Equation (3) in the main text).\nComment 8\n7.I strongly encourage the authors to consider releasing the code, even if only for the MZI-based\nimplementation, which would significantly enhance the impact and reproducibility of this work.\nResponse 8\nWe thank the reviewer for pointing our the issue regarding code availability. Therefore, we have summarized all\nthe simulation codes of different computing systems and store them in the GitHub repository.\nhttps://github.com/cuhkhuangslab/Sharpness-Aware-Training/tree/main\nReference\n[1] Zheng, Z., Duan, Z., Chen, H., Yang, R., Gao, S., Zhang, H., ... & Lin, X. (2023). Dual adaptive training of\nphotonic neural networks. Nature Machine Intelligence, 5(10), 1119-1129.\n[2] Wu, S., Wen, S., & Xue, H. (2024, August). Liquid-Cooled Heat Dissipation Technology for Co-Packaged\nOptics over 12.8 Tbps. In 2024 25th International Conference on Electronic Packaging Technology (ICEPT) (pp.\n1-4). IEEE.\n[3] Bandyopadhyay, S., Hamerly, R., & Englund, D. (2021). Hardware error correction for programmable\nphotonics. Optica, 8(10), 1247-1255.\n[4] Foret, P., Kleiner, A., Mobahi, H., & Neyshabur, B. (2020). Sharpness-aware minimization for efficiently\nimproving generalization. arXiv preprint arXiv:2010.01412.\nResponse To Reviewer #3\nComment 1\nThis manuscript introduces the use of Sharpness-Aware Minimization (SAM), a technique originally developed\nin the machine learning domain, to improve the robustness of physical neural networks. By simultaneously\nminimizing the loss and its sharpness, the authors aim to eliminate the need for in-situ training, thereby enabling\nthe network to tolerate environmental fluctuations and fabrication variance. The effectiveness of this approach\nis experimentally evaluated across three distinct hardware platforms.\nResponse 1\nWe thank the reviewer for the comments, which help us improve our manuscript.\nComment 2\nWhile the manuscript presents a technically sound implementation and provides extensive experimental\nvalidation, I have the following major concerns regarding its novelty and overall impact:\n1. Lack of originality in methodology: The training is entirely performed in silico, with only the deployment\ncarried out on physical platforms. The core idea—SAM—is directly borrowed from machine learning literature\nand is already well-established.\nResponse 2\nWe thank the reviewer for their comments. However, we respectfully disagree with the points raised regarding\nthe in silico training. The reviewer comments that “the training is entirely performed in silico, with only the\ndeployment carried out on physical platforms”. However, our training method, SAT, is applicable to both in-silico\nand in-situ training. In particular, the whole Subsection 2.3 (orignally Subsection 2.4) is dedicated to reporting\nthe results of applying SAT in an in-situ training setting, demonstrating its unique ability to enable transferable\nlearning, that is, parameters trained in one device can be successfully transferred to other devices. In contrast,\nexisting methods such as DAT [1] and PAT [2] lack this capability and are limited by device-specific training\nconstraints.\nRegarding the second comment “The core idea—SAM—is directly borrowed from machine learning literature\nand is already well-established.”, we would like to clarify that that drawing inspiration from machine learning\nshould not be the reason to diminish the research novelty.\nSAM was originally proposed in the machine learning community. Its core insight is the connection between\nthe sharpness of the loss landscape and a model’s generalization performance. The purpose is to address AI\nmodel generalization. Our work is inspired by this conceptual framework, however the originality of our\ncontribution is to establish a new connection—between the sharpness of the loss landscape and the robustness\nof physical systems. Guided by this new perspective, we propose SAT tailored specifically for physical\nsystems, which enables robustness to hardware imperfections, environmental disturbance, and transferrable\nlearning, for both in-silico and in-situ training. Such cross-disciplinary integration is the essence of\ninterdisciplinary innovation.\nIn addition to introducing a new conceptual framework, our work also makes a significant methodological\ncontribution. The original SAM cannot be directly applied to physical systems. SAM is designed to optimize\nweights and biases in AI models that are explicitly defined by precise mathematical functions. However, as\nwe point out in Section 2.1 of the Principle section, the mathematical models of PNNs are often imprecise\nor even unknown. To address this, SAT extends from SAM in two ways: (1) it accommodates cases where the\nunderlying model of the PNN is not explicitly known, and (2) it reformulates the standard weight optimization\ninto a control-parameter optimization problem, allowing direct training of the physical system, thus enabling\nin-situ training. These new developments make SAT highly generalizable and suitable for a wide range of\nphysical systems, whether or not a precise model is available, and support both in-silico and in-situ training\nscenarios.\nTo support the conceptual and methodological novelty and effectiveness of our approach, we have conducted\nexperiments and simulations on three different physical systems: microring circuits, free-space optics, and\nMZI networks. These systems were not chosen to simply repeat the same concept, but rather to demonstrate\nthe broad applicability and benefits of SAT in different contexts below:\n(1) When SAT is implemented in in-silico training setting, it addresses the modeling challenges in in-silico\ntraining by enabling high-performance training even when the digital model is imperfect;\n(2) When SAT is implemented in in-situ training setting, it facilitates transferable training, overcoming a key\nlimitation of in-situ training where trained parameters typically cannot be transferred to other systems due to\nhardware variability;\n(3) SAT enables high-performance inference when environmental disturbances and system noise presents after\ntraining regardless of using in-silico or in-situ training; and\n(4) it is generalizable across a wide range of physical systems, regardless of whether a precise mathematical\nmodel is available.\nWe realize that some key contributions of our work were not clearly visualized in the original manuscript. In\nresponse, we have significantly revised the manuscript to better highlight the novelty and significance of our\ncontributions by adding an additional table (Table 1) in the between Section 3 and 4, and additional\nparagraph at the end of Section 2.1.\nWe sincerely hope that the reviewer will reconsider the evaluation after reading the above clarifications and\nimprovements.\nTable R1. Comparison of training methods regarding training objectives, generality, and performance during the training,\ndeployment, and inference stages.\nRevision 2\nWe have significantly revised the manuscript to better highlight the novelty and significance of our contributions\nby adding an additional table (Table 1) in the between Section 3 and 4, and additional paragraph at the end of\nSection 2.1 (Line 190-201). We clarify the relationship between SAM and SAT as requested in the Abstract (Line\n21-25) and Manuscript (Line 91-95).\n“In the following section, we demonstrate SAT on three distinct PNNs. These demonstrations not only highlight\nthe broad applicability of SAT, but also illustrate its advantages over existing training methods from three\nperspectives: (1) When SAT is implemented in in-silico training setting, it closes the model-reality gap, enabling\naccurate backpropagation training under imprecise models (Section.\\ref{sec2.2}); (2) When SAT is implemented\nin in-situ training setting, SAT addresses the challenge of limited transferability caused by fabrication variances.\nParameters trained on one device using SAT can be reliably deployed to other devices without accuracy\ndegradation, even when hardware discrepancies exist (Section.\\ref{sec2.3}); (3) In both in-silico and in-situ\ntraining settings, PNNs trained by SAT can continuously operate accurately after training under perturbations\nwithout retraining (Section.\\ref{sec2.2}--Section.\\ref{sec2.4}); (4) SAT is broadly applicable to physical systems\nregardless of whether an accurate physical model is explicitly available (Section.\\ref{sec2.2}--\ref{sec2.4}). The\nadvantages of SAT over existing in-silico and in-situ training methods, including those designed to improve\nrobustness, such as NAT and optical pruning, are further demonstrated in the following sections, and summarized\nin Table.\\ref{tab:training_comparison}.”\n“Here, we address the challenges with both in-silico and in-situ training through Sharpness-Aware Training\n(SAT), \\textcolor{blue}{inspired by a machine learning technique Sharpness-Aware Minimization\n(SAM)~\\cite{foret2021iclr} that links loss landscape geometry with model generalization to enhance model\ngeneralization. In this work, we establish a new link between loss landscape geometry with robustness of physical\nsystems, and then leverage this connection to tackle above challenges in training physical systems.”\n“Furthermore, SAT develops methodologies for automatically locating flat minima in physical systems and\nenables direct training in physical systems.”\nComment 3\n2. Limited conceptual advancement: While the experimental implementation differs from the authors’ earlier\nwork, the underlying idea and overall conceptual direction remain highly similar. This raises concerns about a\nlack of novelty in terms of scientific framing and motivation. Such conceptual continuity may be appropriate\nfor an incremental study, but may not meet the threshold for a high-impact general-interest journal.\nResponse 3\nWe thank the reviewer for raising the question regarding the comparison with our prior work.\nIn our prior pruning work [3], we also aimed to train the system to a robust region. However, the method we\nemployed was limited to relatively simple systems where such robust regions could be easily identified. For\ninstance, in microring circuits, where each component operates almost independently, the robust region\ncorresponds to the flat region of each individual microring's transfer function.\nHowever, in more general physical systems, robust regions are not so easily defined. For more complex systems\nsuch as MZI networks (Section 2.3), although their behavior can be approximated by mathematical models, the\nrobust region of the entire network is neither explicit nor easily identifiable. In even more complex scenarios,\nsuch as free-space optical systems (Section 2.4), even the system model cannot be explicitly formulated at all,\nmaking it practically impossible to identify robust regions using the approach from our previous work. Moreover,\nthrough this work, we found that even in simple systems, the robust region of individual components does not\nnecessarily overlap with the robust region of the system as a whole.\nIn contrast, the core contribution of SAT is that it is broadly applicable to all physical systems: SAT can\nautomatically identify robust minima without requiring explicit knowledge of the underlying physical dynamics.\nOur experimental and simulation results on different systems support this generality claim. Furthermore, as\nillustrated in Figure 1(d) in the main text, SAT not only enhances the stability of individual components but also\nidentifies configurations that improve the robustness of the system as a whole, thus leading to superior\nperformance even when operating on the same hardware.\nFigure R8 shows the improved performance of SAT on microring circuits compared to our previous optical\npruning method. As mentioned earlier, the key reason is that SAT not only enhances the stability of individual\ncomponents but also identifies configurations that improve the robustness of the system as a whole. In contrast,\nthe optical pruning approach focuses solely on stabilizing individual components.\nFigure R8. Inference accuracy on microring circuits and confusion matrix when removing the thermoelectric\ncontroller (TEC).\nRevision 3\nWe revise our manuscript to highlight the difference of our proposed method with our prior work (Line 174-181).\n“The loss minimization process highlights the difference with our prior work optical pruning [27]. Optical pruning\nfocuses solely on improving the stability of individual devices inside the PNNs. In contrast, SAT not only\nimproves the stability of individual components but also enhances the overall system's stability by ensuring that\nthe entire system's loss function remains minimized in response to weight perturbations. Furthermore, optical\npruning is limited to relatively simple systems where such stable regions could be easily identified. For instance,\nin microring circuits, where each component operates almost independently, the robust region corresponds to the\nflat region of each individual microring's transfer function. In contrast, SAT is broadly applicable to different\nphysical systems because it can automatically identify robust minima without requiring explicit knowledge of the\nunderlying physical dynamics.”\nComment 4\n3. Absence of full neural network functionality: Despite being framed as work on “physical neural networks,”\nthe experimental demonstrations are limited to simple linear functions. No physical implementation of\nnonlinear activation or full inference capability is presented.\nResponse 4\nWe thank the reviewer for raising the question of the full neural network functionality.\nThe principle described in Section 2.1 demonstrates that SAT is universally applicable to systems with or without\nnonlinear functions. SAT identifies flat minima based on gradients, which can be approximated from an imperfect\nmodel or measured using finite-difference methods. Therefore, even for systems with nonlinearities, the gradient\napproximation process is the same as for systems with linear functions.\nTo further address the reviewer’s concern, we have included an additional simulation study on a deep diffractive\nneural network (D2NN). Here we follow the setting in [1] and include nonlinear functions between the linear\ndiffractive neural network, as illustrated in Figure R9. The simulated system consists of two programmable linear\nlayers and two nonlinear layers. The first nonlinear layer is implemented using a saturable absorber, which\nintroduces an intensity-dependent transmission while preserving the optical phase [4]. And the second nonlinear\ncomponent is the photodetector (PD) array, which introduces a quadratic nonlinearity during the optical-to-\nelectrical conversion process. The input image is encoded via a spatial light modulator (SLM), after which light\npropagates through the layered structure (linear → nonlinear → linear → nonlinear) and is ultimately captured\nby the PD array. We apply this system to perform handwritten digits classification, where the output class is\ndetermined by the index of the maximum PD response.\nWe train the physical network using both standard backpropagation (BP) and our SAT. During training, we\nemploy the finite difference method to approximate the gradient of the loss with respect to the rotation angle\n(which acts as trainable control parameters), and we intentionally introduce perturbations in rotation to evaluate\nrobustness.\nAs shown in Figure R10, our method maintains high classification accuracy under rotation angle variations.\nThese results clearly demonstrate that SAT remains effective when both linear and nonlinear operations are\nphysically implemented. The detailed simulation parameters are depicted in Table.R1. The simulation code will\nbe provided in the open-source code file [5].\nTable.R2 Simulation hyperparameters\nParameter name Value\nBatch size 200\nLearning rate 0.002\nLinear layers size 128*128\nLinear layer pixel resolution 20 𝜇m\nLaser wavelength 532 nm\nFigure R9. Schematic diagram of the simulated Deep diffraction neural network (D2NN).\nFigure R10. Inference accuracy with rotation angle change from 0.0° to 1.0°.\nRevision 4\nWe revise the manuscript to add a discussion on our method’s compatibility with nonlinear systems (Line 463-\n474), and incorporate the simulation results into the Supplementary Note 6.4.\nComment 5\n4. Presentation issues: The manuscript contains several typographical issues and instances of repeated text. A\ncareful proofreading is necessary to ensure clarity and professionalism in presentation.\nResponse 5\nWe thank the reviewer for pointing out the presentation issues. We have done a careful proofreading and revised\nthe typographical issues and unintended redundancies. However, certain repetitions were intentionally retained\nas they serve as reminders for readers, given the paper’s length and technical complexity.\nComment 6\nThat said, I appreciate the authors’ effort in generating high-quality figures and conducting thorough, cross-\nplatform experimental validation. These aspects reflect a strong technical execution.\nIn summary, while the paper demonstrates a useful application of an existing training technique to physical\nsystems, the work currently does not meet Nature Communications' expectations for innovation and conceptual\nsignificance. Substantial improvements in both novelty and experimental depth would be required for this work\nto be considered for publication.\nResponse 6\nWe sincerely thank the reviewer’s appreciation in our large amount of effort in experiments and simulations to\nproof the effectiveness and advancement of our proposed method.\nHowever, we sincerely hope that, after re-reviewing the response letter (particularly the difference and connection\nbetween SAM and SAT in Response 1 and reviewing the in-situ training section 2.3) and the revised manuscript,\nthe reviewer may reconsider the assessment of the novelty and constribution of our work. To the best of our\nknowledge, we are the first to establish a new conceptual framework that links the sharpness of the loss landscape\nto the robustness of physical systems. This framework leads to a new, highly effective, and universally applicable\napproach to training physical neural networks, effectively addressing the fundamental challenges inherent in both\nin-silico and in-situ training methods, as outlined below:\n1. SAT allows PNNs to achieve reliability, efficiency, and precision even when facing all kinds of errors crossing\ndesign, manufacturing, and deployment stages.\n• Modeling Errors: SAT does not require precise modeling of the physical system or consideration of side\neffects or system imperfections, yet it still outperforms methods that do.\n• Fabrication Variances: The SAT-trained integrated PNNs are highly resilient to fabrication variances\nand errors, addressing a critical challenge in scaling functional large-scale PNNs.\n• Ambient Fluctuations: PNNs trained with SAT maintain high accuracy even under temperature\nfluctuations, eliminating the need for power-hungry temperature control.\n• Deployment Errors: PNNs trained with SAT show high resilience to deployment errors, such as\nalignment errors commonly seen in free-space PNNs including the “x” and “y” axes, rotation, and scaling.\n2. SAT outperforms state-of-the-art training methods in both in-silico and in-situ training regarding accuracy,\nefficiency, and robustness.\nFirst, SAT achieves the highest accuracy while significantly reducing training time compared to the most\neffective training methods currently based on backpropagation (Nature 601(7894), 2022; Science 380(6643),\n2023; Nature Machine Intelligence 5(10), 2023) [6-8]. Second, SAT demonstrates robustness and tolerance to\nsystem imperfections, particularly those encountered after training during deployment, such as temperature\nfluctuations and alignment errors. This capability positions SAT as a leading solution in the field.\n3. SAT facilitates transferable learning across different devices, a feature that has not been achieved with other\ntraining methods.\nIn modern AI hardware, transferable learning is essential and commonly practiced for digital computers.\nTypically, training occurs on high-end computers or cloud-based clusters, with the trained parameters then\ndeployed across multiple edge devices for inference.\nHowever, in physical systems, whether through in-silico or in-situ training, transferable learning has remained\nelusive. This is because physical systems, even with identical designs, can exhibit different behaviors due to\nfabrication variances and environmental noise. In contrast, SAT enables transferable training across diverse\ndevices regardless of these imperfections––training is done ONCE in one PNN, while trained parameters can be\nreliably deployed to numerous edge PNNs for inference.\n4. SAT is a universally applicable method for various PNNs regardless of whether their models are explicitly\nknown or unknown, linear or nonlinear. It is versatile across diverse applications, including image classification,\ncompression, reconstruction and generation.\nIn summary, SAT offers a practical, effective, and computationally efficient solution, not only for training PNNs,\nbut also for deploying them in real-world environments and applications. We respectfully hope that the reviewer\nwill find the clarifications and improvements satisfactory, and reconsider the overall assessment.\nReference\n[1] Zheng, Z., Duan, Z., Chen, H., Yang, R., Gao, S., Zhang, H., ... & Lin, X. (2023). Dual adaptive training of\nphotonic neural networks. Nature Machine Intelligence, 5(10), 1119-1129.\n[2] Wright, L. G., Onodera, T., Stein, M. M., Wang, T., Schachter, D. T., Hu, Z., & McMahon, P. L. (2022). Deep\nphysical neural networks trained with backpropagation. Nature, 601(7894), 549-555.\n[3] Xu, T., Zhang, W., Zhang, J., Luo, Z., Xiao, Q., Wang, B., ... & Huang, C. (2024). Control-free and efficient\nintegrated photonic neural networks via hardware-aware training and pruning. Optica, 11(8), 1039-1049.\n[4] Lin, X., Rivenson, Y., Yardimci, N. T., Veli, M., Luo, Y., Jarrahi, M., & Ozcan, A. (2018). All-optical machine\nlearning using diffractive deep neural networks. Science, 361(6406), 1004-1008.\n[5] https://github.com/cuhkhuangslab/Sharpness-Aware-Training/tree/main\n[6] Wright, L. G., Onodera, T., Stein, M. M., Wang, T., Schachter, D. T., Hu, Z., & McMahon, P. L. (2022). Deep\nphysical neural networks trained with backpropagation. Nature, 601(7894), 549-555.\n[7] Pai, S., Sun, Z., Hughes, T. W., Park, T., Bartlett, B., Williamson, I. A., ... & Miller, D. A. (2023).\nExperimentally realized in situ backpropagation for deep learning in photonic neural networks. Science,\n380(6643), 398-404.\n[8] Zheng, Z., Duan, Z., Chen, H., Yang, R., Gao, S., Zhang, H., ... & Lin, X. (2023). Dual adaptive training of\nphotonic neural networks. Nature Machine Intelligence, 5(10), 1119-1129."}, {"round": 2, "reviewer_comments": {"reviewer_1": "I have reviewed the responses the authors gave to my questions and suggestions, and the revisions they have made to the manuscript as a result. They have satisfactorily responded in all cases, so I recommend publication."}, "author_response": ""}]

Decision letter