Paper: A Spitting Image: Superpixel Transformers (ICLR-2024)

Reviewer comment (Experiments):
**Limited baselines:** The paper primarly evaluates three methods: `ViT-{S,B}/16` (the base model), `SPiT-{S, B}/16` (the proposed model) and `RViT-{S,B}/16`, which, from my understanding, is a baseline introduced in this paper using random Voronoi cells as tokens. I understand that improving classification accuracy is not the main goal of the paper hence the choice of baselines, but I do still think the paper would benefit from comparing to baselines which have a similar goal (i.e. either methods that propose more flexible tokenization, or more interpretable attention maps).

Author response:
> Limited baselines: The paper primarily evaluates three methods: ViT-{S,B}/16 (the base model), SPiT-{S, B}/16 (the proposed model) and RViT-{S,B}/16, which, from my understanding, is a baseline introduced in this paper using random Voronoi cells as tokens. I understand that improving classification accuracy is not the main goal of the paper hence the choice of baselines, but I do still think the paper would benefit from comparing to baselines which have a similar goal (i.e. either methods that propose more flexible tokenization, or more interpretable attention maps)

> On the aspect of interpretability, models such as DINO trained with finer-grained tokens (Emerging Properties in Self-Supervised Vision Transformers, Caron et al, CVPR 2021) would be a stronger baseline than ViT trained with coarser patch size

We agree with the reviewer that one must have a comparison with methods with similar goals to establish fair comparisons. While the reviewer points out several interesting works, they are geared towards scalability and efficiency in reducing the number of tokens, while our approach targets semantically aligned patches with pixel level granularity. Under this light, comparing against other methods doing tokenization with the goal of efficiency or other training paradigms; e.g., DINO (Caron et al. 2021), will not constitute a fair comparison. On the other hand, we agree that a more general evaluation of the method could provide insights about the separation of the tokenizer and the transformer backbone. 
To address the concern, we establish a taxonomy of adaptive methods in our related work section, and include another downstream task evaluation for unsupervised salient semantic segmentation, which compares our method to DINO to demonstrate the robustness of our proposal.

Reviewer comment (Novelty):
**Omission of works on flexible tokenization in introduction:** There have been many works investigating more flexible tokenization than the standard ViT "uniform square" assumption, which is not reflected in the introduction. For instance: **Token merging** approaches (e.g. "*`ToME`: Your ViT but faster, Bolya et al, ICLR 2023*") also alleviate the issues that "complexity and memory scales quadratically with the number of tokens in self-attention"; other works on **dynamic tokenization** (e.g. "*Vision Transformers with Mixed-Resolution Tokenization, Ronen et al, 2023*" or "*MSViT: Dynamic Mixed-Scale Tokenization for Vision Transformers, Havtorn et al, 2023*") also propose to incorporate mixed-resolution information directly at the input tokens stage and also show that ViT "can be succesfully trained under irregular tokenization"; `Swin` transformers and other **multi-scale works** gradually build multi-scale tokens, hence they do not have the issue that "the scale of the partitions are rigidly linked to the model architecture by a fixed patch size."

Author response:
> Other works on dynamic tokenization (e.g. "Vision Transformers with Mixed-Resolution Tokenization, Ronen et al, 2023" or "MSViT: Dynamic Mixed-Scale Tokenization for Vision Transformers, Havtorn et al, 2023") also propose to incorporate mixed-resolution information directly at the input tokens stage and also show that ViT "can be succesfully trained under irregular tokenization"

> My main concern is that it is not clear to me why using superpixels as tokens is advantageous with respect to other methods implementing some form of input-dependent tokenization, either in terms of performance or computational cost. The main gain seems to be in obtaining more fine-grained interpretability maps, but these are also only compared to the standard ViT-backbone (and the random voronoi cell baseline)

