Paper: Cooperative Hardware-Prompt Learning for Snapshot Compressive Imaging (ICLR-2024)

Reviewer comment (Novelty):
The major contribution of the paper is not well-justified. i.e., the crucial need for the proposed method (as opposed to training from scratch for every new measurement model) is not well-supported. For example, the reviewer still finds it very convenient to train a model for every new set of measurement models for a new organization of interest from a practical perspective. All it takes is a few hours (days) of training for the new set, which also automatically addresses the other concern raised by the authors regarding privacy constraints, in that each organization has access to its own data and device sets.

Author response:
Thanks for this useful question! 

The conventional setting where each organization employs its local data and device to perform independent training has practical limitations for the current SCI research study and prevents its further real-world applications due to the following reasons.
* Data-starving challenge of the client. Some clients may have limited training data (e.g., only several scenes) that the training cannot converge.
* Efficiency concern. We find training a local model from scratch takes $3.54$ days ($10.62/3$ as shown in Table 3) on our platform. It only takes less than $1$ day ($2.86/3$ in Table 3) to adapt the model to a new device/client using FedHP.  We have put this illustration into the manuscript. The training time can be even longer if the reconstruction backbone becomes larger. 

In summary, considering the data-starving nature of the client, it may be impractical for a new client to train a model. Besides, there is a training efficiency concern when there is enough data for the client. Plus, the reconstruction model trained with a single well-calibrated hardware instance is hard to adapt to new hardware. This work proposed FedHP to practically enable the deployment of reconstruction models on new clients.

Reviewer comment (Experiments):
The comparisons are not fair in Fig. 1. Clearly, joint training should serve as an upper bound on the performance when the test set contains the same measurement models as training.

Author response:
We appreciate the reviewer’s valuable insight in improving this work!

In general, joint training serves as an upper bound for federated learning when the test set contains the same measurement models as training. 

However, in Fig. 1, the test set consists of random measurement models that are unseen for training. Specifically, all results (1 to 5) are evaluated by sampling from the same mask pool to test the trained models. Note that this mask pool contains masks from three different distributions, i.e., $P_1$, $P_2$, $P_3$ as we plotted in Fig. 1. Since we always sample non-overlapped masks, the masks used for testing are generally unseen to the models. This is quite challenging and can well simulate a real-world scenario.  We have added more illustrations about the settings to the caption of Fig. 1.
In this case, our experiments find that simply combining all data from different hardware to jointly train does not work well (as shown in Fig. 1, over $0.6$dB lower than FedHP and $0.3$dB lower than FedAvg).  

We appreciate the reviewer's help in distinguishing this work from general federated learning endeavors!

Reviewer comment (Evaluation):
The results (especially the quantitative ones) do not yield the conclusion that FedAVG is outperformed by FedHP. The major advantage of FedHP seems to be its 4x more efficient training time compared to FedAVG. While improving efficiency is valuable from a practical point of view, the paper isn’t oriented around this conclusion; instead, it emphasizes the value of prompt networks and FedHP in the form of accuracy and robustness gains, whereas FedAVG enjoys those traits too.

Author response:
We thank the reviewer bAjN for helping us summarize this contribution! 

We find that FedAvg serves as a very strong baseline by even working better than recent  FL methods, such as SCAFFOLD and FedProx. Thus, it is non-trivial to achieve further performance boost. By comparison, FedHP can bring a consistent performance boost. For example, $+0.14$dB in Table 1, $+0.35$dB in Table 2, and $+0.27$dB in Table 4 (a) with more clients. 

Besides, by jointly introducing the hardware prompt network and adapter, the proposed method also benefits from the efficiency advantage by only requiring adapting the pre-trained models in the system to the new clients. We have also emphasized this point by mentioning the advantage of efficiency in the revised manuscript.

Reviewer comment (Presentation):
On Tab. 2, FedHP is highlighted as the best-performing method in terms of SSIM (0.8481), whereas FedAVG should be highlighted (0.8496).

