Paper: Learning Hierarchical Image Segmentation For Recognition and By Recognition (ICLR-2024)

Reviewer comment (Evaluation):
The runtime analysis is missing. Compared with the vanilla ViT, the proposed method involves superpixel generation and graph pooling. Does this require much more extra runtime?

Author response:
The computational cost of superpixel generation and graph pooling can be reduced by the decreased number of tokens required for computing self-attention. This advantage becomes more pronounced when using larger models, where self-attention blocks dominate the entire cost. To validate this, we analyze the latency of model inference and superpixel generation. Our system comprises a 32GB Nvidia Titan V GPU card and two Intel(R) Xeon(R) CPU E5-2630 v4 processors, totaling 20 CPU cores. We utilize the PyTorch machine learning framework with 24 workers, a batch size of 64, and an image resolution set to 224x224.

In our system, CAST-B achieves a lower average inference latency of 217 ms compared to ViT-B with 273 ms. SEEDS takes 73 ms to generate superpixels from the batch of images. However, we remark that the current SEEDS implementation is not fully optimized. Employing GPU implementation or parallelizing the process with more CPU cores can alleviate the bottleneck in superpixel generation. Furthermore, the cost of superpixel generation becomes less significant with larger models, which are commonly used in practice. We have included these discussions in Appx. B.1 of our revised manuscript.

Reviewer comment (Experiments):
Ablation study of using different superpixel methods. Is this proposed method sensitive to the finest segmentation?

Author response:
CAST performs well with reasonable superpixel methods. To verify this, we conducted an additional ablation study on superpixel methods. Specifically, we train CAST-S models using SEEDS (Bergh et al., 2012) and SLIC (Achanta et al., 2012) superpixels with MoCo objectives. We report the linear probing accuracy on ImageNet-1K using the models trained on ImageNet-1K and mIoU on Pascal VOC using the models trained on COCO, both before (left) and after (right) fine-tuning.

The table below compares the classification and segmentation performance of SEEDS and SLIC. Both superpixels performed reasonably well, with SEEDS outperforming SLIC by capturing better boundaries. We included deeper discussions and visual examples of superpixels in Appx. B.8 of our revised manuscript.

| Superpixel | Classification | Segmentation |
|--------|--------|-------|
| SEEDS  | 67.0   | 38.4/67.6 |
| SLIC   | 65.6   | 37.7/65.7 |

Nevertheless, we acknowledge that CAST indeed depends on the quality of superpixels. We discussed this limitation in Appx. B.7, illustrating that superpixels may not capture thin structures like light poles. Addressing this issue, perhaps by jointly learning superpixels, would be an interesting future direction.

Reviewer comment (Evaluation):
Does this require much more extra runtime compared with classical ViT.

Author response:
The computational cost of superpixel generation and graph pooling can be reduced by the decreased number of tokens required for computing self-attention. This advantage becomes more pronounced when using larger models, where self-attention blocks dominate the entire cost. To validate this, we analyze the latency of model inference and superpixel generation. Our system comprises a 32GB Nvidia Titan V GPU card and two Intel(R) Xeon(R) CPU E5-2630 v4 processors, totaling 20 CPU cores. We utilize the PyTorch machine learning framework with 24 workers, a batch size of 64, and an image resolution set to 224x224.

In our system, CAST-B achieves a lower average inference latency of 217 ms compared to ViT-B with 273 ms. SEEDS takes 73 ms to generate superpixels from the batch of images. However, we remark that the current SEEDS implementation is not fully optimized. Employing GPU implementation or parallelizing the process with more CPU cores can alleviate the bottleneck in superpixel generation. Furthermore, the cost of superpixel generation becomes less significant with larger models, which are commonly used in practice. We have included these discussions in Appx. B.1 of our revised manuscript.

Reviewer comment (Experiments):
Is this proposed method sensitive to the finest segmentation?

Author response:
CAST performs well with reasonable superpixel methods. To verify this, we conducted an additional ablation study on superpixel methods. Specifically, we train CAST-S models using SEEDS (Bergh et al., 2012) and SLIC (Achanta et al., 2012) superpixels with MoCo objectives. We report the linear probing accuracy on ImageNet-1K using the models trained on ImageNet-1K and mIoU on Pascal VOC using the models trained on COCO, both before (left) and after (right) fine-tuning.

The table below compares the classification and segmentation performance of SEEDS and SLIC. Both superpixels performed reasonably well, with SEEDS outperforming SLIC by capturing better boundaries. We included deeper discussions and visual examples of superpixels in Appx. B.8 of our revised manuscript.

| Superpixel | Classification | Segmentation |
|--------|--------|-------|
| SEEDS  | 67.0   | 38.4/67.6 |
| SLIC   | 65.6   | 37.7/65.7 |

Nevertheless, we acknowledge that CAST indeed depends on the quality of superpixels. We discussed this limitation in Appx. B.7, illustrating that superpixels may not capture thin structures like light poles. Addressing this issue, perhaps by jointly learning superpixels, would be an interesting future direction.