We kindly thank the reviewer for bringing the pioneering work by Ronen et al. (2023) and Havtorn et al. (2023) to our attention, as mixed-scale tokenization ultimately shares similar goals as our own. We emphatically agree that these methods are highly interesting, but note that these have their own distinct focus and objective. Mixed-scale tokenization approaches focus on more clever partitioning strategies with the goal of improving efficiency, scaling, and complexity of transformer models. Our approach, as noted by the reviewer, is to construct data driven partitions independent of the general transformer backbone that semantically align with the image content. In our opinion, **the key difference** that distinguishes our work from the Quadtree approach (Ronen et al.) as well as the gating mechanism in MSViT (Havtorn et al.) is that **we use pixel-level granularity with irregular (non-convex, non-square) partitions that fully adapt to the semantic content**. We see inherent benefits in different tokenization strategies for different problems, and note that given their modularity, these works instead contextualize and augment our own contributions, and demonstrate the importance of research into novel tokenization approaches.

Reviewer comment (Experiments):
**Overall concern on advantage of superpixels over other tokenization methods:** It is not clear why using superpixels as tokens is advantageous with respect to other methods implementing some form of input-dependent tokenization, either in terms of performance or computational cost. The main gain seems to be in obtaining more fine-grained interpretability maps, but these are also only compared to the standard ViT-backbone (and the random voronoi cell baseline).

Author response:
> Other works on dynamic tokenization (e.g. "Vision Transformers with Mixed-Resolution Tokenization, Ronen et al, 2023" or "MSViT: Dynamic Mixed-Scale Tokenization for Vision Transformers, Havtorn et al, 2023") also propose to incorporate mixed-resolution information directly at the input tokens stage and also show that ViT "can be succesfully trained under irregular tokenization"

> My main concern is that it is not clear to me why using superpixels as tokens is advantageous with respect to other methods implementing some form of input-dependent tokenization, either in terms of performance or computational cost. The main gain seems to be in obtaining more fine-grained interpretability maps, but these are also only compared to the standard ViT-backbone (and the random voronoi cell baseline)

We kindly thank the reviewer for bringing the pioneering work by Ronen et al. (2023) and Havtorn et al. (2023) to our attention, as mixed-scale tokenization ultimately shares similar goals as our own. We emphatically agree that these methods are highly interesting, but note that these have their own distinct focus and objective. Mixed-scale tokenization approaches focus on more clever partitioning strategies with the goal of improving efficiency, scaling, and complexity of transformer models. Our approach, as noted by the reviewer, is to construct data driven partitions independent of the general transformer backbone that semantically align with the image content. In our opinion, **the key difference** that distinguishes our work from the Quadtree approach (Ronen et al.) as well as the gating mechanism in MSViT (Havtorn et al.) is that **we use pixel-level granularity with irregular (non-convex, non-square) partitions that fully adapt to the semantic content**. We see inherent benefits in different tokenization strategies for different problems, and note that given their modularity, these works instead contextualize and augment our own contributions, and demonstrate the importance of research into novel tokenization approaches.

Reviewer comment (Novelty):
**Post-rebuttal note on novelty/impact:** As some other reviewers mentioned, there are still some reservations about the novelty/impact of the proposed method; however, in light of the rebuttal addressing several concerns and showing promising results on dense prediction tasks, the reviewer is raising their score from 5 to 6.

Author response:
Thank you for your detailed and constructive feedback, and for acknowledging the strengths in our revised manuscript. Your comments have been instrumental in improving our work.

We note your concerns regarding the novelty and impact of our method, and while we strived to address these aspects, we are intrigued by your perspective. Could you please elaborate on your specific reservations in this regard? Understanding your viewpoint will be beneficial for our continued research. We respect your expert opinion and look forward to any additional insights you might offer.

Reviewer comment (Theory):
**Notion of equivalence in Propositions 1 and 2:** What does it mean for the two operators to be equivalent? The proof of proposition 1 only seems to show that both operators have the same input and output dimensions. Generally, it is not clear what the conclusion of Section 2.4 is and how significant it is.