Author response:
Thanks for the useful comments! We modified the annotation typos in Table 1  in our manuscript.

Reviewer comment (Theory):
The idea of the prompt network is interesting; however, there is a question of whether that network induces another source of instability to the overall model. Specifically, what guarantees that the prompt network doesn’t do a terrible alignment for measurement models deviating from the training distribution?

Author response:
Thanks for the valuable insight! 

We provide an ablation study in Table 3 to discuss the effect of prompt networks. FedHP experiences large performance degradation without using a prompt network. One thing to note is that the prompt network is learned on local but will be aggregated to the server, and then distributed to clients at each global round. This makes sure that the prompt network does not deviate from the global learning objective.

Reviewer comment (Evaluation):
What is the source of inconsistency between Tab. 1 and Fig. 3? Fig. 3 indicates a higher PSNR difference than 0.14 dB between FedAVG and FedHP. It is understandable to argue that quantitative metrics such as PSNR or SSIM do not perfectly capture the true quality, but the visual difference in Fig. 3 is too large not to be captured by those metrics.

Author response:
The difference of $0.14$dB is an averaged result on ten testing scenes as provided in Table 1. In Fig.3, we select one example of Scene $7$ to visually compare between different methods. This should correspond to a PSNR gap of $0.43$dB and $0.0073$ in SSIM. Besides, the visual quality in different regions can vary.  In the supplementary material, we provide more visual comparisons of different scenes.

Reviewer comment (Experiments):
Why isn’t the deep unfolding network included in all the results and only reported as a short paragraph at the end?

Author response:
Thanks for the insightful questions!

The deep learning-based reconstruction methods and model-based methods (e.g., deep unfolding) methods demonstrate promising performance for the SCI study. This work mainly focuses on deep learning-based methods. 

Besides the performance comparison in main tables, we also would like to compare with a state-of-the-art deep unfolding method [3]. As shown in Table 4b, FedHP also brings a notable performance boost ($+0.28$dB/$0.0038$ in PSNR/SSIM) with a much smaller model size.

Reviewer comment (Evaluation):
Any intuition on why FedAVG is so much slower to train than FedHP?

Author response:
Thanks for the useful question!

There are two designs that make FedHP more efficient than FedAvg. 

Firstly, FedHP does not need to train every local client model from scratch, as long as there is a pre-trained model, we can easily adapt it to new clients under FedHP. By comparison, for any new client, FedAvg is required to train the model from scratch, this can be computationally cumbersome.

The other reason is that the proposed method has a much lower communication burden. FedHP only needs to aggregate and distribute a lightweight hardware prompt network and adaptors. By comparison, FedAvg has to do full model training and communication to adapt to different masks, otherwise, the performance will be significantly degraded (e.g., $<20$dB). In contrast, we introduce a hardware prompt to address this issue in a smart way.

Reviewer comment (Experiments):
It remains unclear if the proposed method can adopt a larger client number. A detailed discussion on the number of clients should be given to demonstrate the practicality of the proposed method and to enhance the soundness of the work.

Author response:
`A1`: We provide an experiment with more clients, e.g., C=8, as follows. We train the proposed FedHP under the same setting as Table 4a. 

| #Clients      | FedAvg   | FedHP | $\Delta$ |
| :-------------:|:-------------:|:-----:|:-----:|
| 3      | 31.21 $\pm$  0.10 \ 0.8959 $\pm$ 0.0017  | 31.35 $\pm$ 0.10 \ 0.9033 $\pm$  0.0014| 0.14dB/0.0074  |
| 4      | 31.06 $\pm$  0.10 \ 0.8955 $\pm$ 0.0018  | 31.33 $\pm$ 0.13 \ 0.9023 $\pm$ 0.0018 | 0.27dB/0.0068  |
| 5      | 31.05 $\pm$  0.10 \ 0.9025 $\pm$ 0.0014  | 31.32 $\pm$ 0.19 \ 0.9029 $\pm$ 0.0019 | 0.27dB/0.0004  |
| 8      | 31.17 $\pm$  0.10 \ 0.9033 $\pm$ 0.0014  | 31.42 $\pm$ 0.11 \ 0.9043 $\pm$ 0.0010 | 0.25dB/0.0010  |

