Cooperative Hardware-Prompt Learning for Snapshot Compressive Imaging

We much appreciate that the Reviewer aeby provides valuable comments and finds the method is designed with clear purpose.

R4.1: The biggest weakness is arguably that the paper covers a somewhat very niche topic, which is the application of a federated learning scheme to compressive snapshot spectral imaging. To some extent one would expect the technique to abstract away from the specific case of CASSI, as the solution does not particularly pertain to CASSI.

Table T1. Comparison between FedAvg and FedHP on CACTI (Client:3)

A4.1: 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, 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.

Per reviewer's inspiration, we tried our best to launch additional experiments by applying FedHP into another prevalent SCI system of Coded Aperture Compressive Temporal Imaging (CACTI). The results in Table T1 above present a performance boost of FedHP over FedAvg baseline (under the same setting as Table 1), demonstrating that the proposed FedHP does not particularly pertain to CASSI. Due to the limited rebuttal time and the heavy workload in deploying real CACTI systems, we can only obtain the intermediate results of both methods (e.g., $1.2\times10^4$ out of $4\times 10^4$). We are committed to updating the latest results during the discussion stage. We will add the above discussion in the manuscript.

[1] A simple framework for 3D lensless imaging with programmable masks. CVPR 2021.

[2] LiDAR-in-the-loop Hyperparameter Optimization. CVPR 2023.

[3] End-to-end high dynamic range camera pipeline optimization. CVPR 2021.

[4] DOLCE: A model-based probabilistic diffusion framework for limited-angle ct reconstruction. ICCV 2023.

[5] Deep unfolding for snapshot compressive imaging. IJCV 2023.

R4.2: In addition, due to limited data available in this setup and to very limited size datasets, it is difficult to ascertain the significance of the findings.

A4.2: This work has constructed the Snapshot Spectral Heterogeneous Dataset (SSHD), which includes multiple practical spectral SCI systems. This dataset will be released to support further research and validation. Despite the current dataset size, our experiments demonstrate consistent performance improvements using FedHP, even under challenging heterogeneous settings. This provides a strong indication of the robustness and potential of the proposed method.

R4.3: Can the authors extend this to any other compressive imaging scheme? Or perhaps disentangle the improvements in terms of FedHP, from those specific to the application? This would also broaden the data available for validating the experiments.

A4.3: As shown in Table T1 and A4.1, we tried our best to launch additional experiments by extending FedHP into another prevalent SCI system of Coded Aperture Compressive Temporal Imaging (CACTI), which is used to compress and recover temporal information. By comparison, FedHP outperforms FedAvg by 0.52dB/0.0018 in terms of the PSNR/SSIM despite both only training for $1.2\times10^4$ iterations out of $4\times 10^4$. The results indicate that FedHP demonstrates the generalization ability on new SCI systems. We are committed to updating the latest results during the discussion stage. We will add the above discussion in the manuscript.

R4.4: The addressed setup assumes that the problem the authors propose to tackle is meaningfully posed, I.e., that federated learning in the chosen formulation is practically meaningful. The reviewer is not sure whether this is a practically relevant problem considering that CASSI systems are arguably scientific instrumentation/experimental devices whose calibration is likely done per case anyways.

A4.4: While CASSI systems are specialized scientific instruments, the challenge of hardware heterogeneity and data privacy remains significant. Training models for each new hardware configuration can be impractical due to data-starving clients and high computational costs, and the cross-silo cooperation can be impractical without solving privacy concerns for the hardware. FedHP provides a solution and can facilitate broader and more practical deployment of SCI systems in diverse settings.

R4.5: In addition the topic would appear to be more meaningful for publications that cover CASSI systems such as IEEE TGARRS or the like. It is hard for this reviewer to disentangle the margin of novelty of this paper in terms of federated learning approach vs. impact on the target application.

A4.5: Our work intersects federated learning and computational imaging, addressing challenges in both fields. The novelty of our approach lies in the integration of a hardware prompt network within a federated learning framework to handle hardware heterogeneity and data privacy issues. This contribution is relevant to both federated learning and SCI communities, and we believe it provides valuable insights and advancements applicable to various imaging systems.

The reviewer is willing to acknowledge the amount of work done by the authors to prove their point, and has decided to update the recommendation.

