Paper: Perceptual Group Tokenizer: Building Perception with Iterative Grouping (ICLR-2024)

Reviewer comment (Experiments):
The experiment is not sufficiently comprehensive and the results are weak. Since a new model architecture is proposed, the performance under both supervised learning and self-supervised learning should be presented to show its’ universal. In performance comparison, the listed comparable architectures should be compared with the same pre-training/supervised-training strategy, a similar amount of parameters. It will be better to do experiment with more model sizes (e.g., PGT-B, PGT-S, PGT-Ti or more) and compare with counterparts under different sizes. The performance on downstream tasks are not clear; although segmentation performance is reported, more comparisons with relative methods are not sufficient.

Author response:
**W1.1. Multiple model sizes**

We agree and are working on the results with various model sizes, the performance of variants will be included in the final version. Providing one data point here, we develop a PGT-S with 34M parameters, the model is still in early training (1/3 course) and show similar performance (76.8%) with PGT-B (77.3%)

**W1.2 Results and others**

We have updated our results in the paper. The table 1 is arranged for more clear fair comparison. Note that the performance of 80.3% is a strong result under the perceptual grouping paradigm, especially, we show that PGT has a much lower memory cost in inference.

Reviewer comment (Evaluation):
Besides the number of parameters, for fair comparison, the computational cost (Gflops or MACs) should be shown. The inference time is missing, which is important for evaluating models’ efficiency.

Author response:
**W2. Profiling computational costs and stats**

Following the reviewer’s suggestion, we have updated our paper with following modifications: peak memory usage in table 3, inference time in A.2.3, and Gflops in A.2.4. Especially to highlight, PGT has a very low memory cost compared to ViT under fair comparison, summarized in table 3.

Reviewer comment (Theory):
Do the group tokens in each block generate independently? If not, maybe the relevant description is ambiguous. If so, why don't the group tokens take its’ output of the previous block as the next block's input?

Author response:
**Q1. Independent group tokens**

This is a great question. Yes, each layer’s group tokens are generated independently. Our paper aims to show that grouping operations can be used as a more general operator to exchange information among input tokens, compared to self-attention. Both operations are per-layer refinement. Adding cross layer connections can complicate the process. 

Regardless, we indeed have tried reusing the previous layer’s group tokens as a variant (e.g., residual connections, or cross attention), and found that the performances are similar. This indicates that the grouping results are already stored in the refined patch token embedding vectors and can be easily recovered in the new round.

Reviewer comment (Novelty):
The proposed method is essentially a combination of ViT and slot attention (or its token grouping variants), and the concept of token grouping has been extensively studied in prior work. Relevant examples can be found in the related work section of ToMe [1] and CoCs [2].

Author response:
**W1. Concerns about the novelty of the paper**

We disagree with this point. The mentioned papers are relevant, but we are delivering fundamentally different messages in our paper, which we clarify below:

- The core message

We replace the self-attention operation with grouping operations, showing that grouping operations can perform feature learning in self-supervised learning. This alone separates our paper with other works [1,3,4,5,6], which still rely on ViT’s self attention operations or convolution layers. 

Specifically, [1,4,5] are ad-hoc adaptations for pruning ViTs and orthogonal to our work. The techniques like patch token drop can directly be applied to PGT input tokens x as an efficiency improvement method. CAST[6] uses object proposals from traditional vision methods as input, this makes it not comparable to most backbones, including PGT. [2] is indeed relevant, but focuses on supervised learning, uses non-standard operations (fixed center feature pooling, aggregation with sigmoid activations, patch pooling for reduction, etc.).  Instead, in our paper, we are the first to cleanly use *iterative grouping to replace the self attention* operation and show a strong performance.

We also show that the grouping tokens as communication channels among input tokens is a more general form of information-exchange operator. This generalizes self-attention and can potentially lead to more advances in vision models.

We do not think the mentioned works [1-6] deliver these messages.

