Cooperative Hardware-Prompt Learning for Snapshot Compressive Imaging

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-07 05:45:55

Abstract: Existing reconstruction models in snapshot compressive imaging systems (SCI) are trained with a single well-calibrated hardware instance, making their perfor- mance vulnerable to hardware shifts and limited in adapting to multiple hardware configurations. To facilitate cross-hardware learning, previous efforts attempt to directly collect multi-hardware data and perform centralized training, which is impractical due to severe user data privacy concerns and hardware heterogeneity across different platforms/institutions. In this study, we explicitly consider data privacy and heterogeneity in cooperatively optimizing SCI systems by proposing a Federated Hardware-Prompt learning (FedHP) framework. Rather than mitigating the client drift by rectifying the gradients, which only takes effect on the learning manifold but fails to solve the heterogeneity rooted in the input data space, FedHP learns a hardware-conditioned prompter to align inconsistent data distribution across clients, serving as an indicator of the data inconsistency among different hardware (e.g., coded apertures). Extensive experimental results demonstrate that the proposed FedHP coordinates the pre-trained model to multiple hardware con- figurations, outperforming prevalent FL frameworks for 0.35dB under challenging heterogeneous settings. Moreover, a Snapshot Spectral Heterogeneous Dataset has been built upon multiple practical SCI systems. Data and code are aveilable at https://github.com/Jiamian-Wang/FedHP-Snapshot-Compressive-Imaging.git

Main paper: 
 introduction: 
 The technology of snapshot compressive imaging (SCI) (Yuan et al., 2021) has gained prominence in the realm of computational imaging. Taking an example of hyperspectral image reconstruction, the spectral SCI (Gehm et al., 2007) can fast capture and compress 3D hyperspectral signals as 2D measurements through optical hardware, and then restore the original signals with high fidelity by training deep neural networks (Meng et al., 2020; Miao et al., 2019) . Despite the remarkable performance (Cai et al., 2022a; b; Lin et al., 2022; Huang et al., 2021; Hu et al., 2022) , existing deep SCI methods are generally trained with a specific hardware configuration, e.g., a well-calibrated coded aperture (physical mask). The resulting model is vulnerable to the hardware shift/perturbation and limited in adapting to multiple hardware configurations. However, directly learning a reconstruction model cooperatively from multi-hardware seems to be infeasible due to data proprietary constraint. It is also non-trivial to coordinate heterogeneous hardware instances with a unified model.To elaborate, we firstly recap previous research efforts of centralized learning solutions. A naive treatment is to jointly train a single reconstruction model with data collected from different hardware configurations, i.e., coded apertures. As shown in Fig. 1 right, this solution enhances the ability of reconstruction (> 0.5dB+) by comparison to single hardware training scenario. However, the performance on inconsistent coded apertures is still non-guaranteed since the model only learns to fit coded apertures in a purely data-driven manner. Followed by, Self-tuning (Wang et al., 2022) advances the learning by explicitly approximating the posterior distribution of coded apertures in a variational Bayesian framework. Despite the significant performance boost, it is only compatible with the coded apertures drawing from homogeneous hardware (same distribution) yet cannot handle heterogeneous hardware. Nevertheless, centralized learning presumes that hardware instances and hyperspectral data A series of hardware (3) Self-tuning (Wang et al., 2022) are centralized training solutions. Both (4) FedAvg and the proposed (5) FedHP adopt the same data split setting. We compare the performance gain of different methods over (1) on the right. Note that P 1 , P 2 , and P 3 present the practical mask distributions. All results are evaluated by unseen masks (non-overlapping) randomly sampled from {P 1 , P 2 , P 3 }.! # ~# % % ! % $ % % ! # ! ~#! ! # $ ~#$ ! # % ~#% % ! % $ % % ! # ! ~#! ! # $ ~#$ ! #are always publicly available, which hardly holds in practice -both the optical systems (with different confidential configurations, e.g., coded apertures) and data samples (i.e., measurements captured from non-overlapping scenes) are generally proprietary assets across institutions, adhering to the strict privacy policy constraints (Vergara-Laurens et al., 2016; Li et al., 2021) , while considering the multi-hardware cooperative training confining to this concern remains unexplored in SCI.In this work, we leverage federated learning (FL) (Kairouz et al., 2021; Li et al., 2020a; Wang et al., 2021) for cross-platform/silo multi-hardware reconstruction modeling without sharing the hardware configurations. Firstly, the FL benchmark, FedAvg (McMahan et al., 2017) , is adopted and brings performance boost (compared by 3 and 4 in Fig. 1 right). However, FedAvg has been proven to be limited in solving heterogeneous data (Hsu et al., 2019; Karimireddy et al., 2020) -the heterogeneity in SCI substantially stems from the hardware, which is usually absorbed into the compressed data and governs the network training. Thus, different configurations, e.g., coded apertures, yield different data distributions. Besides, we consider a more practical scenario by extending the sample-wise hardware difference into distribution-wise, i.e., not only the different coded apertures yield heterogeneity, but also coded apertures from different clients may follow different distributions (see P 1 ∼ P 3 in Fig. 1 ).To further approach the heterogeneity issue, this work proposes a Federated Hardware-Prompt (FedHP) framework to achieve multi-hardware cooperative learning with privacy piratically preserved.Prevalent FL methods handle the heterogeneity by regularizing the global/local gradients (Karimireddy et al., 2020; Li et al., 2020b) , which only take effect on the learning manifold but fail to solve the heterogeneity rooted in the input data space. Differently, FedHP traces back to the source of the data heterogeneity of this application, i.e., inconsistent hardware configurations, and devises a prompt network to solve the client drift issue in input data space. By taking the coded aperture as input, the prompter better accounts for the underlying inconsistency and close the gap between input data distributions across clients. Besides, the prompter explicitly models the correlation between the software and hardware, empowering the learning by following the spirit of the co-optimization (Goudreault et al., 2023; Zheng et al., 2021; Robidoux et al., 2021 ) in computational imaging. In addition, FedHP directly operates on pre-trained reconstruction backbones with locally well-trained models and keeps them frozen throughout the learning, which potentially improves the efficiency than directly optimizing the reconstruction backbones in FL from scratch. The contributions are as follows:• We introduce and tackle an unexplored practical problem of hardware cooperative learning in SCI, under the presence of data privacy constraint and the heterogeneous configurations. Our proposed FedHP, to the best knowledge, bridges FL and the filed of spectral SCI for the first time. • We uncover the data heterogeneity of SCI that stems from distinct hardware configurations. A hardware prompt module is developed to solve the distribution shift across clients and empower the hardware-software co-optimization in computational imaging. The proposed method provides an orthogonal perspective in handling the heterogeneity to the existing FL practices. • We collect the first Snapshot Spectral Heterogeneous Dataset (SSHD) from a series of practical spectral snapshot imaging systems. Extensive experiments demonstrate that FedHP outperforms both centralized learning methods and classic federated learning treatments. This work can inspire future works in this novel research direction of hardware collaboration in SCI. 
