Paper: LUMEN-PRO: Automating Multi-Task Learning on Optical Neural Networks with Weight Sharing and Physical Rotation (ICLR-2024)

Reviewer comment (Theory):
The method depends on the rotation of the physical layer. However, the physical layer has at most four directions. Therefore, the method only supports most four tasks.

Author response:
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.

Reviewer comment (Novelty):
The method is mainly derived from the AutoMTL [1] method.

Author response:
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.

Reviewer comment (Presentation):
The presentation is not clear enough. Some details are not included in the paper. This can be seen in the questions.

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

Reviewer comment (Presentation):
The Figures are not annotated; thus, it is difficult to understand the method directly by looking at them.

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

Reviewer comment (Presentation):
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.

Author response:
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.

Reviewer comment (Evaluation):
Why the LUMEN-PRO’s performance can exceed the single task in Figure 5? Normally, it works best to use a separate model for each task.

Author response:
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.

Reviewer comment (Novelty):
The application of designing multi-task learning method for the DONN is too narrow. DONN is just one type of optical neural network and there is no evidence that this approach generalizes to other physical neural networks. The method may not have much practical usages in real life.

Author response:
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.

Reviewer comment (Experiments):
The experiments conducted seems only from a mathematical perspective. If we put this solution to produce physical systems, will there be accuracy degeneration due to imperfect fabrication? And is the proposed rotation-based sharing method practical in real fabricated system? The authors did not address these issues.

Author response:
***Q2: Will there be accuracy degeneration due to imperfect fabrication? Is the proposed method practical in real fabricated system?***

R2: There are several approaches to deal with the accuracy degradation after physical deployment: (1) As shown in [2], we can use adaptive training to fine-tune the model regarding the optical devices. (2) We can integrate physics-aware codesign algorithms in [1] in our numerical emulation and training process, where the algorithm takes the imperfect fabrications into the training process for better hardware-software correlations. [1] provide the physical prototype for physical deployment, which shows the impressive correlation between numerical emulation and physical measurements. 

As for the rotation-based sharing method, it is practical after the fabrication system. As shown in Figure 2 in our paper, the diffractive layers (i.e., the weight matrix) are fabricated as 3D printed phase masks with square shape. To perform the rotation method, for example, in the system in Figure 2 of [3], the fabricated layers are inserted in slots. We can take the phase mask out of the slot, rotate an angle, and insert the layer back to the slot (these movements can be accomplished by a fast robotic arm designed), which is straightforward and realizes weights sharing at no extra fabrication cost for multiple tasks in DONN systems. 


[1] Y. Li, R. Chen, W. Gao, and C. Yu. Physics-aware differentiable discrete codesign for diffractive optical neural networks. In Proceedings of the 41st IEEE/ACM International Conference on Computer-Aided Design, pp. 1–9, 2022.

[2] T. Zhou, X. Lin, J. Wu, Y. Chen, H. Xie, Y. Li, J. Fan, H. Wu, L. Fang, and Q. Dai, “Large-scale neuromorphic optoelectronic computing with a reconfigurable diffractive processing unit,” Nature Photonics, vol. 15, no. 5, pp. 367–373, 2021.

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

Reviewer comment (Writing):
Due to my lack of experience with this field, I do not understand a lot of technical details in this work. I believe the authors can improve on the explanation of the key concept to make it easy to understand. For example, the Figure 3 is really confusing. What does a node mean? What does the numbers in the node blocks mean? Are they network weights? There are a bunch of switches on the figure. What is the functionality of these switches and how do they work? Another key concept is the rotation-based layer sharing. What exactly does rotation mean in this scenario? How does such rotation facilitate weight sharing?

Author response:
***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.

Reviewer comment (Evaluation):
Table 3 is kind of confusing. It seems to contain both ASIC-based solution and physical neural networks. How do you measure the throughput of an optical neural network? The proposed framework has very high throughput but is it really possible in a real system? Since you need to switch the input image physically at such a fast rate. And what does “Accuracy” mean in this table? Is it just the testing accuracy on MNIST?

Author response:
***Q4: Confusing about Table 3.***

R4: Thank you for your questions. We will address your questions separately.

- ***How do you measure the throughput of an optical neural network?***

Table 3 showcases the theoretical throughput of the free-space DONN. The presented data is comparable to the theoretical throughput figures shared by HolyLight paper, which encompass the performance of ISAAC and DaDianNao (simulated results with theoretical bounds). We follow the same setting. The distinct characteristic of the free-space DONN lies in its ability to utilize all-optical processes. Leveraging the inherent properties of light, including natural light diffraction and swift light signal phase modulation at the speed of light, the free-space DONN exhibits the (theoretically) potential to achieve a throughput of $10^5$. 

- ***The high-throughput framework's practical feasibility raises concerns due to the physical input switching required.***

The "rotation" action does not occur during inference of each individual task. Instead, it takes place after completing 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 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. It's important to note that the discussion of "cost" and "time" in this context is not reflected in Table 3. 

- ***What does "Accuracy’’ mean in Table 3?***

As mentioned in the "Energy Efficiency Comparison" paragraph under *Section 5.2.2 Accuracy-Cost Comparison*, Table 3 presents the accuracy results specifically for MNIST. For LUMEN-PRO, this implies that we have already assembled the physical diffractive layers and related hardware components of the DONN. The measurements provided in Table 3 focus solely on the throughput and power requirements associated with MNIST.