Paper: LUM-ViT: Learnable Under-sampling Mask Vision Transformer for Bandwidth Limited Optical Signal Acquisition (ICLR-2024)

Reviewer comment (Reproducibility):
Code and pre-trained weights are not submitted. The reproducibility of this work cannot be checked.

Author response:
We acknowledge your concern regarding the reproducibility of our work, particularly the absence of code and pre-trained weights. We are actively working on organizing and annotating the code for LUM-ViT's model construction and will make it available on the current platform as soon as possible. Furthermore, following the publication of our paper, we plan to open source all related codes and checkpoints on GitHub.

This effort to provide detailed and well-documented code reflects our commitment to the scientific community's standards for reproducibility and transparency. We recognize the importance of your emphasis on these aspects and are dedicated to upholding these essential scientific principles.

Reviewer comment (Reproducibility):
The multi-stage training pipeline is tedious, which makes the whole technical route unreliable. What's worse, the finetuning details about stage 3 are missing. Other researchers cannot re-implement this complex approach.

Author response:
We appreciate your concern regarding the complexity of the multi-stage training pipeline of LUM-ViT. LUM-ViT, as a binarized-full precision hybrid model, requires distinct training configurations for each precision type. The binary layer, forming the foundation of our network, significantly influences subsequent layers, particularly mask training. Our attempts to train the binary layer and mask concurrently have been unsuccessful, leading us to adopt a multi-stage training approach.

We have experimented with an alternative strategy of binarizing first and then training the mask, but found that it did not yield results as effective as our current three-stage training process. Therefore, we consider the three-stage training necessary for optimal performance.

Despite the appearance of complexity, the training settings for the first and third stages are essentially the same. The actual process involves inserting a training phase with a different strategy for the binary layer in the middle. The total training duration is not excessively long, encompassing only 110 epochs, with the relatively unstable second stage requiring just 10 epochs. The detailed settings of the three-stage training are presented in the following table.

---
**Training Parameter settings.**

| config            | Stage 1-value              | Stage 2-value             | Stage3-value              |
| ----------------- | -------------------- | ------------------- | ------------------- |
| optimizer         | AdamW                | AdamW               | AdamW
| optimizer momentum($\beta_1,\beta_2$)|0.9, 0.999| 0.6,0.9999  |0.9,0.999
| base learning rate| 5e-4                 | 1e-4               |1e-4
| weight decay| 0.03                 | 5e-5               |0.03
| batch size| 4096                 | 4096               |4096
| learning rate schedule| cosine decay        | cosine decay      |cosine decay
| warmup epochs| 5                 | 1               |5
| frozen layers|none|mask|bi-conv-kernels,mask
| label smoothing| 0.1                 | 0.1               |0.1
| random erase| 0.25                 | 0               |0.25
| mixup| 0.8                 | 0               |0.8
| cutmix| 1.0                 | 0               |1.0
| drop path| 0.1                 | 0.1               |0.1


---

For a comprehensive understanding, we have provided detailed explanations of the three-stage training process in Section E of the Appendix. This should facilitate the replication of our approach by other researchers.

Reviewer comment (Writing):
The writing should be further improved, especially the mathematical notations in Section 3.3. The formula for binary compression is not formal.

Author response:
We thank you for highlighting the need for improvement in the writing, particularly regarding the mathematical notations in Section 3.3. We have thoroughly reviewed the manuscript, focusing on refining any verbose, confusing, or ambiguous expressions. Special attention was given to improving Figure 1 to enhance its clarity and intuitiveness.

In Section 3.3, we have revised the explanation of the binarization process. This includes a clearer presentation of the step function and a more rigorous definition of the binarization of the weight matrix $\boldsymbol{W}$, denoted as $\boldsymbol{W}^b$. We have clarified the relationship between $\boldsymbol{W}^b$ and $\boldsymbol{\Theta}$ and refined the calculation formula for s to be more precise. Additionally, we have included the formula for computing $\boldsymbol{\theta}$ during binarization.

We believe these modifications have made the content in this section more rigorous and easier to understand, addressing your concerns effectively.

Reviewer comment (Experiments):
The experiments are not sufficient and many critical comparisons are missing. For example, the backbone is ViT variants. However, the ViT is very computationally expensive because its computational complexity is quadratic to the input spatial size. This also embeds the real-time applications of HSI processing. In contrast, MST [1] or MST++ [2] are specially designed for HSI processing. They treat spectral feature maps as a token to capture the interdependencies between spectra with different wavelengths. Most importantly, they are very efficient with linear computational complexity regarding spatial resolution. Thus, it is better to add a comparison with the spectral Transformer in Figure 5.

Author response:
We greatly appreciate your recommendation of the MST and MST++ works. These studies indeed offer significant insights that can aid our research.