method 2.1 preliminary knowledge: 
 In this work, we study the cooperative learning problem in SCI by taking the representative setup of coded aperture snapshot spectral imaging system for hyperspectral imaging as an example, due to its recent advances (Cai et al., 2022a; b; Lin et al., 2022) . Given the real-world hyperspectral signal X ∈ R H×W ×N λ , where N λ denotes the number of spectral channels, the hardware performs the compression with the physical coded apterture M of the size H ×W , i.e., M hw ∈ [0, 1]. Accordingly, the encoding process produces a 2D measurement Y M ∈ R H×(W +∆) , where ∆ denotes the shiftingEQUATIONwhere ⊙ denotes the pixel-wise multiplication and Ω presents the measurement noise. For each spectral wavelength λ, the corresponding signal X(:, :, n λ ) is shifted according to the function d(λλ * ) by referring to the pre-defined anchor wavelength λ * , such that ∆ = d(N λ -1). Following the optical encoder, recent practices train a deep reconstruction network f (•) to retrieve the hyperspectral data X ∈ R H×W ×N λ by taking the 2D measurement Y M as input. We define the initial training dataset of the hyperspectral signal as D and the corresponding dataset for the reconstruction asEQUATIONwhere X i serves as the ground truth and Y M * i is governed by a specific real physical coded aperture M * . The reconstruction model finds the local optimum by minimizing the mean squared lossEQUATIONwhere θ expresses all learnable parameters in the reconstruction model. Thanks to the sophisticated network designs (Cai et al., 2022a; Huang et al., 2021) , pre-trained reconstruction model demonstrates promising performance when is compatible with a single encoder set-up, where the measurement in training and testing phases are produced by the same hardware using a fixed coded aperture of M * .Motivation. However, previous work (Wang et al., 2022) uncovered that most existing reconstruction models experience large performance descent (e.g., > 2dB in terms of PSNR) when handling the input measurements encoded by a different coded aperture M † from training, i.e., M † ̸ = M * . This is because the coded aperture M * implicitly affects the learning as shown by equation 3. To this end, the well-trained reconstruction model can be highly sensitive to a specific hardware configuration of coded aperture and is hardly compatible with the other optical systems in the testing phase. A simple solution of adapting the reconstruction network to a different coded aperture M † is to retrain the model with corresponding datasetD M † = {Y M † i , X i } i=N i=1and then test upon M † accordingly. However, this solution does not broaden the adaptability of reconstruction models to multi-hardware and can introduce drastic computation overhead. In this work, we tackle this challenge by learning a reconstruction model cooperatively from multiple hardware with inconsistent configurations.Jointly Train. To solve the above problem, Jointly train (Fig. 1 part 2) serves as a naive solution to train a model with data jointly collected upon a series of hardware. Assuming there are total number of K hardware with different coded apertures, i.e., M 1 , M 2 , ..., M K . Each hardware produces a training dataset upon D asD M k = {Y M k i , X i } i=N i=1 . The joint training dataset for reconstruction is D M 1∼K = D M1 ∪ D M2 ∪ . . . ∪ D M K ,(4) where different coded apertures can be regarded as hardware-driven data augmentation treatments toward the hyperspectral data representation. The reconstruction model will be trained with the same mean squared loss provided in equation 3 upon D M 1∼K . Wang et al. (2022) demonstrated that jointly learning brings performance boost compared with single mask training (Fig. 1 right ). However, this method adopts a single well-trained model to handle coded apertures, failing to adaptively cope with the underlying discrepancies and thus, leading to compromised performances for different hardware.Self-tuning. Following Jointly train, recent work of Self-tuning (Wang et al., 2022) recognizes the coded aperture that plays the role of hyperprameter of the reconstruction network, and develops a hyper-net to explicitly model the posterior distribution of the coded aperture by observing D M 1∼K . Specifically, the hyper-net h(σ; M k ) approximates P (M|D M 1∼K ) by minimizing the Kullback-Leibler divergence between this posterior and a variational distribution Q(M) parameterized by σ. Compared with Jointly train, Self-tuning learns to adapt to different coded apertures and appropriately calibrates the reconstruction network during training, even if there are unseen coded apertures. However, the variational Bayesian learning poses a strict distribution constraint to the sampled coded apertures, which limits the scope of Self-tuning under the practical setting.To sum up, both of the Jointly train and Self-tuning are representative solutions of centralized learning, where the dataset D and hardware instances with M 1 , ..., M K from different sources are presumed to be publicly available. Such a setting has two-fold limitations. (1) Centralized learning does not take the privacy concern into consideration. In practice, both the hardware instances and hyperspectral dataset are proprietary assets of institutions and thus, corresponding hardware configuration and data information sharing is subject to the rigorous policy constraint. (2) Existing centralized learning methods mainly consider the scenario where coded apertures are sampled from the same distribution, i.e., hardware origin from the same source, which is problematic when it comes to the coded aperture distribution inconsistency especially in the cross-silo case. Bearing the above challenges, in the following, we resort to the federated learning (FL) methods to solve the cooperative learning of reconstruction considering the privacy and hardware configuration inconsistency.FedAvg. We firstly tailor FedAvg (McMahan et al., 2017) , into SCI. Specifically, we exploit a practical setting of cross-silo learning in snapshot compressive imaging. Suppose there are C clients, where each client is packaged with a group of hardware following a specific distribution of EQUATIONEQUATIONwhere N c represents the number of hyperspectral data in D c . The local learning objective isEQUATIONwhere X i = f ( θ; Y M c i ), M c ∼ P c ,EQUATIONwhere C ′ denotes the number of clients that participate in the current global round and α c represents the aggregation weight. Compared with the centralized learning solutions, FedAvg not only bridges the local hyperspectral data without sharing sensitive information, but also collaborates multihardware with a unified reconstruction model for a better performance (Fig. 1 right comparison between 3 and 4). However, FedAvg shows limitations in two-folds. (1) It has been shown that FedAvg is hard to handle the heterogeneous data (Karimireddy et al., 2020; Khaled et al., 2020; Hsu et al., 2019) . ( 2) Directly training and aggregating the reconstruction backbones from scratch would introduce prohibitive computation. In the following, we firstly introduce the hardware-induced data heterogeneity in SCI. Then we develop a Federated Hardware-Prompt (FedHP) method to achieve cooperative learning without optimizing the client backbones.Data Heterogeneity. We firstly consider the data heterogeneity stems from the different coded apertures samples, i.e., hardware instances. According to Section 2.1, the optical hardware samples the hyperspectral signal constitutes D M and further serves as the input data for the reconstruction model. To this end, the modality of {Y M i } i=1 i=N is vulnerable to the coded aperture variation. A single coded aperture M defines a unique input data distribution for the reconstruction, i.e., Y M i ∼ P M (Y M i ). For arbitrary distinct coded apertures, we haveX i from D = {X i } i=N i=1 and encodes it into a 2D measurement Y M i , which 2. Local Update (Prompt) 3. Local Update (Adaptor) 4. Aggregation 1. Initialize ⋯ ⋯ Client 1 Client 2 Client # " ! " " " # " $ # $ # # # " " % # % $ % ! ! " Prompter Aligned input for reconstruction learning ! " % # % $ % Global model CONV GELU CONV MHSA LN Adaptor Adaptor MLP LN ! ⋯ " ! " " " # " $ ! $ # % $ % #P M * (Y M * i ) ̸ = P M † (Y M † i ) if M * ̸ = M † .In federated learning, data heterogeneity persistently exists since there is no identical coded aperture across different clients. We name this heterogeneous scenario as Hardware shaking, which potentially can attribute to physical perturbations such as lightning distortion or optical platform fluttering.We take a step further to consider the other type of data heterogeneity stemming from the distinct distributions of coded apertures1 . As formulated in equation 6, each client collects a coded aperture assemble following the distribution P c for c-th client. We have P c differs from one another, i.e., P c1 ̸ = P c2 for c1 ̸ = c2, c1, c2 ∈ {1, ..., C}. We name this heterogeneous scenario as Manufacturing discrepancy, where hardware instances from different clients (institutions) are produced by distinct manufacturing agencies, so that the distribution P c1 and P c2 drastically differs as demonstrated in Fig. 1 . This turns out to be a more challenging scenario than Hardware shaking. As presented in Section 3.2, classic federated learning methods, e.g., FedProx (Li et al., 2020b) and SCAFFOLD (Karimireddy et al., 2020) hardly converge in this setting. By comparison, the proposed method enables an obvious performance boost.Hardware-Prompt Learning. Bearing the heterogeneous issue, previous efforts (Li et al., 2020b; Karimireddy et al., 2020) mainly focus on rectifying the global/local gradients upon training, which only takes effect on the learning manifold but fail to solve the heterogeneity rooted in the input data space, whose effectiveness in this low-level vision task may be limited. Since we explicitly uncover and determine two types of the heterogeneity in snapshot compressive imaging stemming from the hardware inconsistency (Section. 2.3), this work opts to tackling the client drift issue by directly operating in the input data space. This can be achieved by collaboratively learning the input data alignment given different coded apertures. In light of the visual prompt tuning in large models (Liu et al., 2023b; Bahng et al., 2022) , we devise a hardware-conditioned prompt network in the following.Step 2 of Fig. 2 , given the input data {Y M i } i=N i=1 of the reconstruction, the prompt network aligns the input samples, i.e., measurements Y Mi , by adding a prompter conditioned on the hardware configuration. Let Φ(ϕ; M) denote the prompt network (e.g., attention block) parameterized by ϕ and Y M i is produced upon coded aperture M. Then the resulting input sample is aligned asEQUATIONIn the proposed method, the prompt network collaborates different clients with inconsistent hardware configurations. It takes effect by implicitly observing and collecting diverse coded aperture samples of all clients, and jointly learns to react to different hardware settings. The prompter regularizes the input data space and achieves the goal of coping with heterogeneity sourcing from hardware.Training. As shown in Fig. 2 , we demonstrate the training process of proposed FedHP by taking one global round as an example2 . Since the prompt learning takes effect on pre-trained models, we initialize the c-th backbone parameters with the pre-trained model θ p c on local data D M c with equation 7. The global prompt network ϕ G is randomly initialized and distributed to the c-th clientEQUATIONwhere ϕ c is the local prompt network and C ′ denotes the number of clients participated in current global round. To enable better response of the pre-trained backbone toward the aligned input data space, we also introduce the adaptors into the transformer backbone. As shown in Fig. 2 Step 3, we show the architecture of the proposed adaptor, which is a CONV-GELU-CONV structure governed by a residual connection. We insert the adaptors behind the LN layers throughout the network.We then perform local update in each global round. It is composed of two stages. Firstly, we update the local prompt network ϕ c for S p iterations, all the other learnable parameters of backbone and adaptors are frozen. The loss then becomesEQUATIONwhere we use ϵ c to represent learnable parameters of all adaptors for c-th client. Secondly, we tune the adaptors for another S b iterations. Both of the pre-trained backbone and prompt network are frozen. The loss of c-th client shares the same formulation as equation 11. After the local update, FedHP uploads and aggregates the learnable parameters ϕ c , c = 1, ..., C of prompt network. Since the proposed method does not require to optimize and communicate the reconstruction backbones, the underlying cost is drastically reduced considering the marginal model size of prompt network and adptors compared with the backbone, which potentially serves as a supplied benefit of FedHP.Compared with FedAvg, FedHP adopts the hardware prompt to explicitly align the input data representation and handle the distribution shift attributing to the coded aperture inconsistency (hardware shaking) or coded aperture distribution discrepancy (manufacturing discrepancy). 