- Self-supervised learning

Perceptual grouping is a general vision framework with a long history, and the root for many vision concepts and algorithms, *especially unsupervised*. There is a natural connection between pixel similarity measurement, grouping, and providing training signals from images themselves. Many classical vision tasks, such as segmentation, tracking, and 3D multiview problems can be formulated as grouping processes. Showing grouping works, especially under self-supervised learning framework, can be a critical stepstone to unify many vision tasks. 

- We do not combine ViT and Slot attention.

 The main components of ViT are tokenization, self attention, and mlp, where self attention is the core. We use tokenization, grouping and mlp. There is no ViT in our model.

Slot attention combines the grouping procedure with ConvNet, showing the feature maps from the last layer of convolutional nets can be grouped. We instead show grouping can be used as a core operation to drive the entire feature learning.


We appreciate that the reviewer raises these relevant and inspiring works, and have added paragraphs discussing them in detail in section 2.

Reviewer comment (Evaluation):
In addition to the technical novelty, the empirical merit of the proposed method is unclear in terms of performance, as it is not better than prior works, such as DINO.

Author response:
**W2. Empirical benefits**

  Note that we do not aim to propose an all-round model that beats every possible aspect of state-of-the-arts from all related works, and most papers do not either. As discussed above, the mentioned papers [1,4,5] are orthogonal to our work and can be applied to PGT. The intriguing part of PGT is that the performance, memory efficiency, adaptive computation and interpretability are all in a single model cleanly driven by one principle.

We hope the explanation and response above address the reviewer’s questions.

Reviewer comment (Experiments):
The motivation is not clear. Since we have powerful vision transformers already, what is the advantages of the proposed perceptual group tokenizer. The authors claim that the perceptual group tokenizer have good properties such as adaptive computation without re-training and interpretability. However, I don't see the experimental results or the visualization that can provides proof of this claim.

Author response:
As mentioned in the introduction, our motivation and contribution is on the fundamental mechanisms of vision backbones. We would like to both refer the reviewer to our first half of introduction, and explain the motivation here. 
- Fundamental paradigms. 

Perceptual grouping in computer vision has deep roots and many vision concepts/methods are developed under that paradigm, for example, segmentation, object tracking, action recognition, etc. Showing that by simply applying one principle - perceptual grouping - can lead to state-of-the-art performance, is very encouraging. We believe this principle can be transferred to related fields and applications given their deep connections with grouping. 

- link to ViT

As discussed in the paper, PGT with grouping can be considered as a more generalized form of ViT. ViT will always have fixed number of grouping slots (equal to number of patches), while PGT separates the concept out and allows for flexible customization on number of groups. This makes it more generalizable to a broader class of vision tasks which require flexible designs in number of objects/groups (e.g., tracking). We also build connections between grouping and self-attention (section 3.4), serving as a generalization in the operator design and an explanation on why ViT works. This perspective can spur more variations of grouping ops to go beyond self attention and advance the vision backbone designs. 

- Generating better tokens. 

One big advantage of ViT over ConvNet is its meaningful tokens[b], although it has been consistently shown that ConvNet can achieve similar performance[b]. The tokens can be easily aligned with language tokens in multimodal training[b]. Perceptual grouping models have the potential to produce even better tokens by design. The grouping process is constructed to produce a set of tokens that flexibly summarize the “detailed small patches” into tokens with more semantic meanings. 

- Efficiency.  

As shown in table 3, the peak memory usage of perceptual grouping models is much less compared to ViT. The memory advantage is critical in larger resolutions, videos or multiview images.

**Q1.2 Adaptive computation and grouping results**

Our results of adaptive computation are in section 4.3, specifically table 3. 

The grouping interpretability results are shown in figure 5 and 6. For example, in figure 6, each large attention image corresponds to a group token, an image is parsed into smaller meaningful regions, such as apple, jar and handle, or camel, camel legs and person. More explanation is in section 4.5 and figure 6 caption.