*Table R1: Comparison between FedAvg and FedHP with different number of clients. The last column denotes the performance gap.*

As shown in Table R1, the proposed method can achieve a consistent performance boost over FedAvg at a larger number of clients.

Reviewer comment (Experiments):
Is it possible to apply the proposed method to other hyperspectral image datasets?

Author response:
`A2`: We perform experiments on another hyperspectral dataset [1] with 24 spectral channels with the number of client C=3. As the table R2 shown below, the proposed method enables a performance boost over FedAvg. 

| Metrics     | FedAvg   | FedHP | 
| :-------------:|:-------------:|:-----:|
| PSNR      | 29.97 $\pm$  0.31    | 30.55 $\pm$ 0.20  |
| SSIM       |   0.8442 $\pm$ 0.0026 |  0.8471 $\pm$  0.0035  |   

*Table R2: Comparison between FedAvg and FedHP on different hyperspectral dataset [1].*

Reviewer comment (Experiments):
It seems that a competitive method of FedGST for comparison was a centralized learning strategy, is it a fair comparison or what are the modifications toward this method? Please provide more details.

Author response:
`A3`: Thanks for the valuable suggestion!

GST [2] is a centralized learning strategy to handle various hardware masks from the same distribution. In this work, we insert this model directly into the federated framework to enable hardware cooperation. We have added more illustrations about this method into the manuscript. Since all of the methods are compared under the federated learning framework and adopt the same hardware instances, we can thus perform a fair comparison. Actually, considering GST adopts a self-tuning network besides the reconstruction backbone, it requires more training cost to converge.

Reviewer comment (Writing):
There are some typos in the manuscript, for example, Fig.3 caption.

Author response:
`A5`: Thanks for the useful comments! We have revised the typos in Fig. 3 caption.

Reviewer comment (Novelty):
While federated learning is a good choice, it remains unclear if the proposed problem setting can be directly solved by some other simple solutions, such as meta learning or deep ensemble.

Author response:
One key reason we adopt federated learning is to solve the privacy concern. To the best knowledge, both meta learning and deep ensemble are centralized learning strategies and require seeing all of the data to train each client, which makes it hard to perform the hardware corporations among different institutions. 

Besides, we additionally conduct a new experiment using meta learning. Specifically, we integrate MAML [1] into the federated learning framework, termed as FedMAML. We perform experiments using the same setting as Table 1, such that #client=3.   As shown in Table R1 below,  FedMAML method gives limited performance compared with FedAvg and the proposed FedHP. 

| Metrics     | FedMAML   | FedAvg | FedHP |
| :-------------:|:-------------:|:-----:|:-----:|
| PSNR      | 29.00 $\pm$  1.44    | 31.21 $\pm$ 0.10   | 31.35 $\pm$ 0.10   |
| SSIM       | 0.8532 $\pm$ 0.0304 |  0.8959 $\pm$  0.0017 | 0.9033 $\pm$ 0.0014 | 

*Table R1: Comparison between FedAvg and FedHP, and FedMAML*

Reviewer comment (Experiments):
Despite the large improvement given by the hardware prompt, it lacks further analysis of how this design works for different hardware. For example, will different hardware lead to different prompts? What these “hardware prompt” look like? Is the prompt network only implemented by an attention block?

Author response:
Thanks for the detailed suggestion! 

There will be only one prompt network obtained as a function to handle different input  masks, as shown in Fig. 2. Thus, different input hardware will lead to different prompts, which will have the same dimensionality as the input mask but with different pixel values. The prompt network only contains one attention block, which per our observation, can effectively cooperate among clients.