implementation details: 
 Dataset. Following existing practices (Cai et al., 2022b; Lin et al., 2022; Hu et al., 2022; Huang et al., 2021) , we adopt the benchmark training dataset of CAVE (Yasuma et al., 2010) , which is composed of 32 hyperspectral images with the spatial size as 512 × 512. Data augmentation techniques of rotation, flipping are employed, producing 205 different training scenes. For the federated learning, we equally split the training dataset according to the number of clients C. The local training dataset are kept and accessed confidentially across clients. Note that one specific coded aperture determines a unique dataset according to equation 2, the resulting data samples for each client can be much more than 205/C. We employ the widely-used simulation testing dataset for the quantitative evaluation, which consists of ten 256 × 256 × 28 hyperspectral images collected from KAIST (Choi et al., 2017) . Besides, we use the real testing data with spatial size of 660 × 660 collected by a SD-CASSI system (Meng et al., 2020) for the perceptual evaluation considering the real-world perturbations.Hardware. We collect and will release the first Snapshot Spectral Heterogeneous Dataset (SSHD) containing a series of practical SCI systems, e.g., CASSI, from three agencies, each of which offers a series of coded apertures that correspond to a unique distribution3 as presented by federated settings in Fig. 2 . No identical coded apertures exists among all systems. For the case of manufacturing discrepancy, we directly assign hardware systems from one source to form a client. We simulate the scenario of hardware shaking by distributing coded apertures from one source to different clients.Implementation details. We adopt the popular transformer backbone, MST-S (Cai et al., 2022a) for the reconstruction. Besides, the prompt network is instantiated by a SwinIR (Liang et al., 2021) Table 1 : PSNR(dB)/SSIM performance comparison. 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. We report mean ±std among 100 trials for FedAvg (McMahan et al., 2017) , FedProx (Li et al., 2020b ), SCAFFOLD (Karimireddy et al., 2020) , and FedGST (Wang et al., 2022) . block. Limited by the computational resource, we set the number of clients as 3 in main comparison.We empirically find that collaborate such amount of clients can be problematic for popular federated learning methods under the very challenging scenario of data heterogeneity (see Section 3.2). For FL methods, we update all clients throughout the training, i.e., C ′ = C = 3. For the proposed method, we pre-train the client backbones from scratch for 4 × 10 4 iterations on their local data. Notably, the total training iterations of different methods are kept as 1.25 × 10 5 for a fair comparison. The batch is set as 12. We set the initial learning rate for both of the prompt network and adaptor as α p = α b = 1 × 10 -4 with step schedulers, i.e., half annealing every 2×10 4 iterations. We train the model with an Adam (Kingma & Ba, 2014) optimizer (β 1 = 0.9, β 2 = 0.999). We implement the proposed method with PyTorch (Paszke et al., 2017) on an NVIDIA A100 GPU.Compared Methods. We compare the proposed FedHP with mainstream federated learning methods, including FedAvg (McMahan et al., 2017) , FedProx (Li et al., 2020b), and SCAFFOLD (Karimireddy et al., 2020) . Besides, GST (Wang et al., 2022) paves the way for the robustness of the reconstruction toward multiple hardware. Thereby, we integrate this method into the federated learning framework, dubbed as FedGST. Notably, all the compared methods require to train and aggregate the client backbones without a prompter or adaptor introduced. By comparison, FedHP updates and shares the prompt network, outperforming the others with smaller amount of parameters being optimized and communicated. We adopt PSNR and SSIM (Wang et al., 2004) for the quantitative evaluation.Simulation Results. We quantitatively compare different methods in Table 1 by considering the data heterogeneity stems from hardware shaking. FedHP performs better than the classic federated learning methods. By comparison, FedProx and SCAFFOLD only allows sub-optimal performance, Table 2 : PSNR(dB)/SSIM performance comparison. 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. We report mean ±std among 100 trials for FedAvg, FedProx (Li et al., 2020b) , SCAFFOLD (Karimireddy et al., 2020) , and FedGST (Wang et al., 2022) . which uncovers the limitations of rectifying the gradient directions in this challenging task. Besides, FedGST works inferior than FedHP, since FedGST approximates the posterior and expects coded apertures strictly follows the identical distribution, which can not be guaranteed in practice. In Fig. 3 , we visualize the reconstruction results with sampled wavelengths. FedHP not only enables a more granular retrieval on unseen coded aperture, but also maintains a promising spectral consistency.Challenging Scenario of Heterogeneity. We further consider a more challenging scenario of manufacturing discrepancy, where the data heterogeneity is highly entangled to the large gap among the coded aperture distributions. We compare different methods in Table 2 . By observation, all methods experience large performance degradation, among which FedProx and SCAFFOLD becomes nearly ineffective. Intuitively, it is hard to concur the clients under the large distribution gap, while directly adjusting the input data space better tackles the problem.Real Results. In Fig. 4 , we visually compare the FedAvg with FedHP on the real data. Specifically, both methods are evaluated under an unseen hardware configuration, i.e., coded aperture from an uncertain distribution. The proposed method introduces less distortions among different wavelengths. Such an observation endorses FedHP a great potential in collaborating hardware systems practically.We conduct model discussion in Table 3 . Specifically, we accumulate the total cost (e.g., number of parameters, GMACs, and training time) of all clients in a federated system.Ablation Study. We firstly consider a scenario that trains three clients independently without FL (FedHP w/o FL). For a fair comparison, each client pre-trains the backbone by using the same procedure as FedHP and are then enhanced with a prompt network and adaptors for efficient fine-tuning. By comparison, FedHP enables an obvious improvement (0.6dB) by implicitly sharing the hardware and data. We then investigate the effectiveness of the prompter and adaptor to the reconstruction, respectively. By observation, directly removing the adaptor leads to limited performance descent. Using prompt network brings significant performance boost. The hardware prompter aligns the input data distributions, potentially solving the heterogeneity rooted in the input data space. This is determined by the fact that learning manifold is highly correlated with the coded apertures.Discussion of the client number. In Table 4a , we discuss the power of FedHP with more real clients under the scenario of Hardware shaking. The performance gap between FedHP and FedAvg consistently remains with the client number increasing, which demonstrates the practicability of the FedHP for the cross-silo spectral system cooperative learning, e.g., 3 ∼ 5 clients/institutions. Comparison with a deep unfolding method. We also compare the proposed FedHP with a representative deep unfolding method of GAP-Net (Meng et al., 2023) as deep unfolding methods can be adaptable to various hardware configurations. Specifically, we use three clients and keep training and testing settings of GAP-Net the same as FedHP. As shown in Table 4b , FedHP improves by 0.28dB with only 7% model size. In fact, despite the adaptability, deep unfolding still shows limitations in solving hardware perturbation/replacement for a given system as discussed in Wang et al. (2022) . 