Improving the efficiency of signal acquisition is a key aspect of enhancing the real-time capabilities of hyperspectral imaging systems. In our work, the DMD signal acquisition system's primary time-consuming factor is the signal acquisition process. Therefore, our focus has been on enhancing signal acquisition efficiency, which directly translates to achieving lower undersampling rates in our experiments.

While MST and MST++ are commendable for their efficiency in processing hyperspectral data, their focus is primarily on data processing efficiency. Our experiments, however, are centered on evaluating performance under low undersampling rates. This difference in focus means that a direct comparison with MST and MST++ might not fully align with the specific objectives of our study. In Figure 5, we have chosen to compare methods that utilize learnable masks on patches, as with LUM-ViT, to more effectively validate our mask strategy. We acknowledge the significance of MST and MST++ in their respective areas, although their designs differ from the scope of our experiments.

Nevertheless, we recognize the importance of these works and the role they play in advancing data processing speed, a key direction for improving system real-time capabilities. We have introduced MST and MST++ in the introduction of our manuscript as significant references for future exploration.

Regarding the concern about the sufficiency of experiments, we have added experiments with Random-mask-ViT and Mag-mask-ViT as baseline methods to validate the effectiveness of our learnable mask strategy. These results are presented in Figure 4 and detailed in Section B of the Appendix. Additionally, we have included classification experiments on three hyperspectral remote sensing datasets to further validate LUM-ViT's capabilities in processing actual hyperspectral data. These experiments are presented in the newly added subsection 'The Hyperspectral Image Classification Experiments' and detailed in Section A of the Appendix.

Reviewer comment (Experiments):
The binarization mechanism is out of fashion. BiSCI [3] provides a specially designed binarized convolution unit BiSR-conv block to process HSI data cubes. It is also better to add a comparison with this new technique.

Author response:
We are grateful for your suggestion regarding the BiSCI work, which offers innovative approaches to the binarization process in models. The use of the Tanh function in this work is particularly inspiring and provides valuable insights for our research.

We recognize that there have been significant advancements in binarization strategies recently, indicating substantial progress in this field.

Our method primarily focuses on the design and experimental validation of the learnable mask. Given that DMD can only perform binary on/off modulation, we employed binarization techniques to adapt to DMD-based computing. Binarization, while not the central focus of our approach, was implemented using a basic binarization method due to this requirement.

We understand that improved binarization techniques can enhance training stability. In future work, exploring more effective binarization strategies will be a priority for us. The binarization strategy used in BiSCI, particularly intriguing to us, is mentioned in the subsection 'Kernel-level weight binarization'. We see great potential in incorporating such innovative techniques into our future research to further improve our model's performance.

Reviewer comment (Novelty):
The technical route and core idea in this paper is to accelerate the HSI processing, which is similar to Coded Aperture Spectral Snapshot Imaging (CASSI). Could you please analyze the differences, advantages, and disadvantages of these two systems?

Author response:
Your question about the similarities and differences between our approach and Coded Aperture Spectral Snapshot Imaging (CASSI) is intriguing. The primary distinction lies in the hardware systems used. Our system employs a DMD signal acquisition system, which is detailed in Section D of the Appendix. This system captures signals using a spectrometer with high spectral resolution, in our case, 2048 spectral channels. However, a spectrometer can only capture single-point pixel information at a time, requiring the DMD to reflect different areas of the target object to the spectrometer with various masks. This process, if done point-by-point, is time-consuming. Originally, this system was mainly used in single-pixel imaging based on compressed sensing principles, to reduce the number of required samplings.

In contrast, the CASSI system superimposes information from multiple spectral channels onto a single grayscale image using a disperser and reconstructs the original image's spectral channels algorithmically based on a random mask applied before dispersion.

Both systems aim to create a three-dimensional hyperspectral cube. The DMD system has a fixed spectral resolution, while CASSI has a fixed spatial resolution. Both require algorithmic optimization to acquire their respective to-be-capture dimensions. Deep learning methods can significantly contribute to this process.

When using compressed sensing algorithms in the DMD system, random masks displayed on the DMD are common for signal acquisition, followed by algorithmic reconstruction. This shares similarities with the random mask usage in CASSI. In our work, we replace the random mask in the DMD system with network parameters trained through deep learning, and we speculate that this approach might also have potential in the CASSI system.

Cost-wise, the primary expense in the DMD system is the DMD itself, which may be more economical compared to the CASSI system.


We hope this response has addressed your question and clarified the differences, advantages, and potential of these two systems.

Reviewer comment (Presentation):
The description of the entire system and methods is not clear and intuitive enough, and the figures are also misleading. The RGB image is used as an example in Figure 1, which does not reflect the characteristics of hyperspectral imaging.

