Paper: On the Role of Discrete Tokenization in Visual Representation Learning (ICLR-2024)

Reviewer comment (Writing):
I found the name "Mixture Tokenizer" a bit confusing since it almost seems adversarial in nature, and wasn't sure if this naming is common in another sub-field. It seems to me that this would be an adversarial/worst-case tokenizer; one that implicitly learns a class structure that is "orthogonal" to the desired one (eg, one that learns backgrounds instead of foreground objects in a perfectly balanced datasets). If so, I would suggest naming it something that denotes this quality better simply a "mixture tokenizier."

Author response:
A1. Thanks for your thoughtful suggestion! We rename "Mixture Tokenizer" with "Cross-class Tokenizer", which reflects the essence of the tokenizer, emphasizing that each equivalence class comprises elements from multiple classes.

Reviewer comment (Writing):
The introduction goes directly into specific methods that might not be known to the readers. Providing some smoother transition (eg, explaining what those methosd are doing first) could improve the readability of the introduction.

Author response:
A3. Thanks for your thoughtful suggestion! We add more explanations into the introduction part in the revision.

Reviewer comment (Evaluation):
The TCAS metric is only shown in Table 2, I think it would be nice to include it as a column in Tables 3 and 4 to see how well it explains variation in performance. Particularly, I am curious how much it changes with the choice of K in table 4, and whether that would explain some of the patterns there.

Author response:
A4. Thanks for your suggestions! Since Table 3 also contains several methods without using discrete tokenization, we do not include TCAS in Table 3. Following your suggestions, we organize the TCAS score and performance of all the tokenizers in Appendix E in the revision. It can be indicated from Figure 4 that there is a strong linear correlation between TCAS and Linear Probing Accuracy in general. This indicates that TCAS does offer a guidance on assessing the performance of a tokenizer.

Reviewer comment (Evaluation):
He et al (CVPR 2022) reported a fine-tuning performance of 83.6, yet Table 3 reports 82.9. Could you please comment on this discrepancy?

Author response:
A6. In [3], the fine-tuning performance of 83.6 is based on the checkpoints of 1600 epochs pretraining, where our setting is 200 epochs pretraining. Our 200-epoch result is consistent with the result in [4].

[3] He et al. Masked Autoencoders Are Scalable Vision Learners, CVPR 2022.

[4] Zhange et al. How Mask Matters: Towards Theoretical Understandings of Masked Autoencoders, NeurIPS 2022.

Reviewer comment (Evaluation):
Table 3 notes that K-MIM DINO achieves a linear probe accuracy of 67.4, which is significantly lower than 78.2 reported by Caron et al (ICCV 2021), while outperforming them on fine-tuning (83.8 vs. 82.8). I was curious why you think the model underperforms this much despite being given being trained using the equivalency structure learned by DINO.

Author response:
A7. Indeed, MIM methods will underperform contrastive methods in the linear probing accuracy and outperform them in the finetuning accuracy. There has been a previous work studying this phenomenon [5]. The following statements are quoted from [5]:

>Contrastive learning methods outperform MIM in linear probing and small model regimes Contrastive learning will capture the shapes of the images, which will  help recognize objects and distinguish images. In contrast, MIM preserves texture and the diversity of representation, which may not correlate as strongly with objects or content when compared to the emphasis on shapes.

Therefore, the phenomenon that K-MIM DINO will underperform DINO in the linear probing accuracy is due to the different training dynamics (MIM v.s. Contrastive Learning). 

[5] Park et al. What Do Self-Supervised Vision Transformers Learn? ICLR 2023.

Reviewer comment (Evaluation):
The results reported in Table 4 are quite interesting as they indicate that performance deteriorates quickly for larger token books, while a larger number of tokens seems to benefit DINO. Could you please comment on this result? I would be curious if TCAS could shed some light on this.

Author response:
A8. Indeed, the results reported in Table 4 indicate that the optimal clustering number for K-MIM DINO will be larger than K-MIM PIXEL. arises from the distinct separability of Kmeans performing in the pixel space and K-means performing in the DINO feature space [6]. To be specific, the dataset cannot be clustered well into real classes using K-means in the pixel space. Therefore, when increasing the number of clusters in K-means, it leads to the model overfitting noise or local features in the data, resulting in low-quality clusters and larger labeling error. Hence, the optimal clustering number for K-MIM is smaller. In the contrary, each class can be seperated well in the DINO feature space [7]. Therefore, a larger clustering number $K$ is optimal for the K-MIM DINO.