related work: 
 Hyperspectral Image Reconstruction. In hyperspectral image reconstruction (HSI), learning deep reconstruction models (Cai et al., 2022a; b; Lin et al., 2022; Huang et al., 2021; Meng et al., 2020; Hu et al., 2022; Miao et al., 2019) has been the forefront among recent efforts due to high-fidelity reconstruction and high-efficiency. Among them, MST (Cai et al., 2022a) devises the first transformer backbone by computing spectral attention. By observation, existing reconstruction learning strategies mainly considers the compatibility toward a single hardware instance. The learned model can be highly sensitive to the variation of hardware. To tackle this practical challenge, GST (Wang et al., 2022) paves the way by proposing a variational Bayesian learning treatment.Federated Learning. Federated learning (Kairouz et al., 2021; Li et al., 2020a; Wang et al., 2021) collaborates client models without sharing the privacy-sensitive assets. However, FL learning suffers from client drift across clients attributing to the data heterogeneity issue. One mainstream (Karimireddy et al., 2020; Li et al., 2020b; Xu et al., 2021; Jhunjhunwala et al., 2023; Reddi et al., 2021) mainly focus on regularizing the global/local gradients. As another direction, personalized FL methods (Collins et al., 2021; Chen & Chao, 2022; Fallah et al., 2020; T Dinh et al., 2020; Jiang & Lin, 2023) propose to fine-tune the global model for better adaptability on clients. However, customizing the global model on client data sacrifices the underlying robustness upon data distribution shift (Wu et al., 2022; Jiang & Lin, 2023) , which contradicts with our goal of emphasizing the generality across hardware and thus is not considered. In this work, we propose a federated learning framework to solve the multi-hardware cooperative learning considering the data privacy and heterogeneity, which to the best knowledge, is the first attempt of empowering spectral SCI with FL. Besides, the principle underlying this method can be potentially extended to broad computational imaging applications (Zheng et al., 2021; Liu et al., 2023a; Goudreault et al., 2023; Robidoux et al., 2021) 5 CONCLUSIONSIn this work, we observed an unexplored research scenario of multiple hardware cooperative learning in spectral SCI, considering two practical challenges of data proprietary constraint and heterogeneity stemming from the inconsistent hardware configurations. We developed Federated Hardware-Prompt (FedHP) learning framework to solve the distribution shift across clients and empower the hardwaresoftware co-optimization. The proposed method serves as a first attempt of exploiting the power of FL in spectral SCI. Besides, we collect a Snapshot Spectral Heterogeneous Dataset (SSHD) from multiple real spectral SCI systems. Future works may theoretically derive the convergence of FedHP and exploit the behavior of FedHP under a large number of clients (e.g., > 100). We hope this work will inspire future works in this novel direction of hardware collaboration in SCI.& & & & ! & $& % % &We presume that the hyperspectral single dataset Dc, c = 1, ..., C, shares the same distribution by generally capturing the natural scenes. Heterogeneity stems from the hyperspectral signal is out of the scope of this work.Due to the limited space, we provide an algorithm of FedHP in supplementary.More illustrations and distribution visualizations of real collected coded apertures are in supplementary. 
 
 Appendix:  


