On the Role of Discrete Tokenization in Visual Representation Learning
Authors: Tianqi Du
Yifei Wang
Yisen Wang
Venue: ICLR 2024 spotlight (2024)
Keywords: Self-supervised learning
Masked image modeling
Discrete visual token
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: In the realm of self-supervised learning (SSL), masked image modeling (MIM) has gained popularity alongside contrastive learning methods. MIM involves reconstructing masked regions of input images using their unmasked portions. A notable subset of MIM methodologies employs discrete tokens as the reconstruction target, but the theoretical underpinnings of this choice remain underexplored. In this paper, we explore the role of these discrete tokens, aiming to unravel their benefits and limitations. Building upon the connection between MIM and contrastive learning, we provide a comprehensive theoretical understanding on how discrete tokenization affects the model's generalization capabilities. Furthermore, we propose a novel metric named TCAS, which is specifically designed to assess the effectiveness of discrete tokens within the MIM framework. Inspired by this metric, we contribute an innovative tokenizer design and propose a corresponding MIM method named ClusterMIM. It demonstrates superior performance on a variety of benchmark datasets and ViT backbones. Code is available at \url{https://github.com/PKU-ML/ClusterMIM}.

Published as a conference paper at ICLR 2024
ON THE ROLE OF DISCRETE TOKENIZATION IN VI- SUAL REPRESENTATION LEARNING
Tianqi Du Yifei Wang** Yisen Wang
I National Key Lab of General Artificial Intelligence,
School of Intelligence Science and Technology. Peking University
2 School of Mathematical Sciences, Peking University
3 Institute for Artificial Intelligence, Peking University
ABSTRACT
In the realm of self-supervised learning (SSL), masked image modeling (MIM) has gained popularity alongside contrastive learning methods. MIM involves re- constructing masked regions of input images using their unmasked portions. A notable subset of MIM methodologies employs discrete tokens as the reconstruc- tion target, but the theoretical underpinnings of this choice remain underexplored. In this paper, we explore the role of these discrete tokens, aiming to unravel their benefits and limitations. Building upon the connection between MIM and con- trastive learning, we provide a comprehensive theoretical understanding on how discrete tokenization affects the model's generalization capabilities. Furthermore. we propose a novel metric named TCAS, which is specifically designed to as- sess the effectiveness of discrete tokens within the MIM framework. Inspired by this metric, we contribute innovative tokenizer design and propose a cor- responding MIM method named ClusterMIM. It demonstrates superior perfor- mance on variety of benchmark datasets and ViT backbones. Code is available at htt
com/ P-U-MMLLL MLLCllstteri
INTRODUCTION
Solf-Supervised Learning (SSL) has recently emerged promising paradigm for learning mean-
ingful data representations without access to labels. Besides the popular contrastive learning meth-
ods (Chen et al. 2020 Chen & He 2020 Grill et al. 2020 Zhuo et al. 2023 Wang et al.) 2021
2023 Cui et al. |023] Zhang et al. 2023a b). there is growing interest in masked image modeling
(MIM) techniques. MIM requires masking parts of input images and subsequently attempting their
reconstruction using the remaining unmasked parts. Notable examples like MAE (He el al., 2022).
BEiT (Bao et al. 2002). PeCo (Dong al., (023). MaskFeat (Wei al., 2022). and MAGE
et al.
2023) have demonstrated state-of-the-as art performance on various downstream tasks.
Within the MIM domain, various reconstruction targets exist. For instance, MAE (He et al. 2022) employs the raw pixel values of the unmasked parts as its reconstruction target and MaskFeat Weii jet al. (022) leverages features from other pretrained models as reconstruction target, whereas some alternative methods adopt visual tokenizers to generate the reconstruction target (Bao et al.) 2022 Dong et al. (023). Such tokenizers convert image patches into predefined discrete tokens. As an illustration, BEIT (Bao et al. 2022) deploys visual tokens generated by discrete tokenizer known as dVAE (Vahdat et al., 2018). PeCo (Dong et al. 2023) refines the token quality by introducing a tokenizer enriched with perceptual loss. More related work about MIM can be found in Appendix Nevertheless, we notice that different tokenization schemes may bring quite different performance. For example, Table suggests, PeCo with perceptual tokenizer and MaskFeat with HOG targets outperform MAE, whereas BEiT with dVAEE and VQGAN tokenizer underperform MAE using raw pixel. These observations naturally raise the question:
What is the role of tokenization in MIM? How does it affect the downstream generalization?
"Equal Contribution. Yifei Wang has graduated from Peking University, and is currently a postdoc at MIT. "orresponniig Author: Yisen Wang (yisen.wang @ @ku.edu.en).

---