Author response:
We acknowledge that the description of our method could have been clearer, and we appreciate your feedback highlighting this issue. We suspect the lack of clarity stemmed primarily from the insufficiently distinct introduction of the two main stages of our method: the Training Phase and the Real-World Application Phase. The Training Phase involves training and inference on electronic computing hardware, while the Real-World Application Phase involves real-world performance testing with the DMD signal acquisition system. These two phases are sequential.

To address this, we have made the following modifications:

1. **Terminology**: We have capitalized the first letters of “the Training Phase” and “the Real-World Application Phase” throughout the manuscript to signify their specific importance.

2. **Preliminaries Subsection**: We have added detailed descriptions of these two phases at the beginning of the "Preliminaries" subsection to help readers understand the distinction and sequence clearly.

3. **Figure 1 Revision**: We realized that the original differentiation between the two phases in Figure 1, marked by grey and blue colors, was not intuitive due to the low visibility of grey. We have now changed this to green and blue for clearer distinction and added annotations directly in the figure, in addition to the caption, to explain the meaning of these colors.

We also acknowledge that the initial lack of detail about the DMD hardware could have contributed to the confusion. To remedy this, we have included a comprehensive description of DMD in Section D of the Appendix.

Regarding the use of an RGB image in Figure 1 and its potential for misrepresentation, we have made revisions. We now overlay the original image with its color maps under different spectral masks to represent the inclusion of information across numerous spectral channels. This modification clarifies that the image is indeed hyperspectral, with annotations added directly on the figure for emphasis.

We hope these modifications adequately address the weaknesses you identified and improve the overall clarity and accuracy of our method's presentation.

Reviewer comment (Experiments):
Lack of experiments on real hyperspectral imaging. The author has emphasized hyperspectral imaging in the introduction section, but in reality, it has not been verified using real hyperspectral images. The performance of this method in real hyperspectral imaging tasks still needs to be discussed.

Author response:
Thank you for your insights regarding the scope of our experimental validation. We appreciate your point on the necessity of a broader range of tests to comprehensively assess the performance of our model, LUM-ViT, particularly in the context of hyperspectral imagery.

---

**Results on HSI classification.**

| Dataset           | OA (%) - before mask | OA (%) - after mask |
| ----------------- | -------------------- | ------------------- |
| Indian Pines      | 87.4                 | 86.2                |
| Pavia University  | 89.5                 | 88.6                |
| Salinas           | 99.4                 | 99.1                |

---

Recognizing the importance of diversifying our validation scenarios, we have conducted additional experiments on three widely-used hyperspectral remote sensing datasets: Indian Pines, Pavia University, and Salinas. These datasets are commonly utilized for hyperspectral pixel and region classification tasks, and we have tested LUM-ViT on these tasks. The results on these three datasets were promising, demonstrating LUM-ViT's capability in processing actual spectral data effectively.

We have added a new subsection titled 'The Hyperspectral Image Classification Experiments' in the 'Experiments' section of our manuscript to report the results of LUM-ViT on these three hyperspectral datasets. Due to space constraints, details regarding the datasets and parameter adjustments for these experiments have been included in Section A of the Appendix.

Once again, thank you for your insightful feedback, which has prompted us to further strengthen our manuscript with these additional experiments.

Reviewer comment (Experiments):
This acquisition system seems to have to work together with vit to obtain intermediate features of images. Can it reconstruct a complete hyperspectral image in a real-world environment?

Author response:
This is an intriguing question indeed. Theoretically, it's possible to achieve reconstruction of the input data by modifying the downstream task, as the requirements for classification and reconstruction tasks are similar in terms of model demands. The success of self-supervised learning in pre-training tasks, which perform well, supports this possibility.

However, we suspect that the quality of reconstruction may not be satisfactory under stringent under-sampling conditions. The reason LUM-ViT achieves satisfactory results in classification tasks at low undersampling rates is due to the sparsity of information required for these tasks in image or hyperspectral data. The amount of information needed for classification is significantly less than the total information in the original image, which makes the feature extraction process feasible. In contrast, reconstruction tasks demand a larger amount of information. At low undersampling rates, the information captured by the model may not be sufficient to accurately reconstruct the complete input image.

Despite these challenges, we still believe that pursuing this direction for reconstruction is viable and investigating this reconstruction process is an intriguing aspect of future research.

Reviewer comment (Reproducibility):
Technical details should be clear. How to achieve the learnable under-sampling mask? How is this learnable achieved? Is the learning accurate? Relevant visualization results should be provided.

Author response:
---
> **W-1: Technical details should be clear. How to achieve the learnable under-sampling mask? How is this learnable achieved? Is the learning accurate? Relevant visualization results should be provided.**

---

**REPLY**: Thank you for highlighting the need for clearer technical details regarding the learnable under-sampling mask in our LUM-ViT model. We acknowledge that our initial manuscript may not have adequately explained the implementation of the learnable mask.