Decision: The paper presents a hardware focused method for performing snapshot compressive imaging in cases where differences in hardware are adjusted in the trained model. Federated learning is used towards this end. Four reviewers considered the paper, with two voting Weak Accept and two Borderline Accept. There was good engagement in the discussion period to indicate that they think the idea a good one and the results are potentially of interest. While the average score for this paper technically puts it on the threshold between accept and reject, the small variance between reviewers, and the fact that all consider the paper above threshold for accept, is a strong signal that it will contribute to the conference and is worth accepting.

Reviews: review: paper_summary: paper_summary: The paper introduces FedHP, a reconstruction method for snapshot compressive imaging systems, which addresses the challenge of cross-hardware learning by proposing a federated learning approach. The key contribution lies in using a hardware-conditioned prompter to align data distributions across different hardware configurations, thereby enhancing the adaptability of pre-trained models without compromising data privacy.
strength_weakness: strength_weakness: 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.
2. The proposed FedHP method is relatively straightforward and lacks deeper insights. Moreover, it does not show a significant performance improvement compared to FedAvg.
3. The experiments are not comprehensive enough. Given that this work aims to address the snapshot compressive imaging (SCI) problem, I suggest adding experiments to test the applicability of other SCI systems, such as Coded Aperture Compressive Temporal Imaging (CACTI).
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. Including reconstruction results obtained from these real-world systems is recommended.
questions: questions: 1. 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?
2. Although the proposed method outperforms other algorithms in terms of performance metrics, 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. Additionally, 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.
3. Furthermore, if possible, I hope the authors can also address the concerns I raised in the Weaknesses section.
limitations: limitations: Yes
paper_summary: paper_summary: Most existing reconstruction models in snapshot compressive imaging systems are trained using a single hardware configuration, making them highly susceptible to hardware variations. Previous approaches attempted to address this issue by centralizing data from multiple hardware configurations for training, but this proved difficult due to hardware heterogeneity across different platforms and privacy concerns. This paper proposes a Federated Hardware-Prompt Learning (FedHP) framework, which aligns data distributions across different hardware configurations by correcting the data distribution at the source, thereby enabling the trained model to adapt to multiple hardware configurations. The performance on existing datasets shows an improvement compared to previous popular training frameworks. Additionally, the authors have released their own created dataset and code.
strength_weakness: strength_weakness: 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.
2.Although good results were reported on simulated data, more results on real data should be included to evaluate the effectiveness of the proosed method.
questions: questions: 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.
limitations: limitations: The generalization to different hardware systems is crucial for deep learning based methods. The current form of this manuscript only reported on a small scale real dataset captured by several systems. A larger dataset captured by more systems is necessary to evaluate the method.
paper_summary: paper_summary: The paper addresses the challenges faced in snapshot compressive imaging (SCI) systems due to hardware shifts and the need for adaptability across multiple hardware configurations. By introducing a hardware-prompt network and leveraging federated learning, the framework enhances the adaptability and performance of SCI models across different hardware configurations.
strength_weakness: strength_weakness: 1. The literature review on federated learning (FL) heterogeneity in the Introduction section lacks comprehensiveness. There are numerous recent papers addressing heterogeneity in FL that are not cited here. Additionally, the references included are somewhat outdated. Including more current and diverse references would strengthen the review and provide a more accurate context for the study.
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.
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.
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.
questions: questions: 1. What is the rationale for using adaptors, and what is their function?
2. What network models are used in the comparison methods? It is necessary to clearly state the fairness of the validated methods.
3. The explanation for Figure 3 is not detailed enough. For example, what is "Patch"?
limitations: limitations: 1. 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.
2. 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.
paper_summary: paper_summary: The authors present a Federated Hardware-Prompt learning (FedHP) framework to address the fact that compressive snapshot spectral imaging devices may not be easily tuneable against changes in the coded aperture, and that in fact the said access to coded apertures may not be possible due to privacy reasons. The authors solve this by hardware prompt learning, which essentially learns from observing diverse coded aperture samples of all clients, regularizing the input data space and achieving the goal of coping with heterogeneity sourcing from hardware. The results show on a specific dataset improvement across all 10 samples in terms of spectral reconstruction quality. The comparison is primarily in the sense of federated learning approaches.

