LUMEN-PRO: Automating Multi-Task Learning on Optical Neural Networks with Weight Sharing and Physical Rotation

Your time in evaluating our work is sincerely appreciated. We would like to address your concerns and questions as below.

Q1: What is unique for DONN? Is this method the general method that can be applied to other CNN or Transformer models?

R1: Thanks for your question and concern regarding the practical usage of our designed algorithm. We would address your questions in two folds.

• Unique for DONN.

In addition to the AutoMTL method, we also want to emphasize the significance of the rotation of the physical layer for the DONN system. While AutoMTL enables the discovery of an efficient multi-task framework, the rotation of the physical layer further facilitates a reduction in memory storage, allowing it to match that of a single-task model. This achievement establishes a memory lower bound, making the rotation of the physical layer an integral and noteworthy aspect.

When doing inference on the physical optical device, the "rotation" action takes place after completing the inference for one task, facilitating the system's transition to conducting inference on another task. Due to the rotational capability of the physical system, there is no need to disassemble and reassemble the entire system or replace all the diffractive layers when switching tasks. Rather, we can simply rotate the corresponding layer or layers based on the emulated MTL model. This rotational approach offers cost savings by eliminating the need for additional printing materials and reduces the time required for hardware reconfiguration.

• Generalization of LUMEN-PRO.

The AutoMTL method is applicable to any CNN and Transformer model with ease. The physics-aware weight sharing rotation algorithm in our LUMEN-PRO is specifically designed for DONN systems due to the non-reconfigurable feature of the optical hardware system. However, we believe that with certain adjustments, this algorithm can also be applied to any model that has a square-shaped feature map.

Q2: You mentioned that the rotation mechanism has a physical meaning, what is that? Why is the rotation different from spatial shifts in CNNs?

R2: “Physical meaning” meaning that when switching to different tasks, we can simply rotate the corresponding layer or layers based on the searched multi-task configurations. There is no need to disassemble and reassemble the entire system or replace all the diffractive layers when switching. In practical deployment of DONN systems, the trained weights in the model (diffractive layers) are fabricated and deployed for all-optical inference, which is non-reconfigurable. Thus, to further reduce the system cost for multiple tasks learning with the deployed DONN system, we proposed the physical rotation mechanism for weights sharing where the cost efficiency is improved by 4x when dealing with 4 tasks compared with single task systems (Table 2). More importantly, the rotation mechanism works for symmetric but non-configurable physical systems.

The rotation of diffractive layers and spatial shifts in CNNs are two distinct concepts used in different contexts. Spatial shifts of CNNs refer to the translation or shifting of input data within the network. It involves moving or shifting the positions of input images or feature maps in the horizontal and vertical directions. Spatial shifts are typically used to achieve spatial invariance, allowing the network to recognize objects or patterns regardless of their specific position in the input image. In summary, the rotation of diffractive layers pertains to modifying the physical properties of the diffractive optics itself, whereas spatial shifts in CNNs involve shifting or translating input data within the network to achieve spatial invariance and robustness to translation.

We sincerely appreciate the time you've taken to evaluate our work. We would like to address your concerns and questions as below.

Q1: The method only supports most four tasks.

R1: While it is true that the physical layer typically has four directions, the method does not necessarily limit the number of tasks to only four. Firstly, not all tasks are necessarily to choose the task-specific option simultaneously at a specific layer. Furthermore, during training, we allocate more weights to the "sharing" option compared to the "task-specific" option to encourage sharing operators across tasks. Secondly, even if all tasks select the task-specific option at the same layer, resulting in some of the tasks sharing the same rotation angle and therefore having identical weights at that layer, these tasks would still exhibit distinct operator selection or rotation angles in other layers, which makes them still have different model structures and weights. The rotation of the physical layer used in our LUMEN-PRO enables efficient sharing of weights among tasks, allowing for a larger number of tasks to be supported and also greatly reduces the memory overhead.

Q2: The method is mainly derived from the AutoMTL method.

R2: Thank you for highlighting this important aspect. We emphasize the significance of the rotation of the physical layer in the DONN system. The rotation of the physical layer further facilitates a reduction in memory storage, allowing it to match that of a single-task model. This achievement establishes a memory lower bound, making the rotation of the physical layer an integral and noteworthy aspect.

When performing inference on the physical optical device, the "rotation" action takes place after completing the inference for one task, facilitating the system's transition to conducting inference on another task. Due to the rotational capability of the physical system, there is no need to disassemble and reassemble the entire system or replace all the diffractive layers when switching tasks. We can simply rotate the corresponding layer or layers based on the searched multi-task configurations. This rotational approach offers cost savings by eliminating the need for additional printing materials and reduces the time required for hardware reconfiguration.