Author response:
> Notion of equivalence in Propositions 1 and 2: What does it mean for the two operators to be equivalent ? the proof of proposition 1 only seems to show that both operators have the same input and output dimensions. Generally, I am not sure what the conclusion of Section 2.4 is and how significant it is.

The reviewer addresses the perceived trivial nature of Proposition 1. In the provided proof, we ourselves comment on this while pointing out that it is necessarily instructive. *It anchors interpolating feature extractors to canonical ViT features since they are equivalent given similar patch sizes.* Proposition 1 is formulated in the context of Proposition 2, where we show that our proposed method for positional embeddings (which is designed for irregular patches) exactly coincide with learnable positional embedding in ViTs. This posits all models in the same framework; ViTs uses fixed square “superpixels”, RViT uses stochastic Voronoi “superpixels”, and SPiT uses our proposed superpixel tokenizer, but all are trained with the same framework. **The purpose is to clearly show the commensurability between our framework and feature extraction in canonical ViTs.** In light of the reviewer's comment, we relabel the proposition as a lemma.

Reviewer comment (Reproducibility):
**Number of tokens and batched executions:** In Section 2.2, it is said that "*We empirically verify that setting T = 4 produces comparable numbers of tokens to a ViT-B16*"; does this mean that each image has a different number of tokens in practice, and how does this impact batched execution during training/inference?

Author response:
> Number of tokens and batched executions: In Section 2.2, it is said that "We empirically verify that setting T = 4 produces comparable numbers of tokens to a ViT-B16"; 

The reviewer congruously brings attention to a clear gap in our original submission. These are clearly experimental results that are of interest to readers in our field. **Full experimental results on the validity of this claim have been included in Table F.1 in our revised submission.** 

> Does this mean that each image has a different number of tokens in practice, and how does this impact batched execution during training/inference ?

Exactly right. In our work, we discuss how we can regularize the number of superpixels, but not exactly control it, similar to nearly all other superpixel approaches. In our framework, all images are processed as a large supergraph consisting of images as a smaller subgraph, similar to the approach in GNNs. This allows us to handle variable token sizes effectively. To apply optimized computations in self attention frameworks (e.g. FlashAttention), we apply similar techniques as language transformers using masks. We realize that this point is not immediately clear in our original submission, and have revised Appendix C with more details to address this issue.

Reviewer comment (Evaluation):
**Evaluation on segmentation:** The choice of super-pixels seems particularly relevant for more dense tasks such as segmentation; this might be a more favorable task to evaluate on than image classification.

Author response:
> Evaluation on segmentation: The choice of super-pixels seem particularly relevant for more dense tasks such as segmentation: This might be a more favorable task to evaluate on than image classification.

While we acknowledge that segmentation is well aligned with superpixel tokenization, we note that classification models (alongside contrastive learning paradigms) are often considered more or less fundamental. Very few, if any practitioners, leverage segmentation models for classification. On the other hand, starting with pre-trained classification models for segmentation is not only common, but generally advisable. Furthermore, we had a scientific curiosity in determining the properties of different tokenization strategies in ViTs, in particular in relation to graph based methods. Given comparable results on classification tasks, improved interpretability and explainability, and our results on generalizability of tokenization across models, we considered our work as both interesting and novel enough for publication. 

However, we need to extend our sincerest gratitude to the reviewer for suggesting applications with segmentation tasks, since they prompted us to reevaluate how we could demonstrate feasibility for dense prediction tasks within the current framework. In considering zero-shot and unsupervised segmentation approaches, we decided that unsupervised salient segmentation could be well aligned with this goal, and we revised our submission to include experiments with this task in rev. Table 2 and Appendix G. **We happily report SotA results on three datasets for unsupervised salient segmentation tasks**, and report our results here for convenience.