Typo: figure 3  caption -> colorblue shouldn’t be there
strength_weakness: strength_weakness: 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.

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.
questions: questions: 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.
limitations: limitations: 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.

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.

Comments: title: Response to Reviewer 1X5M
Response to Reviewer aeby
Response to Reviewer aeby
Response to Reviewer Jxiy
Response to Reviewer BdgG
Response to Reviewer aeby
Response to Reviewer 1X5M
comment: We appreciate the reviewer's valuable comments. We thank the reviewer's recognition of our rebuttal!
Thanks to the authors for addressing my concerns.
We thank the reviewer for acknowledging the response and for updating the recommendation. 

In line with the reviewer’s suggestions, we will revise the manuscript to outline potential future work in areas such as lensless cameras, LiDAR, HDR cameras, and CT reconstruction, especially concerning hardware-defined forward models. We will also include a comprehensive set of references on related hyperspectral and multispectral imaging strategies, incorporating forward models that could be explored in the future work.

We much appreciate the reviewer's insightful suggestions, which have significantly contributed to the enhancement of the paper's quality and scope!
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 sincerely appreciate the reviewer’s constructive comments and the time in reviewing the submission. Due to the limited rebuttal time and the significant workload involved in deploying new hardware and collecting additional data, we will investigate the potential applications such as Lensless cameras, LiDAR, HDR cameras, and CT-Reconstruction in future work. The proposed FedHP  can address the challenge of hardware heterogeneity in snapshot compressive imaging by integrating hardware prompt learning within a federated learning framework, highlighting the unique intersection of computational imaging and privacy-preserving federated learning. We would appreciate it if the reviewer can further assess the contribution of this work from the above perspective. We thank the reviewer again for the valuable comments and effort in reviewing this work!
We appreciate the reviewer's recognition of our response and support for our work!
We appreciate the reviewer's approval for our response and recognition of our work!
This reviewer appreciates very much the extra trials provided by the authors and believes this paper is borderline, leaning on the previous rating. The reviewer stresses how the newly brought evidence is valuable in making a point on the results with the studied dataset and architectures. 