To clarify, the learnable under-sampling mask in LUM-ViT is achieved through a binary decision mask with trainable parameters. During the training phase, these parameters determine the mask values, which influence the loss function values of two critical metrics: the accuracy of downstream tasks and the under-sampling rate. These parameters are updated through backpropagation, resulting in a mask that adheres to the under-sampling rate criteria and is optimally adapted for the downstream tasks.

It is important to note that, unlike DynamicViT or AdaViT architecture, which dynamically adjust their masks during inference, LUM-ViT requires pre-acquisition determination of the mask scheme. LUM-ViT focuses on under-sampling, which means reducing the number of optical computations and detections. For a single inference, optical computations are carried out before the entire system detects the target object and cannot adjust the masking strategy based on yet-to-know information on the target object. Consequently, LUM-ViT generates a mask that remains fixed post-training, *i.e.*, during the inference stage. This mask is designed to suit the validation set and downstream tasks through training on a compatibly distributed training set.

To further elucidate the principle of the learnable mask, we have revised our manuscript to replace the term 'learnable parameters' with 'trainable parameters'. This change emphasizes that the learnability is achieved through training. Additionally, we have substantially revised the subsection 'Learnable Under-sampling Mask' in the manuscript and revised the description in the introduction regarding the different requirements for learnable masks in under-sampling tasks versus those in pruning tasks.

---

**Top-1 Acc comparison.**

| Model             | Acc (%) - 10% mask | Acc (%) - 2% mask |
| ----------------- | ------------------------ | ----------------------- |
| Random-mask-ViT   | 74.6                     | 62.1                    |
| Mag-mask-ViT      | 76.2                     | 65.4                    |
| **LUM-ViT**       | **81.9**                 | **78.3**                |

---

We have also incorporated new experiments to validate the efficacy of the learnable mask approach. We introduced two alternative mask schemes: Random-mask-ViT, which employs a random mask, and Mag-mask-ViT, which uses a mask based on data magnitude. The results of these experiments, which show a significant performance gap between these methods and LUM-ViT, indirectly corroborate the effectiveness of our learnable mask method.

These experiments are detailed in the subsection 'The Training Phase Experiments', with results showcased in Fig. 4. Due to space constraints, we have included detailed descriptions of the Random-mask-ViT and Mag-mask-ViT models, along with previously base models like DU-ViT and CS-method, in Section B of the Appendix.

Furthermore, we have added a visualization analysis to investigate the internal mechanism of the learnable mask. Since LUM-ViT requires the mask to be fixed before sample detection during inference and does not allow for dynamic adjustment based on individual samples, we are unable to provide visualizations that compellingly illustrate non-essential parts of an individual sample being selectively masked. To compensate for this, we have visualized the masking results of LUM-ViT on convolutional kernels and patches at different mask ratios, using heatmaps and histograms. This analysis reveals LUM-ViT's tendency to select convolutional kernels under low under-sampling pressure and delve deeper into patch selection under high pressure. Due to space constraints, only the heatmap for LUM-ViT-2% is retained in the main text, with a detailed analysis provided in Section C of the Appendix.

Once again, we thank you for your insightful comments, which have significantly contributed to the improvement of our manuscript.

Reviewer comment (Experiments):
The experimental results seem insufficient. The author only conducted validation on the ImageNet-1k classification task, and other tasks should also be further explored.

Author response:
---
> **W-2: The experimental results seem insufficient. The author only conducted validation on the ImageNet-1k classification task, and other tasks should also be further explored.**

---

**REPLY**: Thank you for your insights regarding the scope of our experimental validation. We appreciate your point on the necessity of a broader range of tests to comprehensively assess the performance of our model, LUM-ViT, particularly in the context of hyperspectral imagery.

---

**Results on HSI classification.**

| Dataset           | OA (%) - before mask | OA (%) - after mask |
| ----------------- | ---------------------------------- | --------------------------------- |
| Indian Pines      | 87.4                               | 86.2                              |
| Pavia University  | 89.5                               | 88.6                              |
| Salinas           | 99.4                               | 99.1                              |

---

Recognizing the importance of diversifying our validation scenarios, we have conducted additional experiments on three widely-used hyperspectral remote sensing datasets: Indian Pines, Pavia University, and Salinas. These datasets are commonly utilized for hyperspectral pixel and region classification tasks, and we have tested LUM-ViT on these tasks. The results on these three datasets were promising, demonstrating LUM-ViT's capability in processing actual spectral data effectively.

We have added a new subsection titled 'The Hyperspectral Image Classification Experiments' in the 'Experiments' section of our manuscript to report the results of LUM-ViT on these three hyperspectral datasets. Due to space constraints, details regarding the datasets and parameter adjustments for these experiments have been included in Section A of the Appendix.

Once again, thank you for your insightful feedback, which has prompted us to further strengthen our manuscript with these additional experiments.