Published as a conference paper at ICLR 2024
Table 1: Linear probing and fine-tuning accuracies (%) of several MIM methods using different tokenizers pretrained on ImageNet100 for 200 epochs. The architecture is ViT-small following Touvron et al. (202T). The comparison baseline the tokenizer of identity function in MAE.
Tokenizer
Linear Probing Acc.
Fine-tuning Acc.
Identity function (MAB)
45.9
81.8
dVAE Vahdat al. 2018
41.2 (-4.7)
80.2 1.6)
VQGAN (Esser et al
2021
44.3 (-1.)
80.9 -099)
Perceptual tokenizer (Dngg et al. 2023
53.2 (+7.3)
83.6 (+188
Maskfeat (HOG targets)
(Wer et al. 2022)
49.1 (+3.2)
82.8 (+1.0)
K-means (ours)
52.7 (+6.8)
86.4 (+4.6)
To answer these questions, we leverage the graph perspective of MAE (Zhang et al. 2022 to dissect the influence of different discrete tokenization schemes the downstream generalization. Zhang iet al. (2022) have demonstrated that masking can generate implicit positive correlations between unmasked views that share the same target masked output. However, it's highly improbable for two unmasked views to share precisely the same masked output. This discrepancy results in limited connectivity within the augmentation graph, potentially leading to suboptimal performance as shown in prior studies (Wang et al., [2022] HaoChen et al. 2021). Subsequently, we observe that when employing discrete tokens, masking creates implicit positive correlations between unmasked views as long they share the same discrete class of target masked output. Specifically, we find that discrete tokens that align well with data classes enhance connectivity among intra-class samples, thereby improving the downstream performance. Conversely, incorrectly specified tokenizers may cause confusion between inter-class samples, resulting in poorer downstream performance.
This insight can also guide the selection of proper tokenizers. Specifically, we design a novel met- ric, named token-class alignment similarity (TCAS), by measuring the discrepancy in distribution between true patch labels and the discrete token labels assigned by the tokenizer. TCAS can be used to directly compare the quality of different tokenizers without training. Meanwhile, the princi- ple inspires us design an easy- casy-to-implement tokenizer with its corresponding MIM method named ClusterMIM, which demonstrates its efficacy through enhanced performance across different bench- mark datasets (e.g., ImageNet- 100 and ImageNet-Ii and different ViT backbones.
We summarize our contributions as follows:
We are the first to identify and theoretically formalize the role of discrete tokens in MIM, highlighting how they improve the alignment between unmasked views through the aug- mentation graph analysis.
We offer a thorough theoretical exploration into the impact of discrete tokenization on the downstream generalization, emphasizing its influence on inter-class and intra-class dynam- ics within the augmentation graph.
We propose a novel metric, TCAS, for evaluating tokenizer quality without training. Mo- tivated by this metric, we propose simple yet effective approach ClusterMIM to improve the performance of MIM methods.
2 THEORETICAL UNDERSTANDING OF DISCRETE TOKENIZATION IN MIM
We begin by detailing the mathematical formulation of MIM. Given a natural image, we first reshape it into n patches, denoted as 5 E x with s representing the patch size. Subsequently, we employ a random binary mask m drawn from the set {0, 1}mh to create two complementary views of z::
3 = [(m] E X
82 = 5 m] E TR
where n and n2 are integers satisfying n = n2 + n2. We denote this random masking process as drawing 21 and xy from the joint distribution M( x22 (z) (its marginal distributions are represented as M(x, (2) and M(xz(z) Denote the set of all unmasked views as X, = {z1} and the set of all masked views as X₂ {x2), where the two sets are assumed to be finite, i.e., | N₁ and
2

---

Published as a conference paper at ICLR 2024
[2 N2- The loss function for MIM can be formulated as:
(h) = Re (.((()))))
Here, the network h maps the input 21 to reconstruct the target (22). The function transforms the unmasked image patches into visual tokens, for instance, in MAE (He et al.) 2022), & serves as an identity function. While BEiT (Bao et al., 2022) employs a d'VEE (Vahdat et al., 2018) tokenizer for t. The loss function L can be a mean square loss form for continuous targets 8 ii (He et al. Xie et al., 2022) or cross-entropy loss form for discrete l(x2) (Bao et al., 2022: Dong et al. 2023 In this paper, in order to simplify the analysis, we assume that 7. takes the form of mean
2022
square loss and our primary focus revolves around distinguishing different MIM methods based on their prediction target, specifically the selection of t.
2.1 MATHEMATICAL FORMULATION OF DISCRETE TOKENIZATION
Building upon the graph-based perspective of MIM in ZZhang et al. (2022), we delve into the effect of discrete tokenization on the downstream generalization. The foundational premise of ohang al. 2022) was to interlink X1 and X₂ through a mask graph. This framework enabled them to utilize the 2-hop connectivity within the mask graph Gm. thereby formulating an augmentation graph GA to examine the relationships between element pairs in the input space X1. Subsequently, they derived a bound for the downstream classification error as
B Bownstteam S Σ x + Cza,
(3)
where X; is the eigenvalue of normalized adjacent matrix A of augmentation graph, a is the mask- induced error E _ 1 [y(x1) * y(x+)), and C1, C2 are constants.
Here, our objective is to trace this analytical trajectory, specifically focusing on discerning the effeets of discrete tokenization on each facet of this graph-based methodology.
Tokenization Induces Equivalence Class in the Target Space. In Equation 2 the function rep- resents tokenization, which, given its intricate nature, can be analytically challenging. Therefore, it is crucial to identify the core essence of tokenization to simplify the analysis. Notably, the loss function aims align æi and its complementary tokenized view t(x2) through the reconstruction task. Considering pairs (21, sampled from M(.--la) and (If x5)) from M(..-att), an in- triguing observation emerges: even if 22 + zz. 21 and of can still share the same reconstruction target as long as ((x2) = ((xt). This observation leads us to focus on the equivalence relation denoted by within Xz. We define 22 22 if and only if t(x2) = t(xzt) Accordingly, we denote (S..... S.}. where each S, represents an equivalence class consisting of multiple 22 elements sharing the same discrete tokens. By doing so, we can treat the target view space as S because there will be no distinction between 22 and z. that belong to the same due to tokeniza- tion. This formalizes discrete tokenization using equivalence classes and obviates the need to delve into the specific properties of t. Therefore, we can equate discrete tokenization with a specific S. Building this, the joint probability of xı and S; can be conceptualized as the summation of all joint probabilities of it and each 32 contained in Sü:
~= X,/
M(zi, S.(T) = Σ Maz1,
(4)
FrES.
This definition of the joint probability is valid, since marginalization on S, makes M(r, (z). i.c. Es,es M(x1, S, = Es.es Ezye,, M(1 x) = Ex,ex, xx M(1,, (2))) M [7). is worth noting that MAE refers to extreme case where each S; becomes a single-element set.
Mask Graph with Tokenized Targets. The original mask graph proposed by /Zhang et al. (2022) is defined over the joint set of X; U describe the joint probability of pairs from X x X2. We now consider the mask graph 9M with tokenized targets, which is defined over the joint set X' = X, US:
Node: Each view Φ1 X1 and each set of views St E S.
Edge: Edges solely exist between T1 E X1 and S, E S, which essentially results in bipartite graph structure. The edge weight between F; and S; is defined as their joint prob- ability M(r1, St) = & S.).). Therefore, an edge with non-zero weight connects
3

---

Published as a conference paper at ICLR 2024
Mask Graph over X; U
Augmentation Graph
Mask Graph over X;
Augmentation Graph
X4
X
x1
Xi
x1
5
X1
x1
X1
X2
X1
X₂
X₁
X2
X1
Tokmizztion
(a) Graphs of MAE
(b) Graphs of MIM with discrete tokenization
Figure 1: An illustration of how the discrete tokenization affects the mask graph and the correspond- ing augmentation graph. I2 and z22 share the same discrete token, enabling a connection between 11 and z through z2 and 2z2 whereas such a connection not possible in MAE.
and S; if and only if there exists 22 S such that xı and z2 are complementary views generated by masking.
Figure Eesureevisse visually demonstrates the construction the mask graph with discrete tokenization, where can observe three pairs of complementary views: (æ1, 22), (zf, 222) and (2r 2²)). Notably, since both 22 and z belong to S1. this results the formation of an edge originating from either æ1, Z2, or s1 connecting to S₁. Intuitively, tokenization aggregates all the weights between 21 and 22 S; into the weight between 21 and S. This aggregation subsequently influences the construction of the augmentation graph which will be illustrated in the following part.
Induced Augmentation Graph. According to Zhang et al. (2022), MAE generates implicit con- nections among different input samples through 2-hop connectivity the mask graph, as it enforces a pair of 2-hop neighbor i and x to reconstruct the same output 22- This treatment transforms the 2-hop input neighbors into positive pairs implicitly aligned as in contrastive learning, thereby forming the concept of an augmentation graph to model relationships among all input samples in X1. However, with discrete tokenization, this mechanism is modified, now requiring 2-hop neigh- bor pair iy and x. to match the same tokenized target S,. In Figure there is illustration of how the discrete tokenization affects the mask graph Ям and the corresponding augmentation graph GA- Consequently, the augmentation graph influenced by discrete tokenization can be formulated as:
Node: Each view Φ Xi is represented as node in GA-
xI" E X1. we define their edge weight as the probability of having the same target view. Formally, this edge weight is computed as A(I1,2t) = (ssMM (S:).
Edge: For any two views xI
With the augmentation graph at hand, we can analyze how the discrete tokenization influences the generalization ability, as shown in the following part.
2.2 DOWNSTREAM GENERALIZATION UNDER DISCRETE TOKENIZATION
In this part, we design a toy model to rigorously characterize the generalization bounds under dif- ferent discrete tokenization.
Data Space. We have two classes, each containing n points representing image patches. These point sets are denoted as P1 and Pz. There that P1 n Pg = m. Therefore, there are a total of 2n m points ((P) U Pa| = 2n - m). Assuming define the data distribution such that, when drawing datum, one class, denoted as P,, and uniformly sample an ordered pair (x1, x22 from P; x P,.
overlapping points between the two classes, meaning
t m/n <
randomly select
MIM Task and Discrete Tokenization. In this simulated MIM task, set œ as the unmasked view and Ze as the masked view. Suppose that the equivalence class induced by the discrete tok- enization is S = {S1,...,Sik For the i-th equivalence class St, it comprises n1,1 elements from
4

---

Published as a conference paper at ICLR 2024
Equiralncee Clase
Puents from Class
Polnts from Clas
Class
Class 2
Class
Compardd to Tokesbarion
Langer Edge Whight
(a) MAE-like Tokenization
(b) Class-wise Tokenization
(c) Cross-class Tokenization
Figure 2: Visual illustration of the three tokenization approaches in the toy model. Each orange bounding box represents an equivalence class, whose elements share the same discrete token. Class- wise tokenization exhibits higher intra-class connectivity and lower inter-class connectivity com- pared to MAE-like tokenization. Consequently, it boasts lower downstream error bound. In con- trast, cross-class tokenization leads to lower intra-class connectivity and higher inter-class connec- tivity, resulting in significantly larger downstream error bound.
Pi/P2, n4,2 elements from Pa/P and n1,3 elements from PinP2. Therefore, we have the following conditions: Ei=1 M4,1 Σi=1 R4,2 m, Σ;=1 14,33
Tokenizers. We study on three kinds of tokenization: 1) MAE-like tokenization SMAB (which essentially implies no tokenization). class-wis tokenization Selass and 3) cross-class tokenization genss, Figure& illustrates the three tokenization approaches, By calculating the weight edge and downstream error bound each scenario, we can compare Stlass and geress with the baseline SMAE guantitatively, Specifically, for each tokenizer, calculate their downstream error bounds following Equation 3 (details in Appendix B and obtain the following results:
MAE-like tokenization SMAE, In this scenario, SMAE {S...... m where each Sş is singlc-clement set similar to MAE. In this case, the edge weight between intra-class pairs MMAAF, the edge weight between inter-class pairs MAB, and the downstream error bound BMAE should be
MAE intra
2n m
MAE
m 4n
BMA cl2
15t
+ O(P)).
(5)
4n3
inter
4
These numerical results serve as the baselines for the other two tokenization methods.
Class-wise tokenization Seluss In this scenario, Seluss = S2). The two equivalence classes divide the entire point space evenly by class, with π1, 712,2 12,,2 = - m, R1,3 122,3 = m/2. In this case, the edge weight between intra-class pairs gyclass the edge weight between inter-class pairs class. and the downstream error bound intra
m,, 71,2
Belass should be
(n
m 7
m
m(n
m 2 4n
91
w class
w inter
Belass c(2 O(P)).
(6)
2n 4
In comparison to MAE-like tokenization, class-wise tokenization enhances intra-class edge weights while diminishing inter-class edge weights. As suggested by HaoChen et al. (202I), this variation makes the feature space more distinct and separable by class, ulti- mately leading to improved downstream performance. This assertion aligns with the results of the downstream error bound calculations, where class- wise tokenization Beloss exhibits a lower downstream error bound compared to MAE- like tokenization MAF,
Cross-class tokenization geross In this scenario, n1,1 = 74,2 = (n-m)/L and n;,3 = m/l which means the / equivalence classes split the three sets of points equally. In this case, the edge weight between intra-class pairs worss the edge weight between inter-class pairs wesss and the downstream error bound geross should be
7
27t
cross RU
CrOSS W inter
ReTONS -(
+ +0()).
Intra
4n
4n
In contrast to wlass-wise tokenization, cross-class tokenization diminishes intra-class edge weights and elevates inter-class edge weights, which, in turn, has detrimental effect on downstream performance. This observation is further evidenced by the significantly larger downstream error bound Beross compared to that of MAE-like tokenization MAE
5

---

Published as a conference paper at ICLR 2024
Summary. Our findings reveal that appropriately designed discrete tokens have the potential to enhance the connectivity among intra-class samples induced by masking, leading to better down- stream performance. However, poorly chosen tokenization methods can lead to confusion among inter-class samples, resulting in a considerably larger downstream error bound that may have adverse effects overall downstream performance.
Furthermore, we seek to uncover the characteristics of well-designed visual tokens. Given the con- trast properties exhibited by gelus and gemss we delve deeper into their specific compositions. A notable distinction emerges: each equivalence class in Selas consists of points from the same true class, while each equivalence class in g comprises an equal distribution of points from both classes. This distinction leads us to wonder whether tokenization schemes aligning with the true classes result in superior downstream task performance. Affirmatively, our findings, supported by the ensuing theorem, confirm this hypothesis:
Theorem 1. Assuming that M(x1(z2) > 0 occurs only if y(æ1) = y(22). and let my denote the equivalence relation on X2 where 22 myy 22 if and only if y(22) = v( Then S = N// ~y=
{S%, ) minimizes 01 E X2 C20.
Theorem 1 suggests that, under the assumption that two views of an image come from the same class, the optimal discrete tokenizer is simply the label function. This insight serves as valuable inspiration for the design and selection of improved discrete tokenizers, discussed in the following.
3 DESIGNING AND CHOOSING BETTER DISCRETE TOKENIZER
In Section we have explored the role of discrete tokens in MIM and studied their effects on the downstream generalization. Drawing from these theoretical insights, this section is dedicated to the principles of crafting and selecting more effective discrete tokenizers. We first design a metric to evaluate the discrete tokenizers in MIM. Then, based on the metric, we design a simple but effective discrete tokenizer utilizing K-means.
3.1 TOKEN-CLASS ALIGNMENT SIMILARITY: MEASURING THE DISCREPANCY BETWEEN TOKENIZER CLASS AND TRUE CLASS
Theorem [ Suggests that, under the assumption that two views of an image come from the same class, the optimal discrete tokenizer is simply the label function. This principle underlines our approach to devising a metric for assessing discrete tokenizers. focusing on the disparity between the cquivalence classes generated by the tokenizer and those defined by the actual labels.
In practice, the set of the equivalence class could be exponentially large, since there are potential C distinct discrete representations across all masked views X₂ where C' is the size of the code- book. The vast set intractable for us to manage when designing a practical metric. Therefore, in the following part, we consider bag-of-word word model (Harris 1954) for patch representations to make the analysis tractable. This approach considers each target view (xr2) as a collection of tokens rather than a singular entity, aligning with the localized nature of patch representations and render- ing the analysis more feasible. Based on Theorem [ lower value of this metric will yield improved downstream performance. The metric constructed as detailed below:
we anticipate that a discrete tokenizer with a
1. Computing Token-class Co-occurrence. To quantify the discrepancy between the tok- enizer class and true class, we first establish the token-class co-occurrence matrix. Let R E RR4xls be a matrix, where I represents the size of the codebook of the discrete to- kenizer, and l2 represents the number of true classes. Considering all the patches in the dataset, we define R(i,) j) as the count of patches belonging to the i-th class in the code- book and the j-th true class. The matrix R, which is the normalized version of R along its rows, is given by R(i.j) = softmax, (R(i,)) Σy R(i,j").. Here, the i-th row of the matrix R, denoted R(i), indicates the distribution of true classes within the patches that belong to the i-th discrete class.
2. Measuring Tokenization-induee Class Confusion. Ideally, if we want to minimize the discrepancy between the tokenizer class and the true class, we aim for R(i) to be close to

---

Published as a conference paper at ICLR 2024
Token-Class Alignment Similarity
dVAE
0.85
Table 2: Comparison of TCAS between our proposed tokenizer and previous ones.
0.80
V-
Method
TCAS (1)
0.75
VQGAN K=16384
dVAE (Vahdat al. 2018
0.86
0.70
Perceptual Tokenizer
VQGAN T=2024 Esser et al. 2021
0.77
VQGAN K=16344 (Esse et al. 2021
0.73
42
44
46
48
50
Perceptual Tokenizer (Dong et 2023 ClusterMIM Pixel (ours)
0.69
Linear Probe Accuracy (%)
0,56
ClusterMIM DINO (ours)
0.21
Figure 3: Correlation between TCAS and linear probing accuracy.
one-hot vector, and R(i) and R(7") with + should be dissimilar. In other words, this pushes C = RR close to the identity matrix Txxh- Accordingly, we formulate the metric as follows:
= a1 Cn-Cn)+
- (1- L
E(R())())) (8)
Eiglts
To ensure that the metric is not influenced by the codebook size 11, we A 1/1 and X = 1/1.. In this metric, R(i) approaches a one-hot vector, the first term decreases. Simultaneously, R((i) and R(i') with * i' become orthogonal, the second term reaches the minimal. Thus, this metric effectively measures the discrepancy between the tokenizer class and the true class.
We designate this metric as the Token-Class Alignment Similarity (TCAS). To validate the sound- ness of TCAS, we conduct empirical analyses that explore its correlation with linear downstream task performance utilizing various tokenizers, including dVAE (Vahdat et al. 2018) VQGAN (Esser) (ct al. 2021). and Perceptual tokenizer (Dong et al., 2023). The results are visually represented in Figure 3 which indicates noticeable correlation between TCAS and performance. This signif- icant correlation not only validates the effectiveness of TCAS as a tool for evaluating tokenizer performance but also highlights its unique advantage: enabling the assessment of tokenizers without necessitating pretraining. Consequently, TCAS emerges as a valuable metric, offering insightful guidance for the refinement and development of more advanced tokenization strategies.
3.2 CLUSTERMIM: LEVERAGING CLUSTERING FOR GENERATING DISCRETE TOKENS
Building on the insights learned from theoretical analyses and toy models, we have established that an effective tokenizer should foster label correlation, even in scenarios where explicit label information is absent
common scenario in self self-supervised learning settings. This prompts the need for a viable strategy that can mimic label-correlated tokenization without relying on labeled data. Our above analyses suggest that the association of specific labels to discrete classes is not as crucial as ensuring that patches sharing the same labels converge within the same discrete category. Thus, we propose novel approach that leverages clustering to produce discrete tokens.
Tokenizer Pretraining Stage: Clustering the Patches. Given dataset, we begin with the collec- tion of image patches, akin to the inputs used by ViT (Dosovitskiy et al. (2020). Then we employ a clustering algorithm to group these patches, yielding set of clustering centers. These centers constitute the codebook for the discrete tokenizer.
Tokenizer Inference Stage: Assigning the Nearest Neighbor. To determine the discrete tokens for image patch, identify the nearest neighbor in the codebook, which is then designated as the discrete token for the respective patch.
We term the MIM approach using this tokenizer ClusterMIM (Clustering Masked Image Model- ing), which bifurcates into two practical implementations based on the clustering method: 1) Clus- Mer Pixel, which directly uses K-means in the pixel space, and 2) ClusterMIM DINO, which

---

Published as a conference paper at ICLR 2024
uses K-means in the feature space generated by DINO (Caron al. 2021 Preliminary results,
as showcased in Table reveal that our methods achieve a notable reduction in token-class align- ment similarity (TCAS) compared baselines, underscoring the potential of our clustering bustering-based tokenizer. The following experiments, detailed in Section 4 further validate the efficacy of Cluster- MIM, providing strong empirical evidence for our proposed approach.
EXPERIMENTS
In this section, we first present the main empirical results of our proposed ClusterMIM methods on different real-world datasets with different backbones. Then we conduct a series ablation experiments to discuss the selection of hyperparameters in ClusterMIM.
4.1 EVALUATION ON BENCHMARK DATASETS
To evaluate the effectiveness of the proposed ClusterMIM method, extensive experiments are con- ducted ImageNet-100 (Deng al. 2009) and ImageNet-IK (Deng al. 2009).
Setup. We mainly follow the basic setup of MAE (He et al. 2022]; for the encoder, we adopt several variants of ViT (Dosovitskiy al. 2020), including ViT-Small, ViT-Base, and ViT-Large. For the decoder, we follow the setting of He et al. (2022). The mask ratio is set to 0.75. On both datasets, we pretrain the model for 200 epochs with batch size 4096 and weight decay 0.05. We conduct both linear evaluation and non-linear fine-tuning the pretrained encoder. For linear evaluation, for non-linear fine-tuning, we train
we train linear classifier on the frozen pretrained encoder. both the pretrained encoder and the linear classifier with the soft target cross entropy loss (Peterson) iet al. 2019). For the K-Means algorithm used in the tokenizer pretraining stage, we use K-Means++ initia lizttio (Arthur & Vassilvitskii 2007). We train K-Means for 100 epochs on ImageNet-100 and 10 epochs on ImageNet-IK.
Effectiveness of the Proposed ClusterMIM Method. In Table present performance parison of different MIM methods across two benchmark datasets. Notably, our proposed methods exhibit significant performance advantage over all baseline methods. Specifically, on ImageNet- 100, with minimal extra time constructing discrete tokens through K-Means on the pixel space (5 hours on K-Means and 36 hours on pretraining), our ClusterMIM Pixel utilizing the ViT-S backbone outperform MAE by an impressive 6.8% in linear probing and 4.6% in fine-tuning accuracy. The improvement still holds when using larger backbones and larger ImageNet- 1K dataset.
Table 3: Linear probing accuracy (LP Acc.) and fine-tuning accuracy (FT Acc.) of pretrained models by various MIM methods with different VIT backbones ImageNet- 100 and ImageNet- 1K. ViT-S, ViT-B and ViT-L are abbreviations of ViT-Small, ViT-Base and ViT-Large, repectively. Bold indicates the best performance within the same setting.
Dataset
Backbone
Method
Extra Model
Acc.
FT Acc.
MAE (He et al. 2022
No
45,9
81.8
MaskFeal HOG ver. Wei et al. 2022
No
49.1
82.8
ClusterMIM Pixel (ours)
No
$2.7
$6.4
VIT-S
BEiT Bao al. 2022
dVAE
43.2
81.0
PeCo Dong 2023
VQVAE
53.2
83.6
ImageNet-1(0)
MaskFeat continuous
Wei al. 2022
DINO
57.1
84.3
ClusterMIM DINO (ours)
DINO
59.7
84.7
MAB He et al.
2022
61.2
86.9
VIT-B
BEiT
Bao et 2022
dVAE
55.8
86.1
Cluster Pixel (ours)
No
63.1
MAE He et al.
2022
No
64.4
87.3
VIT-L
ClusterMIM Pixel (ours)
No
67.2
88.9
MAR (e al. 2022
No
55.4
82.9
Cluster Pixel (ours)
No
62.1
83.2
ImageNet-IK
VIT-B
BEiT [Bao et al. 2022
dVAE
53.2
82.7
MaskFeat contmuous ver.
(Wi et al. 2022
DINO
63.2
83.7
ClusterMIM DINO (ours)
DINO
67.4
83.8
8

---

Published as a conference paper at ICLR 2024
Table Ablation study on the selection of clus- tering number K. Linear probing accuracy Fine-tuning accuracy in each box.
Experiments exploring different K- Means training epochs. Linear probing accu- racy Fine-tuming accuracy in each box.
Table 5:
K
ClusterMIM Pixel
ClusterMIM DINO
Epoch
ClusterMIM Pixel
ClusterMIM DINO
50
52.7/86.4
58.2/84.3
47.2/82.1
53.5/84.3
100
50.1/83.7
59.7/84.7
10
52.5/85.9
57.4/84.4
1000
50.6/84.4
56.9/85.1
100
52.7/86.4
59.7/84.7
Furthermore, when using the discrete tokens generated by pretrained DINO features as the recon- struction target (ClusterMIM DINO) instead the pretrained DINO features themselves (MaskFeat continuous version), we achieve a performance boost of 2.6% in linear probing and 0.4% in fine- tuning accuracy. These improvements provide direct evidence of the effectiveness of our proposed tokenizer.
4.2 ABLATION STUDY
In this section, we conduct an ablation study to investigate the impact of two key factors in our proposed ClusterMIM method: the choice of clustering number K and the training epochs of the K-Means algorithm. These experiments are conducted on the ImageNet- 100 dataset using the ViT- small architecture.
Clustering Number K. We vary the clustering number K for the K-Means algorithm while fixing the training epoch to be 100. The results of these ablation experiments are presented in Table 4 As observed in the table, different choices of clustering numbers have varying effects on the downstream performance. In the of ClusterMIM Pixel, achieve the highest linear probing accuracy of 52.7% and fine-tuning accuracy of 86.4% when K = 50. Conversely, for ClusterMIM DINO, the best performance is achieved when K 100, yielding linear probing accuracy of 59.7% and fine-tuning accuracy of 84.7%. Notably, using a very large K = 1000 did not necessarily lead to improved performance. Therefore, is advisable select a moderate clustering number in practice.
Training Epochs in K-Means. We conducted experiments to explore the influence of varying the training epochs for the K-Means algorithm within the ClusterMIM framework. The results of these experiments are presented in Table 5 Across all scenarios, the highest accuracies are consistently achieved when training K-Means for 100 epochs, demonstrating the importance Means model in enhancing ClusterMIM performance. It is noteworthy that the K-Means algorithm computationally efficient. For instance, in the case of ClusterMIM Pixel trained on ImageNet-100 with 200 epochs, while the pretraining stage takes 36 hours, training epoch of K-Means with 100 cluster centers requires only 0.05 hours. Furthermore, it's worth highlighting that even with just 1 epoch of K-Means pretraining, ClusterMIM Pixel outperformed MAE (47 1 vs. 45.9/8 and ClusterMIM DINO outperformed the continuous version of MaskFeat (53.5/84.3 vs. 53. These results underscore the effectiveness and efficiency of our proposed methods.
well-trained K-
5 CONCLUSION
In this paper, draw attention the crucial role of discrete tokens in the efficacy of MIM tech- niques. We introduce the first in-depth theoretical analysis dedicated to understanding the impac of various discrete tokenization strategies on MIM methods. Through a nuanced examination from the mask graph perspective, we have unveiled how discrete tokenization critically influences down- stream model performance. Central to our findings the revelation that tokenizer alignment with actual labels is instrumental in enhancing model efficacy. Building upon these insights, we develop Token-Class Alignment Similarity (TCAS), a novel metric to quantitatively assess the quality of given tokenizer. Leveraging this metric, we further propose straightforward yet effective discrete tokenizer along with an accompanying MIM approach, termed ClusterMIM. Our empirical evalua- tions across range of benchmark datasets demonstrate the superior performance of proposed ClusterMIM, highlighting its effectiveness in enhancing model efficacy through improved discrete tokenization. Hope our work could inspire some new tokenizer designs in MIM methods.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENT
Yisen Wang was supported by National Key R&D Program of China (2022/D0160300), Na- tional Natural Science Foundation of China (62376010, 92370129), Beijing Nova Program (20230484344), and CCT-BaiChuan-Ebech Foundation Model Fund.
REFERENCES
David Arthur and Sergei Vassilvitskii. K-means++ the advantages of careful seeding. In SODA, 2007.
Hangho Bao, Li Dong, Songhao Piao, and Furu Wei. BEit: BERT pre-training of image transform- ers. In ICLR,, 2022.
Shuhao Cao, Peng Xu, and David A Clifton. How to understand masked autoencoders. arXiv preprint (xXiv:2022.03600 2022.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jegou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In ICCV, 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. ICML, 2020.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In arXiv preprint (.t::2011./5500 2020.
Jingyi Cui, Weiran Huang, Yifei Wang, and Yisen Wang. Rethinking weak supervision in helping contrastive learning. In ICML, 2023.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009.
Xiaoyi Dong, Jianmin Bao, Ting Zhang, Dongdong Chen, Weiming Zhang, Lu Yuan, Dong Chen, Fang Wen, Nenghai Yu, and Baining Guo. Peco: Perceptual codebook for bert pre-training of vision transformers. In AAAI, 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020.
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In CVPR, 2021.
Jean-Bastien Grill, Florian Strub, Florent Altché, C. Tallec, Pierre H. Richemond, Elena Buchatskaya, Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, B. Piot, K. Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent: A new approach self-supervised learning. In NeurlPS, 2020.
Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for selfsupervised deep learning with spectral contrastive loss. NewrIPS, 2021.
Zellig S Harris. Distributional structure. Word, 10(2-3): 1954.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. CVPR, 2022.
Tianhong Li, Huiwen Chang, Shlok Mishra, Han Zhang, Dina Katabi, and Dilip Krishnan. Mage: Masked generative encoder to unify representation learning and image synthesis. In CVPR, 2023.
Jihao Liu, Xin Huang, Jinliang Zheng, Yu Liu, and Hongsheng Li. Mixmae: Mixed and masked autoencoder for efficient pretraining of hierarchical vision transformers. In CVPR, 2023.
Joshua Peterson, Ruairidh M Battleday. Thomas Griffiths, and Olga Russakovsky. Human uncertainty makes classification more robust. In ICCV, 2019.
10

---

Published as a conference paper at ICLR 2024
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In ICML, 2021.
Arash Vahdat, Evgeny Andriyash, and William Macready, Dvae#: Discrete variational autoencoders with relaxed boltzmann priors. In NeurIPS, 2018.
Yifei Wang, Zhengyang Geng, Feng Jiang, Chuming Li, Yisen Wang, Jiansheng Yang, and Zhouchen Lin. Residual relaxation for multi-view representation learning. In NewrlPS, 2021.
Yifei Wang, Qi Zhang, Yisen Wang, Jiansheng Yang, and Zhouchen Lin. Chaos is a ladder: A new theoretical understanding of contrastive learning via augmentation overlap. In ICLR, 2022.
Yifei Wang. Qi Zhang, Tianqi Du, Jiansheng Yang. Zhouchen Lin, and Yisen Wang. A message passing perspective on learning dynamics of contrastive learning. In ICLR, 2023.
Chen Wei, Haoqi Fan, Saining Xie, Chao- Wu, Alan Yuille, and Christoph Feichtenhofer.
Masked feature prediction for self-supervised visual pre-training. In CVPR, 2022.
Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: A simple framework for masked image modeling. In CVPR, 2022.
Qi Zhang, Yifei Wang, and Yisen Wang. How mask matters: Towards theoretical understandings of masked autoencoders. In NeurIPS, 2022.
Qi Zhang, Yifei Wang, and Yisen Wang. On the generalization of multi-modal contrastive learning. In ICMI., 2023a.
Qi Zhang, Yifei Wang, and Yisen Wang. Identifiable contrastive learning with automatic feature importance discovery. In NeurIPS, 2023b.
Jinghao Zhou, Tao Kong, Xianming Lin, Rongrong Ji, al. Exploring target representations for masked autoencoders. In ICLR, 2023.
Zhijian Zhuo Yifei Wang, Jinwen Ma. and Yisen Wang. Towards a unified theoretical understanding of non-contrastive learning via rank differential mechanism. ICLR, 2023.

---

Published as a conference paper at ICLR 2024
A ADDITIONAL RELATED WORKS
Masked Image Modeling (MIM). MIM involves learning from portion of masked input signals by predicting these signals using the unmasked part. Different masked prediction objectives have been proposed for MIM. For instance, MAE (He al. 2022), and variants like U-MAE (Zhang et al. 2022) and MixMAE (Liu et al. 2023), directly predict the raw pixels within the masked part. DBot Zhou et al., 2023) can also use representation from another pretrained model as prediction target, but updating the target per certain epochs with current model representation. Notably, the current designs of discrete tokens in these methods do not possess theoretical guarantees. In this paper, we provide theoretical understanding of the working mechanism of discrete tokens in MIM, especially regarding their impact on downstream generalization.
Theoretical Understanding of MIM. Despite the remarkable success of MIM methods, their the- oretical understandings remain rarely explored. Cao et al. (2022) focus on the attention mechanism of MAE through an integral kernel perspective. Zhhang et al. (2022) build connection between MAE and contrastive learning, highlighting the importance of masking. While these theoretical approaches offer meaningful insights into MAE, they sidestep examination of the influence of discrete tokens on MIM. Therefore, our work intends to explore the specific role and impact of discrete tokens within the MIM framework.
B DETAILED CALCULATION IN SECTION 2.2
We denote W21.02 M( =2) the joint probability of one complementary pairs and denote Wx, = -M(x1) as the marginal probability of one view.
Consider the joint probability of complementary pair T2, When the two views both come from the same overlapped region, the joint probability Wan.22 1/n². Otherwise, the joint proba- bility is 1/2n².
For the masked view 222 unmasked view 21 in the non-overlapped region, the marginal probability 1/(2.). For masked views and unmasked views the overlapped region, the marginal probability 1/n.
The normalized augmentation graph is
Wr x2 W 2x W.
Was 22 2 an
Adiscrene (I1, xt)
Way W 2 Ww
W of
Way W (D,)
D3 22 ED,
We calculate the joint probability of a positive pair under each situation. Each item denotes where the two unmasked views of the pair come from:
(Pr/P2,PP/P3
2n
(ne.1 + n4.3) 1
P11
2n
1 + n4. + 2n133
2n
2n
[=1
(10)
( (on1++1,,33+ 13,,))
74,1 + n4,2 + 2n4,3
n22
(Pa/Pa.
2n
(mi.2 + n4,3)² 1 1
D22 2n
n1,1 + n:,2 + 2n,,3
2n² 2n²
(11)
(ni,2 + n2,3)²
n4,1 + n4.2 + 2n4,3
n2
12

---

Published as a conference paper at ICLR 2024
(P.PP,.PP/P)
2n
(meA. + +naa) + nes))
1
P12 = 2n
11,11 + n4,2 + 2n,,3
2n²
2n²
1m1
(12)
2 1 ni.3
74,1 + n4,2 + 2n2,3
n²
(P/P2, Pin.3
P13 2n
2n
( + m1,3)44
+ (mi,1
n1,3)
1
7i,1 + n4, + 2n,,3
2n
2n
n²
2n
( (10.1 + (,,)(mss + 2n1aa
v2
11,1 + n1.2 + 2n4,3
2n
(13)
(PP/P,,PPPP))
P23 = V2n
2n
(n + +4,3)4422
+ (n,,2 +
n1,s)
14,1 n442 + 2n1,3
2n
2n
n²
2n
(ma + ne,))(m,22 + 2n1,3)
v2
n4,1 + 14,2
2n,,3
2n?
(14)
(1
(10.1 + n1, + 222,a)2
1
P33 = n
11,1 + n4, + 2n,,3
2n
2n12
i=1
(15)
(ma.1 + 1122 + 2n4,3)
n1,1 + n1,2 + 2m,,3
2n
The normalized adjacent matrix is:
P11
Pi1
Pi1
P12
P12
P12
P13
Pi1
PI1
PI1
P12
P12
P12
P13
PI1
PI1
P12
P12
P12
P13
A =
P12
P12
P12
P22
P22
P22
P23
(16)
P12
P12
P12
D22
P22
P22
P23
P12
P12
P12
P22
P22
P22
P23
P13
P13
P13
P23
P23
P23
P33
The sum of square of the eigenvalues is equal to the sum of square of the elements in the adjacent matrix:
(n + p22 + 2p(2) + (n m)m(pi + +22)) + m²²p3
(n mm3
(Mi,1 + n2.3)²
(n2,2 14.3))
)²2 2( (n + n1,3) + ne,s) )2)
)2
nd
T4,1 + 14.2 + 2n1433
74,1 + n4.2 + 221433
Mi,1 + 14,22 + 2n153
(n m m
(n + n4,3 ( 2m
(n4.2 + 14.3 (n4,2 + 2m,,s) )2
+ m?
+
2nd
n4,1 + 14.2 + 2m
T4,1 + n4.2 + 2n
n
(17)
13

---

Published as a conference paper at ICLR 2024
The labeling error should be
(n - m)²
(n + n1,3) + -n1,3)
B m ((()))()) 2(n (10-m)pp12(22)
(18)
n
ni,1 + 14i + 2ni,3
We consider the downstream error bound derived in (hang et al.
2022 which is
B Σ X2 + C2a
(n m²
(n nn,,3)
(ni,2 + n1,3)
)2 2 (ni,1 + m,,,) n1,3) )2)
)2
0
n²
n1,1 + n2,2 + 2n1,3
n2,1 + -2,22 + +2n;,3
ni.1 + n;, + 2n
(n mmm
(n + (11,3) + 2n1,3)
(n + n1,3) + 221,) )²
+ m²
+
+ C3a 112
2n
ni.1 + 14,2 + 2n14,3
ni,1 + ni,2 + 2n1,3
n1 m)?
(Nt_1 + n2,33)
)2
(n1.2 + n2.,3) )
-2 + 2(
(n,,1 + (1 + n1,3)
+ Can )2)
nd
14i,1 + 14.2 + 2n1,3
12,11 n4,2
2n,,3
Mi,1 + n4, + 221,3
i=1
(n m)m
(ni.1 n2,33 (11,1 + 2n4,3)
(ni.2 + n1. (ni,) + 221,,) )2
m
)2 +
2n
74,1 + n22 + 2n,,3
71,1 + n1, + 2n,,3
n22
(19)
where C' =n) /882²) and c3 = 5/2. We denote = m/n and assume < 1. Since P11: P12 and P22 are elements from normalized adjacency matrix, the edge weight of intra-class pairs can be computed as (P11 + P22)/ /4n and the edge weight inter-class pairs can be computed as p12/4n.
MAE Tokenization. We have 2n and only of n4,1, f1,2 and n1,3 is and the other two are 0 for each
< i < l. The bound for MAE tokenization should be
(n m) d
(n m)m
m C)
MAE
(2(n-
+
12
(n² + n²²
n1
2
2n
n12
c22 51 + 7 2 41a 1
+ 2
(1 1)?()
(20)
15
11
c22 - -
2
3
+
4
The two types of edge weights are
2n
MAE
MAE
(21)
4n
4m
Class-wise Tokenization. We consider an extreme case where = 2, 11.1 = n 11.2
0, n1,3 = m/2 and 12211 = 0, 2222 = n - m, 12233 = m/2 The bound for this case should
39
29
95
15
(n m) m 7
Beluss = =c(2- -7 7 16/2
3 +
5
+
6
+
+c
2
16
16
n
39
29
95
15
5
t2
= =c(2 71 + 161
3
+
5
6
-
(1
- 1)2
(22)
2
16
16
2
2
39
c(2 - + 2
+ 0(t)))
2
The two types of edge weights are
(n
m2 )
mmy²
m(n
class
w
(23)
2n
4n
14

---

Published as a conference paper at ICLR 2024
Cross-class Tokenization. Here, 74,1
n4,2 (n mmml and n4.3 mm/. The bound should be
Belass
(n
n²² m
n²
2 T6²
(n-m)m
(1 n2
+
- +
)
m²²
n2
(n
m))²
n²
2n
2n2
= =((1 - t)² + (1 - t)t(t + 1)2 + 8
7
c
(1- - -))²
4
2
7 27 +
15 1 + 0(²)).
(
4
(24)
The two types of edge weights are
1
(25)
intra
4n2
4n2"
C TOY MODEL WITH MULTIPLE CLASSES
In this section, We extend the discussion in Section 2.2 from two classes to multiple classes. Here are the setting of the toy models.
Data Space. We have classes, each containing n points representing image patches. These point sets are denoted as P.; i = 1,2, 1.2..... 8. There are m overlapping points between any two classes, meaning that [P, n PJl = m. There does not exist any point belonging to three classes. Therefore, there are a total of sn ms(s 1)/2 points. Assuming = m/n < 1, we define the data distribution such that, when drawing datum, we randomly select one class, denoted as P., and uniformly sample an ordered pair (æ1, from P; x P..
MIM Task and Discrete Tokenization. In this simulated MIM task, set 21 as the unmasked view and Z2 as the masked view. Suppose that the equivalence class induced by the discrete tokenization is S (S Si). For the i-th equivalence class S, it comprises ag elements from Pa n Ps where a + b, ni,a elements from Pa! Ubita Pb- Therefore, we have the following conditions:
Σ niia = - (s - 1) * m. Σ Ri,b = m. a a#1 #b.
(26)
1=1
Tokenizers. We also study on three kinds of tokenization: MAE-like tokenization MAB (which essentially implies no tokenization), class-wise tokenization Selas and cross-class tokenization geress, By calculating the weight edge and downstream error bound in each scenario, compare Selass and Semes with the baseline SMAF, By doing this, are able to explore how tokenization influences the augmentation graph and the downstream error bound as follows:
MAE-like SMAE, In this scenario, SMAE = (S,, Syn 11n-ma(s-1)/2/ where each S single-element set similar to MAE. this case, the edge weight between intra- class pairs w MAE. the edge weight between inter-class pairs w Wnter
MAf. the downstream
"inga error bound MAE should be
MAF Pintra
2n - (8 1)m
MAF inter
m
W
2sn
2sn
(27)
15(s
1)t
BMAE = c(s
+ +(t))
These numerical results serve as the baselines for the other two tokenization methods. - Class-wise gelass In this scenario, Selass = S2..,,,, The s equivalence classes divide the entire point space evenly by class, with ns,s,bb m/2, T1s,s (s m.m. In this case, the edge weight between intra-class pairs wlllass the edge weight between inter-class pairs filles the downstream error bound Belass should be
(n (s 1) m )2 + (s ))
m(n (s (ss-1) 1)
w class
w class
2
2sn
(28)
9(s- 1)t
Belass
c(s
+ O(t2)).
2
15

---

Published as a conference paper at ICLR 2024
In comparison to MAE-like tokenization, class-wise tokenization enhances intra-class edge weights while diminishing inter-class edge weights. This leads to lower down- stream enror bound, which is consistent with the situation where 2.
Cross-class geross, In this scenario, fie (n- (s- 1---1)mm 1m)/ and ns.a,b = m a + b which means the I equivalence classes split each set of points equally. In this case, the edge weight between intra-class pairs wimtra. the edge weight between inter-class pairs weees the downstream error bound Berss should be
W
CrOsS
intra
2sn
CTOIN
inter
2sn
c 2 5ウ + O(t).
(29)
In contrast to class-wis wise tokenization, cross-class tokenization diminishes intra-class edge weights and elevates inter-class edge weights, which, in turn, has a detrimen- tal effect on downstream performance. This observation is further evidenced by the significantly larger downstream error bound Benss compared that of MAE-like kenization BMAE."
In conclusion, the conclusion for multiple classes similar to that for two classes. While class- wise tokenization will obtain larger intra-class connectivity and lower inter-class connectivity, leading lower downstream error bound and potentially better downstream performance. On the other hand, cross-class tokenization that does not align well with the class, obtain larger inter-class connectivity and lower intra-class connectivity, leading to larger downstream error bound.
D EXPERIMENTS ON LONGER TRAINING
We conduct additional experiments using ClusterMIM Pixel involving 800 epochs of pretraining on ImageNet- 1k using ViT-B as the backbone. The comparisons with baseline methods are presented in Table 6 Notably, with only half the number of pretraining epochs, ClusterMIM Pixel reaches the same finetuning accuracy with MAE and ourperforms MAE the linear probing accuracy. The results indicate that ClusterMIM Pixel outperforms the baseline results in the classification tasks, ensuring the better learning ability of ClusterMIM.
Table 6: Linear probing accuracy and fine-tuning accuracy of pretrained models with ViT-B on ImageNet- IK. Bold indicates the best performance.
Methods
Epochs
Linear Probing Accuracy
Fine-tuning Accuracy
МAE
1600
68.0
83.6
BETT
800
N/A
83.2
ClusterMIM Pixel
800
69.4
83.6
E TCAS SCORE OF VARIOUS TOKENIZERS
In this section, will present and analyze the TCAS scores of various tokenizers and the responding performance on ImageNet-100. The results are detailed in Table 7 and we visually organize the outcomes Figure where we use linear regression fit the data. The r-value of the fitting is 0.85, demonstrating a strong correlation between TCAS score and linear probing accuracy. Specifically, the tokenizer's linear probing performance tends to improve as the TCAS metric de- creases. Therefore, TCAS can serve as valuable guidance for assessing the quality of a tokenizer.
16

---

Published as a conference paper at ICLR 2024
Table 7: TCAS scores and performances on classification tasks of various discrete tokenizers on
ImageNet-100 100 with 200 epochs pretraining.
K
Method
TCAS
Linear Probing Acc.
Fine-tuning Acc.
10
ClusterMIM Pixel
0.69
42.7
81.9
25
ClusterMIM Pixel
0.60
48.3
84.6
50
ClusterMIM Pixel (KMEANS epoch)
0.66
47.2
82,1
50
ClusterMIM Pixel (KMEANS 10 epochs)
0.56
52.5
85.9
50
ClusterMIM Pixel (KMEANS 100 epochs)
0.56
52.7
86.4
50
ClusterMIM DINO
0.21
58.2
84.3
100
ClusterMIM Pixel
0.52
50.1
83.7
100
ClusterMIM MAE
0.34
52.3
83.2
100
ClusterMIM DINO
0.15
59.7
84.7
100
ClusterMIM DeiT
0.09
64.8
83.1
200
ClusterMIM Pixel
0.49
50.8
84.8
1000
ClusterMIM Pixel
0.39
50.6
84.4
1000
ClusterMIM DINO
0.12
56.9
85.1
1024
VQGAN
0.77
44.3
80.9
8192
dVAE
0.86
41.2
80.2
8192
Perceptual Tokenizer
0.69
50.3
83.1
Token-Class Alignment Similarity
0.8
0.6
0.4
0.2
40
45
50
55
60
65
Linear Probe Accuracy (%)
Figure 4: The relation between TCAS score and linear probing accuracy on ImageNet-100. 100.
17


Reviews:
[{"review_id": "xztkjxSEqQ", "reviewer": "Reviewer_yDgo", "summary": "This paper explores the role of tokenization in Masked image modeling (MIM). It starts by discussing how discrete tokenization affects generalization performance on downstream tasks. For measuring the proficiency of different tokenization roles, this work designs the token-class alignment similarity (TCAS) metric. Based on the TCAS, they propose a MIM framework (K-MIM) with a cluster-based discrete tokenization scheme. Extensive experiments are conducted to validate the effectiveness of the proposed tokenization strategy.", "questions": "- Could you please show more ablation results about clustering numbers (e.g., K = 10, 25, 200) in Table 4?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The main problem, \"What is the role of tokenization in MIM? How does it affect downstream performance?\", is interesting and necessary.\n- The proposed token-class alignment similarity (TCAS) is a cheap but effective metric to measure the performance of tokenization roles.\n- This paper theoretically and empirically demonstrates the superiority of class-wise tokenization.\n- The paper is well-organized and easy to follow.", "weaknesses": "- It would be better to show more ablation results about clustering numbers in Table 4.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Wl0aa4RyNt", "reviewer": "Reviewer_ss62", "summary": "The paper analyzes the effect of discrete tokens as targets of masked image model (MIM) training. Using a graph-based view for MIM [A], the paper shows that discrete tokens change the graph of MIM. Also, it provides a theorem that discrete tokens similar to image class are the best for MIM targets. Based on the theorem, a metric to measure better tokenizer for MIM is proposed, named Token-Class Alignment Similarity (TCAS). Last, it is shown that simple K-means, which have a better TCAS score than other baselines, could improve the performance of MIM.\n\n[A] How Mask Matters: Towards Theoretical Understandings of Masked Autoencoders, NeurIPS 2022", "questions": "- 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.\n\n[B] Exploring Target Representations for Masked Autoencoders, arxiv\n\n- 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?\n\n- 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?\n\n- Discrete tokenizers like dVAE and VQGAN employ ConvNet or ViT, utilizing the entire image to create tokens. These tokens are interrelated, and a token from one location can incorporate patches from others. However, it looks like the paper handles these tokens as individual local information, which is not correct. Is there any explanation for this?\n\n- Although K-MIM has better TCAS than other tokenizers, it is hard to say that K-MIM is specially designed for TCAS. Is there \bany discrete tokenization method optimized for the TCAS metric?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Theoretical analysis on a discrete tokenization method looks novel and interesting.\n- A metric for tokenization (TCAS) would help a lot of researchers to investigate MIM.", "weaknesses": "- 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.\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2YzNAPF6Nk", "reviewer": "Reviewer_f8wZ", "summary": "This authors of this paper present theoretical analysis on the problem of masked image modeling (MIM), and study the impact of discrete tokenization in the MIM pipeline.", "questions": "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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors present an intriguing tokenization approach using graph representation, and the paper is both technically robust and clearly articulated", "weaknesses": "The theoretical analysis is only considered for two classes. I wonder if this can be extended into multiple classes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "f9sCcjQ1NS", "reviewer": "Reviewer_d6S4", "summary": "Masked image modeling methods (MIM) have started to employ discrete tokenization as a reconstruction target instead of raw pixel values. Yet, the role of tokenization and how it affects downstream performance is not well studied. This paper studies the impact of tokenizer design on downstream performance. The paper finds that that tokenizer that are more aligned with the downstream task labels (ie, more discriminative with respect to downstream classes) result in better performance, and design a metric to measure this similarity. The paper also introduces a new tokenizer based on clustering features to provide discretized reconstruction targets. Methods trained with the new tokenization achieve a better performance that existing approaches. The paper also includes some interesting analysis for the new tokenizer.\n\n**Update (11/20):** I have raised my rating from 6 to 8 in respond to reading the other reviews and the author's response to all reviews. I think the paper presents an interesting study of the impact of discretization and proposes an interesting approach to representation learning.", "questions": "- Could you explain why the bounds in equation 7 do not depend on $l$? Is it because each point is equivalent to an equal number of each class? \n\n- 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? \n\n- 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. \n\n- 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. \n\n- 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.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The framing of the equivalency structure induced by the tokenization was quite interesting. \n\n- The paper poses an interesting question regarding the impact of tokenization and provides several interesting observations and analysis to answer it. \n\n- The proposed metric is intersting and seems to have predictive power. Although, as noted in the weaknesses, it would be great to report it for all experiments. \n\n- The proposed tokenizer is fairly simple, especially for K-MIM pixel which achieves a significant performance gain without requiring any other models.", "weaknesses": "- 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.\" \n\n- The use of a pre-trained SSL models to create the tokenization seems a bit odd. I would argue that if you use the features of model X to generate the labels/tokens to train model Y, then model Y is effectively being supervised by model X. While this is okay since both models have similar training requirements, one would expect model Y to outperform model X for this strategy to be pragmatic. Yet, K-MIM DINO achieves a much lower linear probe accuracy than DINO. Furthermore, the efficiency argument made in Sec 5.2 would need to take into account the time taken to train DINO for the K-MIM DINO results (the argument for K-MIM PIXEL holds and is a very nice finding). \n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)