However to increase the quality of the submission and to broaden its scope, which at present is very narrow, the reviewer invites the authors to investigate the "several potential related applications that might benefit from the proposed method, such as Lensless camera, LiDAR, HDR camera, or CT-Reconstruction". Showing that the technique applies regardless of the forward model, or on a range of forward models could allow the authors to increase the value proposition of this paper to a "Cooperative Hardware-Prompt Learning for Linear Sensing Models", broadening the results with more datasets such as those in Compressive CT-Reconstruction or MRI, for example.

The rating on this reviewer's side is exactly borderline (leaning to borderline reject).
The author has addressed most of my concerns. I will keep my score unchanged.
We appreciate the reviewer's constructive comments and the recognition of our rebuttal. To address the concern regarding the performance comparison between FedHP and FedAvg, we conducted  a statistical analysis using a paired t-test to compare the PSNR and SSIM values from FedHP and FedAvg. 

Specifically, we define the hypotheses as follows: (1) Null hypothesis ($H_0$): there is no significant difference in the PSNR and SSIM values between FedAvg and proposed FedHP. (2) Alternative hypothesis ($H_a$): there is a significant difference in the PSNR and SSIM values between FedAvg and proposed FedHP. 

We calculated the differences based on the averaged PSNR and SSIM values for each scene from both FedAvg and FedHP, resulting in ten differences values for PSNR ($d_{PSNR}$) and SSIM ($d_{SSIM}$). We performed the paired t-test using $t = \frac{\bar{d}}{s\_d/\sqrt{n}}$, where $\bar{d}$ denotes the mean of the difference values for either PSNR ($d_{PSNR}$) or SSIM ($d_{SSIM}$), $s\_d$ is the standard deviations, and $n$ is the number of the paired observations (e.g., $10$). 