Reviewer comment (Experiments):
The performance is still concerned. Because we should also focus on the accuracy despite some good properties, the results shown in Table 1 demonstrate that the method cannot beat the baselines. Moreover, the state-of-the-art methods proposed in 2023 are not compared.

Author response:
We have updated the results of our model. We better aligned with DINO’s preprocessing, where we previously omitted random blur. The current results have shown on par or better on top-1 accuracy.  Note that this performance is achieved with much less inference memory cost. 

We also would like to point out that only the lower half of table 1 are  fair comparisons. We fully follow DINO’s framework and only compare the backbone performance. The upper half are cited as reference for other backbone architectures. We cite works within the similar time range of DINO. With better loss functions and training procedure, PGT can potentially have better performance.

Reviewer comment (Experiments):
Only ViT-B level (70-80M parameter) model is reported. It would justify the effectiveness if the proposed architecture works when scaling the model size up.

Author response:
**W1. Scale up model sizes**

We agree. We are working on more experiments to scale up the model to 300M or larger level sizes and will add the results to the final version (the training takes some time).

Reviewer comment (Evaluation):
The computation cost, peak memory usage, and inference speed comparison with ViT-B are not reported. It would be informative for readers how fast the PGT is since PGT doesn't have memory/computation demanding self-attention operations.

Author response:
**W2. Profiling PGT’s computational stats & Q3**

Thank you for this suggestion! We have updated the draft and added peak memory usage in table 3 and inference time in appendix section A.2.3. For easier read, the results are also listed here:

- Peak memory usage

We show the percentage of peak memory usage in PGT compared to ViT-B with the same patch size (4×4).  Even with a large number of group tokens, due to less number of heads and total attention tokens, we still only use a much smaller percentage of memory compared to the ViT-B backbone. Note that the usage is obtained using forward inference graph, as in practice the underlying complex hardware optimizer is a less accurate measurement and largely varies across infrastructures.


| # group tokens        | 16  | 32  | 64  | 128 | 256 | 384 | 512 | 768  | 1024 | ViT-B |
|-----------------------|-----|-----|-----|-----|-----|-----|-----|------|------|-------|
| Peak memory usage (%) | 4.6 | 4.6 | 4.6 | 4.6 | 4.6 | 6.1 | 8.2 | 12.2 | 16.3 | 100   |

- Inference time

Similarly, we profile our model’s inference time against ViT-B/4. Our model is built upon a complex infrastructure with XLA and other hardware optimizers. We find that ViT has ~680 im/sec/core inference speed, while PGT has ~640 im/sec/core. Note that PGT has a heavier process (3 steps of iterative grouping), and can still maintain similar speed with ViT. With fewer grouping iterations (1 and 2), our model can achieve ~710 im/sec/core and ~820 im/sec/core.

Due to the underlying hardware speed optimizers, changing the number of group tokens currently leads to negligible influence. Theoretically, with less number of group tokens, the inference time should still be reduced. We will look into how to profile the model in a more toy and simpler environment.

Reviewer comment (Experiments):
The model uses patch size 4x4, which means the visual backbone only downsamples the image by 4. Intuitively, this model should be good at dense prediction tasks that require high feature resolution, e.g. semantic segmentation and object detection. The authors reported the results of semantic segmentation on ADE20K in Section 4.4, but it only outperforms the ViT-B by a smaller margin. I understand that the segmentation architecture is different. So it would be interesting to compare ViT-B vs PGT-B with the same segmentation architecture (linear classification layer).

Author response:
**Q1. Segmentation**

We have an updated result of 45.1% for PGT. Using ViT with a linear layer, we got similar results (44.2). We completely agree on that having more dense patches should indeed lead to higher performance. One very likely reason is DINO’s loss is a whole-image holistic loss, where parts of the same image are always matched with each other. The images in ImageNet, although curated, still contain multiple objects in most of them. This loss design might limit the power of our grouping backbone. We think a dense loss function that uses “pixel-group assignment” will further unleash the performance.