Please update, however, your references to mention that significant future work will be done in the direction of lensless cameras, and CT reconstruction, i.e., in the presence of hardware-defined forward models.

This must also include thoroughly all hyperspectral and multispectral compressive imaging strategies that include, among others, random masks applied in a number of ways, including those by random convolution, if the authors see them as feasible.

In absence of other forward models in the current paper, it is important that the authors update the manuscript accordingly with a comprehensive set of references on what forward models could be studied in future work, that should be supported by experimental evidence.

Thanks again for addressing my concerns.

We much appreciate that the Reviewer 1X5M provides valuable comments and finds the proposed work have significant practical relevance to the study and the collected SSHD dataset can benefit the future research. We will release the dataset and the training/testing code.

R3.1: The literature review on federated learning (FL) heterogeneity in the Introduction section lacks comprehensiveness. Additionally, the references included are somewhat outdated.

A3.1: We appreciate the reviewer's commitment in helping improve our draft. We will include more current and diverse references to align with the latest development.

R3.2: the manuscript explains that the coded apertures for each client follow a specific distribution Pc, it does not provide further details about the exact nature or type of this distribution.

A3.2: We visualize the distributions of the real-world masks collected from real CASSI systems in Fig. 5$\sim$7. The coded apertures for each client follow specific distributions and can reflect the practical variations, including perturbations, shaking, and potential replacements.

R3.3: There are many ways to partition data to construct heterogeneous scenarios, such as practical and pathological methods. The approach of equally splitting the training dataset according to the number of clients is not very convincing. The authors should try different partitioning methods.

A3.3: We'd like to clarify that our approach already incorporates both practical and pathological methods to construct heterogeneous scenarios by considering heterogeneity stemming from both the dataset and hardware variations.

• Practical scenario (Table 1): Different clients have non-overlapping masks sampled from the same distribution. This simulates real-world hardware perturbations or imperfections within the same system design, which is a common practical scenario.
• Pathological scenario (Table 2): Each client has masks sampled from a specific distribution. This represents a more extreme case of significant hardware variations or replacements across institutions, which can be considered pathological.

Besides, the equal splitting of the training dataset across clients, combined with these hardware variations, creates a comprehensive heterogeneous environment.

R3.4: It is unclear which datasets were used to obtain the experimental results in Tables 1 and 2. The authors did not specify this, which creates confusion in the experimental analysis.

A3.4: All our experiments use the CAVE dataset for the training and the KAIST dataset for testing, both of which are prevalent benchmarks in the field. We will make it clear in the manuscript about the dataset settings.

R3.5: What is the rationale for using adaptors, and what is their function?

A3.5: The adaptor in FedHP (1) allows efficient fine-tuning of pre-trained backbones, enhancing the model's ability to adapt to new clients. (2) The adaptor significantly reduces communication costs in the FL system by only requiring the transmission of adaptors rather than entire model backbones. As shown in Table 3, the adaptor brings a 0.07dB improvement in PSNR and 0.0029 boost in SSIM. The adaptor is a CONV-GELU-CONV structure governed by a residual connection (L207).

R3.6: What network models are used in the comparison methods? It is necessary to clearly state the fairness of the validated methods.

A3.6: All compared methods use the same network architecture MST [R1] for the reconstruction backbones and a SwinIR [R2] block for the prompt network (see L240~L241). The key differences lie in how each method handles the federated learning aspect and hardware variations. In Table 3, we demonstrate the comparable complexity and cost between FedHP and FedAvg.

[R1] Mask-guided spectral-wise transformer for efficient hyperspectral image reconstruction. In CVPR 2022.

[R2] Swinir: Image restoration using swin transformer. In ICCV 2021.

R3.7: The explanation for Figure 3 is not detailed enough. For example, what is "Patch"?

A3.7: The ``Patch'' corresponds to the small regions (such as a, b in Fig.3 top-left) that are randomly cropped for the spectral accuracy evaluation. We will provide more detailed explanations in Figure 3.

R3.8: In the "Discussion of the client number" section, the number of clients increases very little, and the metrics slightly decline. However, the authors conclude that the performance is stable with the change in the number of clients. The small variation in the number of clients is unconvincing. A larger difference in the number of clients should be set to demonstrate this more effectively.