| TokenCut Backbone | Postproc. | $\max F_\beta$ (ECSSD) | IoU (ECSSD) | Acc. (ECSSD) | $\max F_\beta$ (DUTS) | IoU (DUTS) | Acc. (DUTS) | $\max F_\beta$ (DUT-OMRON) | IoU (DUT-OMRON) | Acc. (DUT-OMRON) |
|-------------------|-----------|-----------------------|-------------|--------------|----------------------|------------|-------------|--------------------------|-----------------|------------------|
| DINO   | ❌        | 0.803                 | 0.712       | 0.918        | 0.672                | 0.576      | 0.903       | 0.600                    | 0.533           | 0.880            |
|                   | ✔️        | 0.874                 | 0.772       | **0.934**    | 0.755                | 0.624      | **0.914**   | 0.697                    | **0.618**       | **0.897**        |
| SPiT              | ❌        | **0.903**             | **0.773**   | **0.934**    | **0.771**            | **0.639**  | 0.894       | **0.711**                | 0.564           | 0.868            |

Reviewer comment (Experiments):
In the experimental section, the authors have provided thorough validation and analysis of the enhancement in feature extraction achieved by superpixel tokens in subsequent transformer modules. However, the paper lacks an analysis of SPiT efficiency. While Table 4 indicates that the superpixel algorithm used in this study performs comparably to state-of-the-art algorithms with substantially lower inference time, it is essential to explore whether the integration of superpixel tokens with subsequent transformer modules introduces a significant computational load. To address this, the authors should consider adding an additional set of comparative experiments to analyze the inference efficiency between ViT, RViT, and SPiT.

Author response:
We fully agree with the reviewer’s concerns that more dimensions for comparison will strengthen our paper. In response to the reviewers' concerns, in our revised submission, we included median images per second from training for all base capacity models to Table 1. This also relates to the reviewer’s next point about model performance, and we further address efficiency there.

Reviewer comment (Experiments):
Have you considered conducting a comparative analysis between your superpixel tokenization method and mixed-scale tokenization approaches [1] [2]? Mixed-scale tokenization can also alleviate the limitations of regular grid partitioning methods to some extent, reducing information loss without introducing excessive computational overhead. What advantages do you offer in comparison to them?

Author response:
We thank the reviewer for the suggestions about highly relevant literature that we missed before. As previously touched upon, **the prospects of comparative analysis between superpixel tokenization and mixed-scale tokenization approaches is both welcome and appealing**, particularly since **both frameworks are posed to tackle similar limitations** in ViT architectures. On the one hand, mixed tokenization approaches encourage more effective partitioning for better throughput, efficiency, and better scaling for larger images, while superpixel tokenization seeks to solicit more granular partitions to improve interpretation, explainability, and improved performance for dense prediction tasks  (see recent addition in Section 3 and Appendix G). In our revised submission, we do indeed add a contrastive discussion on these approaches, given that they are largely commensurable with our own proposed framework, and **very much look in the same direction as our own work in expanding the tokenization process towards a modular cohesive whole**. We consider these as significant contributions to the overall research on modular tokenization, and graciously accept the reviewer’s suggestion. We meticulously expand our related work and discussion with these works in mind.

Reviewer comment (Writing):
Some sentences are not complete, e.g., "We hope that our work inspires more research into the" in page 9.

Author response:
We sincerely thank the reviewer for the thorough evaluation.  We have remediated this issue, and further proofread the manuscript in the revised version.

Reviewer comment (Evaluation):
The runtime analysis is missing. Would replacing the square patches with superpixels for tokens require much extra runtime?

Author response:
This is an astute observation, and we wholeheartedly agree that this is of central importance. In our original submission, we only briefly touch upon the efficiency of our tokenization scheme in Section 3.2. (Quality of superpixels) and Table 4 (rev. Table 5), where we show that our superpixel tokenizer achieves very strong inference runtimes compared to other state-of-the-art approaches (0.0047 sec/img.).  However, this does not address runtime costs during training. In response to the reviewer’s concern, **we have added comparative runtime analysis from the training stage in Table 1**, alongside effective parameter counts for all architectures.  We thank the reviewer for their contribution towards improving the quality of our work.