[6] Caron et al. Deep Clustering for Unsupervised Learning of Visual Features

[7] Caron et al. Emerging Properties in Self-Supervised Vision Transformers, ICLR 2023.

Reviewer comment (Experiments):
I think the finding that discretized features (HOG/DINO) provide better targets than raw pixels. DINO is an interesting target as it has been shown to be good local descriptor. I am curious if other features would be useful as well; eg, imagenet-supervised VIT features or MAE features. The first should be more aligned with the downstream task structure, while the latter is based on a model that seems to perform worse. I am curious if you had tried this or what your thoughts are on other features as targets for clustering.

Author response:
A9. As you suggest, we use DeiT features and MAE features in our K-MIM methods, where DEiT is an supervised ViT method trained on ImageNet. We pretrain DeiT and MAE models using ViT-S on ImageNet-100 for 800 epochs. The evaluation setting is consistent with that in our paper. Since DeiT is more aligned with the downstream task structure and MAE performs worse, we expect K-MIM DeiT will have a smaller TCAS socre and achieve better results than K-MIM DINO. Conversely, we anticipate that K-MIM MAE will have a larger TCAS score and underperform K-MIM DINO. The results are shown below:

|Method|TCAS score|Linear Probing Acc.|
|---|---|---|
|K-MIM DINO|0.15|59.7|84.7|
|K-MIM DeiT|0.09|64.8|83.1|
|K-MIM MAE|0.34|52.3|83.2|

As the table indicates, the TCAS score aligns well with the linear probing accuracy. K-MIM DeiT with the smallest TCAS score achieves the best linear probing accuracy and K-MIM MAE with the largest TCAS score underperforms K-MIM DINO. This conclusion ensures the reliability of our TCAS metric. For more detailed information, please refer to Appendix E, where we organize the TCAS scores and performances of various tokenizers.

Reviewer comment (Theory):
The theoretical analysis is only considered for two classes. I wonder if this can be extended into multiple classes.

Author response:
Sure. We add some discussions on the case of multiple classes in Appendix C. Similarly, we consider the three tokenization settings in this case and discuss how the tokenization influences the intra-class and inter-class connectivity, as well as the downstream error bound. In this situation, we still find that class-wise tokenization will have larger intra-class connectivity and cross-class tokenization will have larger inter-class connectivity. Therefore, similar to the two-class case in the paper, class-wise tokenization will enjoy a lower downstream error bound and cross-class tokenization will have a larger downstream error bound.

Reviewer comment (Theory):
Other the generalization to multi-classes, I have another question regarding the downstream error bound; it would be helpful if you can talk more in detail about the derivation perhaps in the appendix.

Author response:
We provide more detailed derivation process regarding the calculation of downstream error bound in Appendix B.

Reviewer comment (Novelty):
I think using discrete tokenization is not a mainstream of MIM. Representative methods, such as MAE, MaskFeat, and data2vec, demonstrate impressive performance without the discrete tokens. Thus, the contribution of the paper is hard to cover diverse variants of MIM.

Author response:
A1. Although discrete tokenization is not a default choice in MIM, it has been receiving wider and wider attention recently. In fact, one of the first MIM papers BeiT [1] uses tokenization, and so does iBoT [2]. Compared to those relying on raw inputs, the tokenization-based MIMs generally achieve much higher linear probing accuracy, which is known to benefit zero-shot generalization. 

Very recently, there is a phenomenological work [3] showing that tokenization is the key to visual generation. With the help of a good tokenizer, they attain the **best generation quality on ImageNet** simply with language model, surpassing previous sophiscated diffusion models. Similarly, visual tokens are also widely adopted by recent developments of visual foundation models. For example, equipped with visual tokenizers, MAGE [4] achieves **SOTA  generation ability within MIM methods**, and BEiT-3 [5] demonstrates **SOTA transfer performance on both vision-only and vision-language tasks**.

Therefore, the importance of tokenization is increasingly recognized in visual tasks. And our investigation from a theoretical perspective provides timely insights into the mechanism behind it. As the first theoretical work in this direction, we believe that it could help inspire more researchers to investigate the use of tokenization for visual representation learning.