Reviewer comment (Experiments):
Since PGT is self-attention-free, the computation cost is not quadratically increasing with the input resolution. But it is still reasonable to compare to ViT under the same resolution, for example, 16x16 patch size and 8x8 patch size. I am wondering whether authors have done this ablation.

Author response:
**Q2. 8x8, 16x16**

Yes, we have been keeping track of model performance in different resolutions. Since one core of the grouping process is the ability of automatically grouping patches and generating more meaningful “regions” (compared to large strided patchification), we mainly test the model on small patches to allow bigger flexibility in the binding procedure. A simple attempt on 8x8 and 16x16 lead to 79.2 and 77.0. PGT-B uses 384 as patch embedding vector dimension and 384 as the hidden dimension for MLP (ViT uses 768 and 3072). Scaling it to the same dimension lead to a 0.4 improvement on each model.

Reviewer comment (Experiments):
As mentioned in the weakness section, it would be interesting to have the computation cost of PGT. As Table 1 shows, as we increase the number of group token from 256 to 768, the linear probe accuracy increases from 79.3 to 79.7. How about less than 256 tokens and more than 768 tokens? It would be insightful to have a graph of number of tokens vs accuracy and inference speed.

Author response:
**W2. Profiling PGT’s computational stats & Q3**

Thank you for this suggestion! We have updated the draft and added peak memory usage in table 3 and inference time in appendix section A.2.3. For easier read, the results are also listed here:

- Peak memory usage

We show the percentage of peak memory usage in PGT compared to ViT-B with the same patch size (4×4).  Even with a large number of group tokens, due to less number of heads and total attention tokens, we still only use a much smaller percentage of memory compared to the ViT-B backbone. Note that the usage is obtained using forward inference graph, as in practice the underlying complex hardware optimizer is a less accurate measurement and largely varies across infrastructures.


| # group tokens        | 16  | 32  | 64  | 128 | 256 | 384 | 512 | 768  | 1024 | ViT-B |
|-----------------------|-----|-----|-----|-----|-----|-----|-----|------|------|-------|
| Peak memory usage (%) | 4.6 | 4.6 | 4.6 | 4.6 | 4.6 | 6.1 | 8.2 | 12.2 | 16.3 | 100   |

- Inference time

Similarly, we profile our model’s inference time against ViT-B/4. Our model is built upon a complex infrastructure with XLA and other hardware optimizers. We find that ViT has ~680 im/sec/core inference speed, while PGT has ~640 im/sec/core. Note that PGT has a heavier process (3 steps of iterative grouping), and can still maintain similar speed with ViT. With fewer grouping iterations (1 and 2), our model can achieve ~710 im/sec/core and ~820 im/sec/core.

Due to the underlying hardware speed optimizers, changing the number of group tokens currently leads to negligible influence. Theoretically, with less number of group tokens, the inference time should still be reduced. We will look into how to profile the model in a more toy and simpler environment.

Reviewer comment (Evaluation):
In Mask Autoencoder paper, researchers find that higher linear probing accuracy may not necessarily stand for better representation. It would be also interesting to compare with MAE-ViT against PGT under the fine-tuning setting.

Author response:
**Q4. Fine-tuning**

This is a great question. We fine-tuned our PGT and achieved 82.3% (DINO has 82.8%). Interestingly, we find that the standard data augmentation trick mixup negatively affects the model’s performance. Mixup combines two images in a linear combination way and will blur the images. This seems to interfere the grouping process. PGT only uses randAugment[a] for finetuning. We believe perceptual grouping models have their own augmentation pipeline for achieving the best performance.

[a] RandAugment: Practical automated data augmentation with a reduced search space. Ekin D. Cubuk, Barret Zoph, Jonathon Shlens, Quoc V. Le