Reviewer comment (Experiments):
What are the benefits of introducing adaptors? Why not directly update the full model?

Author response:
Directly updating the full model can cause cumbersome computational cost and communication cost. As exemplified by Table 3, directly learning client backbones from scratch under a federated framework (FedAvg) can result in $14$ times training time. Together with the prompt network,  adaptor can help enhance efficient fine-tuning performance of pre-trained backbones. As shown in Table 3, the adapter can bring $0.016$dB improvement in PSNR and $0.0037$ boost in SSIM.

Reviewer comment (Experiments):
What’s the main reason for setting C=3 in the experiment?

Author response:
Thanks for the valuable question!

We previously collected $5$ different real hardware instances. Considering the computational cost and our limited resources, we choose to adopt a number of $3$ clients in the main table. In Table 4a, we also report results under more clients, such as C=$4$ and C=$5$. We are still working on collecting more real hardwares.

Reviewer comment (Theory):
In Eq (9), is there any other way to impose a prompt on the measurements? For example, can the concatenation operation be applied?

Author response:
Thanks for the insightful question! 

Directly using the concatenation will cause the dimensionality inconsistent with the backbones and cannot be used. Besides, there is no learnable module for concatenation, which lacks flexibility in handling different data distributions.

Reviewer comment (Presentation):
It would be better to directly explain the settings of different hardware shits in the captions of Table ½.

Author response:
Thanks for the useful suggestion! 

We have added more explanations of the settings of different hardware shifts in captions of Table 1 and 2. Specifically, for Table 1, we provide the explanation that: For different clients, we sample non-overlapping masks from the same mask distribution to train the model and use unseen masks randomly sampled from all clients for testing. For Table 2, we provide the explanation that: Masks from each client are sampled from a specific distribution for training. We  randomly sample non-overlapping masks (unseen to training) from all distributions for testing.

Reviewer comment (Novelty):
The technical contribution to more general computational imaging seems to be limited or at least not well supported by this paper’s current state.

Author response:
Thanks for the valuable comments! 

Our key technical contribution is to provide a new multi-hardware optimization framework adapting to hardware shift by only accessing local data.  The principle underlying the proposed FedHP can be potentially extended to broad SCI applications. However, due to the practical cost of data acquisition and building optics systems, this study explores one specific direction following previous works in the field, where we focus on spectral SCI and collecting optical masks and real data from multiple hardware. 

Exploiting the hardware collaboration of computational imaging systems is still in an early stage. This work serves as a proof of concept to inspire future endeavors in a more general scope. We list several potential related applications that might benefit from the proposed method, such as Lensless camera [1], LiDAR [2], HDR camera [3], or CT-Reconstruction [4], cooperating multiple imaging systems via aligning forward models, etc. We have put this discussion into the related work.

Reviewer comment (Novelty):
Likewise, the main deep learning technic behind this proposal, FedAvg, is already well known, which makes the technical contribution to deep learning community also marginal.

Author response:
Thanks for the useful comments!

We find that FedAvg serves as a very strong baseline by even working better than recent  FL methods, such as SCAFFOLD and FedProx. Thus, it is non-trivial to achieve further performance boost. By comparison, FedHP provides an encouraging performance. 

We kindly summarize the technical contributions of this work as follows:
* We introduce a hardware prompt network to capture the hardware perturbations/replacement. 
* We improve the efficiency of the  federated learning in SCI, achieving performance boost with much lower training time cost. 
* We collect and will release a heterogeneous dataset that covers multiple hardware configurations, which to the best knowledge, is the first one in SCI.

Reviewer comment (Presentation):
Figure 1. It is difficult to find differences between 4. FedAvg and 5. FedHP, the method instruction plot.

Author response:
Thanks for the useful suggestion!

Fig. 1 aims to show different settings for different types of solutions. For federated learning methods including FedAvg and FedHP, we make the settings the same to avoid misunderstanding. We have made it clear by modifying the Fig.1 captions.