A3.8: While the performance slightly declines with more clients, FedHP maintains its advantage over FedAvg (0.21dB improvement for C=5 and 0.19dB for C=10). We conjecture that the slight descent tendency might attribute to more complex corporations among the clients and less sufficient training samples for each client. It is non-trivial to collect real hardware systems due to the privacy concern and the cost. FedHP is dedicated to facilitate this process by offering a way of cross-silo cooperation. We are also working on collecting larger-scale data and real hardware.

R3.9: The authors mention the "presence of data privacy" in the contributions, but there is no further discussion or experimental comparison regarding data privacy in the subsequent sections. This makes it difficult to validate their contribution to data privacy protection.

A3.9: Thanks for the valuable comment. While we don't explicitly discuss data privacy in our experiments, it is inherent to the federated learning framework we employ. FedHP, like other FL methods, ensures that raw data never leaves the local clients, addressing privacy concerns by design. We will add a dedicated subsection to discuss how our method preserves data privacy and compare it with centralized approaches in terms of privacy protection.

The authors have addressed most of my concerns; however, two issues remain unresolved, so I will keep my rating unchanged.

1. Regarding the discussion of the client number, Table 4(a) in the paper shows that when C=5, FedHP outperforms FedAvg by 0.27 dB. However, in the rebuttal (Section A3.8), this performance gap is reported as 0.21 dB. The authors have not provided an explanation for this discrepancy.
2. The authors mentioned in the rebuttal that they would add a detailed description of Privacy Protection, but this has not been presented.

We much appreciate that the Reviewer BdgG provides valuable comments and finds the problem novel and the method convincing.

R1.1: There are some typos in the writing. For example, the caption of Figure 3 and the bold parts in the second row of Table 1 and the eighth row of Table 2 are confusing.

A1.1: We will correct the caption of Fig. 3 and the bold parts in Tables 1 and 2 to improve clarity. We will thoroughly proofread the manuscript to eliminate other errors.

R1.2: The proposed FedHP method is relatively straightforward and lacks deeper insights. Moreover, it does not show a significant performance improvement compared to FedAvg.

A1.2: FedHP offers insightful contributions:

• It explicitly models hardware variations in computational imaging systems, which is non-trivial and addresses a key practical challenge.
• The plug-and-play prompt network allows learnable hardware configurations, enabling hardware-software co-optimization.
• The combination of prompt learning and adaptors reduces communication costs in federated learning systems.

We empirically find that FedHP enables a consistent performance boost over FedAvg on different number clients, demonstrating the scalability of the proposed method. Considering that FedAvg is a strong baseline that even outperforms prevalent FL methods such as FedProx and SCAFFOLD, it is non-trivial to achieve performance boost over FedAvg.

R1.3: I suggest adding experiments to test the applicability of other SCI systems, such as Coded Aperture Compressive Temporal Imaging (CACTI).

Table T1. Comparison between FedAvg and FedHP on CACTI (Client:3)

A1.3: We launch additional experiments by applying FedHP to the real CACTI systems. The results in Table T1 above present a better reconstruction performance of FedHP over FedAvg baseline (under the same setting as Table 1), demonstrating the generalizability ability of FedHP over a new SCI architecture. We will add the above discussion in the manuscript.

Due to the limited rebuttal time and the heavy workload in deploying real CACTI systems, we can only obtain the intermediate results of both methods (e.g., $1.2\times 10^4$ out of $4\times 10^4$). We are committed to updating the latest results during the discussion stage.

R1.4: There is a lack of sufficient real-world experiments. It would be beneficial to set up multiple independent SCI systems to test the algorithm's performance.

A1.4: We kindly illustrate that all experiments are performed upon multiple independent real-world SCI systems. For both training and testing, we not only consider sampling non-overlapping masks from the same mask distribution, but also explore a challenging scenario where each client can have real masks sampled from a specific distribution, mimicking the diverse real-world scenarios including hardware perturbation and replacement (L27).

R1.5: All the experiments in this paper are based on the SD-CASSI model. Can the same FedHP model be simultaneously applicable to both DD-CASSI and SD-CASSI architectures, which have significantly different designs?

A1.5: It is challenging to apply FedHP simultaneously to both DD-CASSI and SD-CASSI architectures due to their different optical designs and forward models. DD-CASSI uses two dispersive elements and a coded aperture, while SD-CASSI uses a single dispersive element and a coded aperture, resulting in distinct measurement formation processes. These differences lead to incompatible forward models, making it difficult to unify them under a single reconstruction model. However, we agree this is an interesting direction for future research.

