Cooperative Hardware-Prompt Learning for Snapshot Compressive Imaging

TL;DR

This work addresses the challenge of cross-hardware learning in snapshot compressive imaging under strict privacy constraints. It introduces FedHP, a hardware-conditioned prompt learning framework that aligns input data distributions across heterogeneous clients while keeping reconstruction backbones frozen, enabling privacy-preserving multi-hardware collaboration. The authors also construct the Snapshot Spectral Heterogeneous Dataset (SSHD) from multiple practical SCI systems to evaluate cross-hardware generalization. Empirically, FedHP outperforms standard FL baselines and centralized approaches, achieving notable gains in PSNR/SSIM and robust spectral fidelity under diverse hardware configurations, with practical efficiency due to reduced backbone optimization and communication. This framework paves the way for scalable, privacy-conscious hardware-software co-optimization in computational imaging.

Abstract

Existing reconstruction models in snapshot compressive imaging systems (SCI) are trained with a single well-calibrated hardware instance, making their performance 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 configurations, 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

Figures (9)

• Figure 1: Comparison of hyperspectral reconstruction learning strategies. (1) The model trained with the single hardware (Prevalent treatment) hardly handles other hardware. Both (2) Jointly train and (3) Self-tuningwang2022modeling 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). All results are evaluated by unseen masks (non-overlapping) sampled from the practical mask distributions $\{P_1,P_2,P_3\}$. FedHP learns a prompt network $\Phi(\cdot)$ for cooperation.
• Figure 2: Learning process of FedHP. We take one global round as an example, which consists of (1) Initialize, (2) Local Update (Prompt), (3) Local Update (Adaptor), and (4) Aggregation. For each client, the reconstruction backbone ($\theta_c^p$), is initialized as pre-trained model upon local training dataset $\mathcal{D}_c$ and kept as frozen throughout the training. The prompt net upon hardware configuration, i.e., coded aperture, takes effect on the input data of reconstruction, i.e., $\mathbf{Y}^\mathbf{M}$. Adaptors are introduced to enhance the learning, where $\epsilon_c$ denotes the parameters of all adaptors.
• Figure 3: Reconstruction results on simulation data. The density curves compare the spectral consistency of different methods to the ground truth. We use the same coded aperture for all methods.
• Figure 4: Visualization of reconstruction results on real data. Six representative wavelengths are selected. We use the same unseen coded aperture for both FedAvg and FedHP.
• Figure 5: Reconstruction results on simulation data. The density curves compares the spectral consistency of different methods to the ground truth. We use the same coded aperture for all methods.