Reviewer comment (Experiments):
The major thing which I find missing is more in-depth discussion of the role of super-pixels here - I feel the ablation table is missing one critical row which is a ViT with a super-pixel (+conv net) embedding layer instead of patch extraction. This is also a possible explanation to the results in Table 5. There is a chance that super-pixels are the source of most of the improvements in the paper (quantitatively, at least, clearly there would be no clear hierarchical structure in this case) and this is not discussed enough.

Author response:
**[W1] Discussion on the role of superpixels**

We have shown that superpixels are not the sole reason for improvements. Please refer to the discussion in Table 3a, which reports (non-hierarchical) semantic segmentation performance. Here, superpixels enhance segmentation by providing better boundaries, as you suggested. However, combining hierarchical pooling leads to even greater improvements, enabling the model to grasp the structure of entire objects. This supports that both superpixels and pooling are essential, even for applications that do not require explicit hierarchy. We have further clarified these points on page 8 of our revised manuscript.

We have included the reorganized Table 3a below for your convenience. In this table, we report the mIoU (left) and boundary F-score (right) of segmentation before and after fine-tuning the models. Superpixels improve the original ViT by +5.1% in the boundary F-score. However, hierarchical pooling further enhances ViT + Superpixel by +6.8%.

| Model | Before tuning | After tuning |
|-------|-------|-------|
| ViT | 30.9/16.1 | 65.8/40.7 |
| ViT + Superpixel | 32.2/21.2 | 66.5/46.7 |
| ViT + Superpixel + Hierarchy (CAST) | **38.4**/**27.0** | **67.6**/**48.1** |

Reviewer comment (Theory):
The positional embedding pooling is a bit odd - super pixels can be with quite arbitrary shapes and there is a chance the resulting average PE over the super pixels would bare little connection to the actual "position" of the super pixel. Can the authors comment on that?

Author response:
**[W2] Pooling over positional embeddings**

Pooling over positional embeddings is as natural as pooling over token features, with the former characterizing spatial location and the latter characterizing appearance, which are two aspects of a segment. In fact, averaged positional embeddings serve as summary statistics for representing pooled segments, even when the corresponding location falls outside of a non-convex shaped segment pool. Prior works have also adopted a similar design choice. For example, SLIC (Achanta et al., 2012) calculates centroids by averaging LAB color features and XY locations within pixel clusters.

Reviewer comment (Theory):
While the authors claim there's a top-down element here (and arguably there is through training) I would argue it is not truly top-down at inference. The super pixels and segmentations do not depend and will not change based on what the top layers infer - the network has top-down pathway to inform them (other than gradients in training). This is of course fine, but it should be discussed.

Author response:
**[W3] Top-down recognition during inference**

Thank you for the insightful suggestion! As you noted, CAST only considers the top-down pathway during training, guided by the recognition objectives. This knowledge is encoded in the model and reflects the bottom-up pathway during inference. While this enables the model to learn the general top-down elements, the segmentations will not change based on what the top layers predict at inference time.

To extend CAST by incorporating the top-down pathway during inference, we employ test-time adaptation (TTA) [1], specifically TENT [2], with the classifier trained on top of the CAST backbones. We apply TENT to each sample to adapt the model and maximize its prediction confidence. As a result, CAST refines its object segments to align with its initial belief: If this image depicts the predicted class, which parts contribute to this prediction?

We have included the detailed approach and experimental results in Appx. D of our revised manuscript. We trained a dog vs. non-dog classifier on PartImageNet and evaluated the evolution of segmentation before and after TTA. Figure 16 showcases visual examples where TTA improves the segmentation of CAST by capturing object shapes more effectively (e.g., not missing legs in rows 1-3) while reducing attention to unnecessary details (e.g., window frames in row 3). The improvement is more substantial for challenging samples that the CAST originally fails on, increasing the mIoU from 69.5% to 76.6% (+7.1%) for samples with an original mIoU of less than 80%.

[1] Sun et al. Test-Time Training with Self-Supervision for Generalization under Distribution Shifts. ICML 2020.\
[2] Want et al. Tent: Fully Test-time Adaptation by Entropy Minimization. ICLR 2021.

Reviewer comment (Reproducibility):
How do segmentations are extracted from ViTs in the paper? It's not clear to me how the results in, say, Figure 4, were generated for ViT-S.

Author response:
**[Q1] How to generate segments from ViT?**

We generate segments from ViT by applying K-Means clustering to the final output tokens. We bilinearly upscale the feature map and then apply K-Means to the pixel-level features to align the segments with the input image resolution. Similar to CAST, we cross-reference clustering assignments across different levels to achieve hierarchical image segmentation. To maintain a consistent cluster hierarchy, we iteratively run K-Means, reducing the number of clusters by grouping the clusters from the previous iteration into coarser clusters. We have clarified this procedure on page 5 and in Appendix A.4 of our revised manuscript.

Reviewer comment (Reproducibility):
Is the PE added to the CLS token as well?

Author response:
**[Q2] PE added to the CLS token?**

Following MoCo-v3, we used fixed sine-cosine positional embeddings for patches and an all-zero vector for the CLS token. The all-zero vector matches the length of positional embeddings with tokens but has no impact. We have clarified this in the equation on page 4 of our revised manuscript.