Q3: The Figures are not annotated.

R3: We greatly appreciate your valuable suggestions. We will incorporate your comments into the paper and include them in the final version.

Q4: What is the meaning of the LUMEN-PRO in Figure 5? As I understand, this network can only make the inference function different under different tasks by rotating the layers.

R4: Yes, your understanding is correct. This network's ability to alter the inference function for distinct tasks is achieved by rotating the layers. In Figure 5, we compare LUMEN-PRO w/o rotation and LUMEN-PRO w/ rotation with other state-of-the-art (sota) multi-task learning (MTL) models. LUMEN-PRO w/o rotation is an MTL DONN obtained using the AutoMTL algorithm, with each task-specific layer having its specific weight. LUMEN-PRO w/ rotation refers to an MTL DONN obtained by combining AutoMTL with rotation design. This model achieves the same memory footprint as a single-task model since all task-specific layers are derived by physically rotating the shared layers. The purpose of comparing the accuracy between these models is to conduct an ablation study highlighting two points: (1) the effectiveness of the MTL DONN obtained through AutoMTL, which outperforms both single-task models and other SoTA MTL models in terms of accuracy, and (2) the effectiveness of the added rotation algorithm in LUMEN-PRO. This algorithm not only surpasses the accuracy of both single-task models and other sota MTL models but also achieves a memory lower bound with minimal accuracy sacrifice compared to not incorporating the rotation algorithm.

Q5: Why the LUMEN-PRO’s performance can exceed the single task in Figure 5?

R5: Multi-task learning models can outperform single-task models due to several reasons. Firstly, the MTL model can improve generalization by learning shared representations and identifying common patterns across multiple related tasks. Secondly, multi-task learning enables the model to leverage limited data more efficiently by using information from multiple tasks, resulting in better performance overall. Thirdly, it acts as a form of regularization, preventing overfitting and enhancing the model's ability to generalize. Lastly, by sharing weights, the memory footprint of multi-task learning model is significantly reduced for all tasks compared to having a model for each task. The effectiveness of multi-task learning models depends on task relatedness.

Q3: Confusing about Figure 3.

R3: Thank you for your questions. We will address your questions in two parts.

• What does a node mean? What does the numbers in the node blocks mean? What is the functionality of these switches?

The input DONN backbone model follows a single-task architecture, with numbers representing the weights in each layer. In Figure 3, we utilize a three-layer DONN as an example, where each layer is represented by a 2x2 matrix. To represent the DONN backbone model as a computation graph, we employ nodes to symbolize the layers. In this context, the notation $node_i$ is used to refer to layer $i$ within the backbone model.

In order to transition the single-task backbone into a multi-task model, we enhance each node as a computation unit ($CU$). If there are $N$ tasks in the design, each $CU$ would consist of $N$ blocks, with each block comprising three operators: (1) the basic shared operator ($bsc$), (2) a task-specific copy of the basic shared operator ($spc^t$), or (3) a skip operator ($skp^t$). The basic shared operator ($bsc$) is common and shared across all $N$ blocks (tasks). All these CUs form our multi-task supermodel.

Our objective is to learn the policy for each task, which involves determining the suitable operator to select for each task in every layer within the CUs. The policy is denoted by "switches" (P), such as $P_2^{t_1}$, representing the policy for task 1 in layer 2. After the entire training process, $P_2^{t_1}$ can take on one of three possibilities: (1,0,0) represents the selection of the shared operator (share the same weight with other tasks) in layer 2 for task 1; (0,1,0) represents the selection of the task-specific operator (with its own specific weights in that layer) for task 1 in layer 2, and (0,0,1) indicates the decision to skip that layer for task 1.

• What does rotation mean? How does rotation facilitate weight sharing?

We present the detail of "rotation" in Section 4.2 Rotation Algorithm. In this context, "rotation" refers to obtaining the task-specific layer through a 90-degree / 180-degree / 270-degree rotation of the shared layer. For example, in Figure 4 (a), if we denote layer 1 of task 1 as $L_1^1$, the weights in the first layer of task 2 ($L_2^1$) and task 3 ($L_3^1$) are derived from a 90-degree and 180-degree clockwise rotation, respectively, of the weights in $L_1^1$. In the second layer, both task 1 and task 2 select the shared operator, resulting in them sharing the same weight in this layer ($L_1^2 = L_2^2$), while task 3 chooses the task-specific operator. As a result, the weights in the second layer of task 3 are obtained through a 90-degree rotation of $L_1^2$. In model training, the degrees of rotation are actually randomly assigned.