[1] Bao et al. BEiT: BERT Pre-Training of Image Transformers, Arxiv 2106.08254.

[2] Zhou et al. Image BERT Pre-training with Online Tokenizer, ICLR 2022.

[3] Yu et al. Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation, Arxiv 2310.05737.

[4] Li et al. MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesis, CVPR 2023.

[5] Wang et al. BEiT Pretraining for All Vision and Vision-Language Tasks, CVPR 2023.

Reviewer comment (Theory):
According to theorem 1, image classification training could be the best way to downstream error bound. But, in practice, MIM works better than classification training in a lot of cases. Thus, I doubt the general applicability of this theorem and the metric (TCAS) on diverse MIM tasks.

Author response:
A2. We note that in existing literature, MIM has many downstream tasks, e.g., classification, detection, segmentation, etc. Among classification, there are two settings: linear probing (frozen features) and finetuning (tuning all features). Theorem 1 is established for the linear probing setting, which is the common evaluation setting in SSL theory, e.g., [6,7]. And indeed, the empirical results align well with our theory, since **MAE pretrained features** (no finetuning) underperform supervised learning by a large margin in this setting, as shown below.

|Method|Epochs|Accuracy (%)|
|---|---|---|
|Supervised|300|77.9|
|MAE|200 pretrain + 100 linear probing|55.4|

**Difficulty in theoretical analysis.** Meanwhile, as you pointed out, MIM features often perform better on other settings when finetuning is allowed. However, we are not aware of any existing theory on these tasks. Finetuning is hard for theoretical analysis since the MAE pretraining weights only serve as the initialization, and SGD training on NNs is highly non-contex. 

**Correlation between tasks.** Nevertheless, we note that the downstream performance of MIM on other datasets is often correlated well with its linear classification performance [8]. Therefore, although Theorem 1 is based on linear classification, we believe the theorem still provides valuable guidelines for improving its performance on other MIM tasks.

[6] Zhang et al. How Mask Matters: Towards Theoretical Understandings of Masked Autoencoders, NeurIPS 2022.

[7] Cao et al. How to understand masked autoencoders, Arxiv 2202.03670.

[8] Li et al. Benchmarking Detection Transfer Learning with Vision Transformers, Arxiv 2111.11429.

Reviewer comment (Experiments):
Experimental results are limited to a small dataset (ImageNet-100) and short training (200 epochs on IN-1k). I think it is not enough to validate the effect of K-MIM. Reported numbers on TCAS are also limited.

Author response:
A3. Thanks for your suggestions! We have added additional experiments involving 800 epochs of pretraining on ImageNet-1k using ViT-B as the backbone. The comparisons with baseline methods are presented in the table below:

|Method|Epochs|Linear Acc.|Finetuning Acc.|
|---|---|---|---|
|MAE|1600|68.0|83.6|
|BEiT|800|N/A|83.2|
|K-MIM PIXEL|800|69.4|83.6|

Notably, with only half the number of pretraining epochs, K-MIM PIXEL reaches the same finetuning accuracy with MAE and even ourperforms MAE in the linear probing accuracy. The results indicate that K-MIM PIXEL outperforms the baseline results in the classification tasks, ensuring the better learning ability of K-MIM. We also add this experiments in Appendix D in the revision.

Reviewer comment (Novelty):
Using DINO as a target representation is similar to [B]. Is there any relation between the distillation target and the TCAS metric? I think the paper would be better to cite [B] and add a discussion on it.

Author response:
A4. Our proposed method, MaskFeat, and dBot [9] can all leverage representations from a pretrained model like DINO as the prediction target, with variations in their processing methods. The distinctions are outlined in the following table:

|Method|Process Method|
|---|---|
|MaskFeat|Directly use the representations|
|dBot|Initially employs the representations, updating them per certain epochs|
|K-MIM|Use the K-means results of the representations|

Since the prediction target of dBot remains continuous, and its training is conducted in a multi-stage manner, our TCAS metric is designed specifically for discrete targets. As such, we will leave the analysis of distillation target of dBot to future works. We have cited [9] and incorporated the above discussion into Section 2 in the revision.

[9] Liu et al. Exploring Target Representations for Masked Autoencoders, Arxiv 2209.03917.