R1.6: There are still many artifacts in the reconstructed images. While I understand that this is maybe due to the precision issues of the CASSI system, it is crucial for evaluating the practical usability of the algorithm.

A1.6: Assessing the practical usability of SCI algorithms is inherently difficult due to the complex interplay of hardware limitations, reconstruction algorithms, and application-specific requirements. The proposed FedHP method itself is actually a step towards uncovering and tackling these practical challenges. By explicitly modeling hardware variations and enabling cross-hardware learning, FedHP aims to address real-world issues that arise when deploying SCI systems across multiple devices or institutions.

R1.7: I am not sure whether the spectral accuracy of the reconstructed images is also optimal in statistical terms, which is vital for spectral imaging systems.

Table T2. Spectral accuracy comparison on three testing scenes.

A1.7: We provide the spectral accuracy comparison in Fig.3, 5, and 6. Specifically, we randomly select a visually informative region (patch) on the hyperspectral images and compute the density values for a specific wavelength (i.e., dividing pixel values on this specific wavelength by the summarized pixel values along all the wavelengths). We statistically measure the spectral accuracy by computing the correlation of the density values between the prediction and the ground truth. FedHP enables better spectral accuracy by comparison.

We further perform a statistical experiment on three simulation data, on which we randomly choose 10 spatial regions for the averaged correlation computation. As shown in Table T2 above, FedHP consistently enables a higher averaged correlation over FedAvg, indicating a better spectral accuracy. We will add the above discussions into the final version.

We much appreciate that the Reviewer Jxiy provides valuable comments and finds the proposed method addresses the issue from the new perspective of hardware with good performance.

R2.1: The number of clients used in the experiments is still relatively small. Although a simple comparison of the impact of different numbers of clients was made, there is not much difference in performance compared to other methods when the number of clients is larger.

A2.1: We find FedHP maintains its performance advantage even with increased client numbers (C=5 and C=10), outperforming FedAvg by 0.21dB and 0.19dB respectively as shown in Table 4 (a). These results demonstrate FedHP's scalability and consistent performance improvements across varying numbers of clients. Besides, it is non-trivial to collect real hardware systems due to the privacy concern and the cost. We are dedicated to facilitating this process by offering a way of cross-silo cooperation using FedHP. We are still working on collecting larger-scale data and real-world hardware systems.

R2.2: Although good results were reported on simulated data, more results on real data should be included to evaluate the effectiveness of the proposed method.

A2.2: We provide additional comparisons on real data in Fig. 7$\sim$8. We are still working on collecting more data.

R2.3: Why does the prompter lead to such a significant improvement, while the effect of the adaptor is not as pronounced? Please provide an in-depth analysis.

A2.3: Thanks for the valuable comment. The prompter plays a crucial role in capturing hardware-specific characteristics and aligning inconsistent data distributions across clients. It directly addresses the heterogeneity rooted in the input data space, which is a key challenge in federated learning for SCI systems. By comparison, the adaptor, while important for efficient fine-tuning, has a more subtle effect on performance. Its primary role is to reduce communication costs in the FL system by allowing us to use pre-trained backbones and only communicate adaptors.

R2.4: The generalization to different hardware systems is crucial for deep learning based methods. The current form of this manuscript only reports on a small scale real dataset captured by several systems. A larger dataset captured by more systems is necessary to evaluate the method.

A2.4: We kindly illustrate that current setup can consider a large amount of the real-hardware under two challenging scenarios.

• Same distribution: As shown in Table 1, we sample non-overlapping masks from the same distribution for different clients. This simulates hardware perturbations or imperfections within the same system design.
• Different distributions: In Table 2, we explore a more challenging scenario where each client has masks sampled from a specific distribution. This closely mimics diverse real-world scenarios, including hardware replacement or significant variations across institutions.

These experiments, conducted on real hardware masks with noise (Fig. 9-11 in supplementary), demonstrate FedHP's ability to generalize across various hardware configurations. We appreciate the reviewer's insights on collecting larger dataset. This work laid the foundation for future extension by collecting a Snapshot Spectral Heterogeneous Dataset built upon multiple practical SCI systems for the first time. We are still working on collecting more data and real systems for better evaluation.