By leveraging "rotation," it is possible to construct a multi-task learning model with equivalent parameter size and memory storage as a single-task model. This is achieved by having all tasks within each layer either share the same weights or acquire their weights through rotation from the weights of other tasks. Moreover, when deploying the emulated model onto a physical optical device, such as a 5-layer multi-task model supporting 4 tasks, only 5 phase masks need to be fabricated. By rotating some of these masks, the inference process can support all four tasks efficiently.

We sincerely appreciate the time you've taken to evaluate our work. We would like to address your concerns and questions as below.

Q1: The application of designing multi-task learning method is too narrow and no practical usages.

R1: To address your concern regarding the practical application of our designed algorithm, let us explain in two folds.

• Practical usages of multi-task learning on other physical neural networks.

Firstly, the proposed multi-task learning method has the potential to generalize to other physics-based neural networks, which utilizes nature physics as neural operators and are physically manufactured. For example, Figure 2 in [1] shows a physics-based neural network using broadband optical second-harmonic generation (SHG), and a mechanical version. The proposed methods can be extended beyond the targeted tasks (image classification), as long as the shared physical implementation can be realized, e.g., metasurface based imaging processing unit [2]. Note that in these systems (and ours), once the neural networks are fabricated, the rotation-based weight sharing will be able to share physical “weight” parameters without re-fabricate the system. For example, the optical and mechanical physics neural networks demonstrated in [1] can leverage completely our proposed methods in order to share the physically fabricated layers. Causally, this is a very similar idea to “Lenticular printing” in real life. At this point, we kindly request the reviewer to kindly consider the contributions of this paper fairly from the physical neural network point of view, instead of comparing it to conventional digital neural networks.

Specifically for DONN, the rotation mechanism enables a physics-aware weights sharing. In practical deployment, the trained weights in the diffractive layers are fabricated and deployed for all-optical inference, which is non-reconfigurable. To further reduce the system cost for multiple tasks learning, we propose the physical rotation mechanism for weights sharing where the cost efficiency is improved by 4x when dealing with 4 tasks compared with single task systems (Table 2 in our paper). More importantly, the rotation mechanism works for symmetric but non-configurable physical systems.

• Practical usages of DONN.

Secondly, Free-space Diffractive Optical Neural Networks have practical applications across various domains [3][4]. They excel in high-speed computing tasks, enabling rapid processing for real-time image and video analytics, object recognition, and data classification. In optical image processing, DONNs prove valuable for applications like image denoising, segmentation, and reconstruction[5][6]. They also excel in optical pattern recognition tasks, facilitating real-time recognition of complex patterns, document classification, biometric identification, surveillance systems, and optical character recognition[7][8]. DONNs contribute to neuromorphic computing systems, enabling machine learning, cognitive processing, and brain-inspired computing. Furthermore, they offer advantages in optical signal processing and sensing, enhancing Fourier analysis, beam shaping, spectral estimation, and enabling improved optical sensing capabilities. With their unique optical properties and parallel computing capabilities, DONNs present promising opportunities in the fields of research, industry, and technology. In this case, further exploration of DONNs applications, including multi-task learning, holds significant promise for addressing real-world challenges and achieving breakthroughs in various fields.

[1] Wright, Logan G., Tatsuhiro Onodera, Martin M. Stein, Tianyu Wang, Darren T. Schachter, Zoey Hu, and Peter L. McMahon. "Deep physical neural networks trained with backpropagation." Nature 601, no. 7894 (2022): 549-555.

[2] Li, L., Zhao, H., Liu, C., Li, L., & Cui, T. J. (2022). Intelligent metasurfaces: control, communication and computing. Elight, 2(1), 7.

[3] Lin, X., Rivenson, Y., Yardimci, N. T., Veli, M., Luo, Y., Jarrahi, M., & Ozcan, A. (2018). All-optical machine learning using diffractive deep neural networks. Science, 361(6406), 1004-1008.

[4] Sui, X., Wu, Q., Liu, J., Chen, Q., & Gu, G. (2020). A review of optical neural networks. IEEE Access, 8, 70773-70783.

[5] Luo, Y., et.al. (2019). Design of task-specific optical systems using broadband diffractive neural networks. Light: Science & Applications, 8(1), 112.

[6] Mengu, D., Veli, M., Rivenson, Y., & Ozcan, A. (2022). Classification and reconstruction of spatially overlapping phase images using diffractive optical networks. Scientific Reports, 12(1), 8446.

[7] Liu, J., et.al. (2021). Research progress in optical neural networks: theory, applications and developments. PhotoniX, 2(1), 1-39.

[8] Chen, R., Tang, Y., Ma, J., & Gao, W. (2023). Scientific Computing with Diffractive Optical Neural Networks. arXiv preprint arXiv:2302.10905.