Reviewer comment (Experiments):
According to theorem 1, it looks like an image classifier would be the best tokenizer for MIM. What is the TCAS score and MIM performance when using an image classifier, such as DeiT, as a tokenizer?

Author response:
A5. We list the comparison of different tokenizers on ImageNet-100 using ViT-S with 200 epochs training in the table below for better illustration. The clustering number is set to $K=100$.

|Method|TCAS score|Linear Probing Acc.|
|---|---|---|
|K-MIM PIXEL|0.52|50.1|
|K-MIM DINO|0.15|59.7|84.7|
|K-MIM MAE|0.34|52.3|83.2|
|K-MIM DeiT|0.09|64.8|83.1|

The results indicate that employing DeiT as a tokenizer obtains the lowest TCAS score and outperforms other tokenizers in the linear probing accuracy, which aligns well with Theorem 1. For more detailed information, please refer to Appendix E, where we organize the TCAS scores and performances of various tokenizers.

Reviewer comment (Theory):
Section 3 is explained with tokenization for a group-of-tokens, i.e. $x_2 \in R^{n \times s}$. But, in Section 4, it seems the tokenization is conducted for a single token. Is it possible to generalize a theorem from the group-of-tokens case to the single-token scenario?

Author response:
A6. Indeed, in Section 3, our theory deals with general tokenizers that involve interactions among patches. However, in practice, it leads to a **exponential large space**, since each patch representation can be unique. However, when designing a practical metric in Sec 4, it is intractable for us to deal with. So we consider a bag-of-word model [10] for patch representations during the analysis in Sec 4, making it computationally tractable. Since patch representations are often local, this simplification (as a 1-gram model) also makes sense. In practice, we find the derived TCAS metric has good alignment with downstream performance, indicating that it could provide insights for the general cases. We have added clear explanations on these relationship in the beginning of Section 4.1.

[10] Zellig Harris. Distributional Structure, doi:10.1080/00437956.1954.11659520

Reviewer comment (Experiments):
Although K-MIM has better TCAS than other tokenizers, it is hard to say that K-MIM is specially designed for TCAS. Is there any discrete tokenization method optimized for the TCAS metric?

Author response:
A8. The guiding principle for TCAS is to assess the dissimilarity between the discrete token class and the true label. **Hence, the optimal choice for the TCAS metric is simply utilizing the true label as the tokenization.** In this context, our aim in the paper is to explore how to optimize the TCAS metric **when label information is unavailable**. Employing clustering methods such as K-means is likely to group samples with the same label together. Therefore, K-MIM is designed to optimize TCAS **under the condition of lacking label information.**

Reviewer comment (Experiments):
It would be better to show more ablation results about clustering numbers in Table 4.

Author response:
Sure. We conduct K-MIM PIXEL experiments on ImageNet-100 using ViT-Small, with clustering numbers $K=10, 25, 200$. The results (along with $K=50,100,1000$) are shown below.

|Clustering Numbers $K$| LP Acc | FT Acc|
|---|---|---|
|10|42.7|81.9|
|25|48.3|84.6|
|50|52.7|86.4|
|100|50.1|83.7|
|200|50.8|84.8|
|1000|50.6|84.4|

The results indicate that the performance peaks at $K=50$ and gradually declines as $K$ decreases from this optimal value. The diminished performance for small $K$ is due to the large inter-class connectivity and large labeling errors, where we have discussed in Section 3 that these factors lead to a larger error bound and consequently result in poorer performance.

Reviewer comment (Experiments):
Could you please show more ablation results about clustering numbers (e.g., K = 10, 25, 200) in Table 4?

Author response:
Sure. We conduct K-MIM PIXEL experiments on ImageNet-100 using ViT-Small, with clustering numbers $K=10, 25, 200$. The results (along with $K=50,100,1000$) are shown below.

|Clustering Numbers $K$| LP Acc | FT Acc|
|---|---|---|
|10|42.7|81.9|
|25|48.3|84.6|
|50|52.7|86.4|
|100|50.1|83.7|
|200|50.8|84.8|
|1000|50.6|84.4|

The results indicate that the performance peaks at $K=50$ and gradually declines as $K$ decreases from this optimal value. The diminished performance for small $K$ is due to the large inter-class connectivity and large labeling errors, where we have discussed in Section 3 that these factors lead to a larger error bound and consequently result in poorer performance.