Reviewer comment (Reproducibility):
Do the superpixel tokens remain the same in each layer or are hierarchically grouped together?

Author response:
There seems to be a misunderstanding on how the superpixels are used in our framework, which we are hopefully able to clarify. The **superpixels are constructed hierarchically in the tokenization step** using a graph-based approach, as per Section 2.2. After the tokenization step, we extract features using our proposed approach (non-hierarchically), which are then fed to the general transformer architecture. The hierarchical nature of the superpixel algorithm provides ample opportunity for continued research with multi-scale feature hierarchies, and we discuss this as further work (Section 4), and this is a route we are currently exploring.

Reviewer comment (Novelty):
This paper has limited novelty, as previous works [1][2] have already explored the combination of superpixels and Transformers.

Author response:
Firstly, we emphasize that **our main contribution lies in the decoupling of tokenization from the core transformer architecture**, a distinctive approach from the previous works cited. *Unlike these works which necessitate significant architectural alterations* for different tokenization strategies, our method *does not alter the standard transformer architectures*. This compatibility and its implications are empirically demonstrated in orig. Table 3 of our paper (rev. Table 4). We note that **we already acknowledge and discuss the work by Huang et al. [1] in our original submission**, differentiating our approach. However, in light of the reviewers comments, we expand our discussion on this important work in our revised submission.

Regarding Zhu et al. [2], it is pertinent to note that as far as we have found, the work was released as a preprint shortly after our submission deadline. **Given that our work thus predates Zhu et al. [2]**, which is considered novel by the reviewer, we believe **our work also merits recognition for its novelty on similar grounds**. We are committed to contributing novel insights to the field and hope this clarification addresses the reviewer's concerns regarding the novelty of our research.

Reviewer comment (Theory):
The superpixels are created using only low-level and predefined features in the first layer, which may cause permanent errors. In contrast, [1][2] use the features learned by the network.

Author response:
Unfortunately, the nature of these “permanent errors” is unclear to us, and we would greatly appreciate an elaboration. We hope to engage in dialogue to address the reviewer’s concerns and solicit constructive feedback for improving our work.

Reviewer comment (Evaluation):
Using only low-level and predefined features for superpixel creation may explain why the performance on image classification is worse than the baseline.

Author response:
We must respectfully disagree with the reviewer. Evaluation of our results across tasks **indicates that superpixel tokenization is largely comparable to standard ViTs for classification** without optimized architectures or training regimes , and performance varies across evaluated datasets. We note an increase over ImageNetReal, and a general drop in CIFAR100, which favors square tokenization due to low resolution. Moreover, **we explicitly define our goal as demonstrating feasibility for modular tokenization strategies**, and **our results indicate that decoupling the tokens not only reproduces the original performance, but also allows for flexible ways of combining the tokens in future work** (see prev. Table 3, rev. Table 4).

Reviewer comment (Reproducibility):
Do you use a differentiable online tokenization process?

Author response:
We do not explore differentiable tokenization in this work. **Our focus is on the introduction of the idea of tokenization as a general modular component to the transformer architecture.** We agree with the reviewers overall sentiment that differentiable tokenization is an interesting prospect, but point out that such approaches must be balanced with potential costs in compute, as discussed in our original submission. Moreover, we highlight that the tokenization process in LLMs is generally not differentiable, but still highly effective. However, our own work towards this goal within our proposed framework is in progress, and we are excited to see continued work exploring tokenization in vision models.

Reviewer comment (Experiments):
How does your method compare with conventional superpixel methods like SLIC, since you also use manually defined features?

Author response:
We kindly direct the reviewer’s attention to prev. Table 4 (rev. Table 5.), **which explicitly compares our approach with SLIC, and shows that SLIC turns out to be less effective while providing lower quality superpixel partitions**. As for the discussion of “manually designed features”, we note that our proposed feature extraction is essentially designed to be commensurable and extendable to other ViT architectures to facilitate further research towards new modular tokenization strategies, aligning with what we set out to show.