We calculated the p-value upon the t-distribution for a two-tailed test using the formula p-value$= 2 \times P(T>|t|)$, where $P(T>|t|)$ denotes the probability that a t-distributed random variable with $n-1$ degrees of freedom exceeds the absolute value of the observed t-statistic. 

For PSNR, we observe $t=2.50$ and p-value$=0.034$. Since the p-value is less than the typical significance level of $0.05$. Therefore, we reject the null hypothesis ($H_0$) and conclude that there is a statistically significant difference between the PSNR values of FedAvg and FedHP. For SSIM, we observe $t=7.39$ and p-value$=0.00004$. The p-value of is significantly less than $0.05$, indicating a very strong statistically significant difference between the SSIM values of FedAvg and FedHP. The test results in PSNR and SSIM confirms that the performance gap between FedHP and FedAvg is statistically significant.  

We thank the reviewer for the valuable feedback, which has enhanced the quality of our submission. We will add the above discussions into the final version and look forward to any further suggestions from the reviewer. We sincerely appreciate the reviewer’s time and effort.
Thanks for the explanations. Concerns like real-system evaluation cannot be addressed in such a short period. I keep my initial rating.
We appreciate the reviewer's constructive comments and the recognition of our rebuttal. The presentation of *0.21dB improvement for C=5* in `A3.8` was a typographical error and should be consistent with the manuscript, which correctly states a *0.27 dB improvement for C=5*. Our intention was to highlight the advantage of the proposed FedHP and FedAvg as presented in Table 4 (a) for *C=5*. Additionally, we provide a detailed description of the privacy protection inherent in our approach as follows.

FedHP inherently addresses privacy from different perspectives. (1) **Hardware decentralization**: In the FedHP framework, real hardware configurations (e.g., real masks) remain confidential to the local clients. This design makes it difficult to reverse-engineer the pattern or values of the real mask without direct sharing.  (2) **Raw data decentralization**: FedHP maintains a private hyperspectral dataset for each client. The hyperspectral images are processed locally (e.g., encoding or data augmentation) and never leaves the client, thereby minimizing the risk of exposure. (3) **Training process decentralization**: FedHP only collects the local updates from the prompt network, which are then shared with the central server. The local updates are anonymized and aggregated without accessing underlying data, preventing any tracing back to the data source and thus protecting confidentiality. In Table 3, we quantitatively compared the performance of the proposed “FedHP” and “FedHP w/o FL” under privacy-constrained environments. FedHP demonstrates a  $0.6$ dB average improvement (e.g., $31.35$ *v.s.* $30.75$), showcasing its robust model performance and offering a significant privacy advantage that aligns with regulations restricting data sharing. We will add the above discussions into the manuscript. 

We sincerely expect our response can help solve the reviewer’s concern and expect a future discussion with the reviewer!
The authors have addressed all of my questions with their view on why FedHP is a substantial improvement for CASSI. This reviewer would like to thank them for the time spent in providing further experiments on the CACTI case, as well as the time spent to prepare the replies.

Focusing on the new Table T1, it appears that FedAvg and FedHP on CACTI perform very closely in PSNR and SSIM, notably in PSNR one could argue that the two distributions overlap. 

Have the authors performed an analysis of the residuals so that they would be able to ascertain whether the resulting models are statistically distinguishable from their residuals? I.e., not looking only at the mean and standard deviation but at whether the resulting distribution of residuals from FedHP is similar to that of FedAvg.

It appears that the results are in a close tie on the current dataset and, not having immediate clarity of how sensitive the dataset is, it is difficult to ascertain whether the proposed technique is sufficiently of impact.

The authors report they believe that their technique could be generalized to "several potential related applications that might benefit from the proposed method, such as Lensless camera ,LiDAR, HDR camera,  or CT-Reconstruction". This reviewer fully agrees that additional evidence on other modalities, potentially with larger datasets to ascertain the margin between the baseline and FedHP, could increase very significantly the strength and quality of the present paper submission.

Due to this, the reviewer would lean to leave my previous rating unaltered.
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.