SlowGuess/NeurIPS_2023_OCR huggingface dataset:
CLUSTERFORMER: Clustering As A Universal Visual Learner
James C. Liang
Rochester Institute of Technology
Yiming Cui
Qifan Wang
Tong Geng
Wenguan Wang
Dongfang Liu*
This paper presents CLUSTERFORMER, a universal vision model that is based on the CLUSTERing paradigm with TransFORMER. It comprises two novel designs: ① recurrent cross-attention clustering, which reformulates the cross-attention mechanism in Transformer and enables recursive updates of cluster centers to facilitate strong representation learning; and ② feature dispatching, which uses the updated cluster centers to redistribute image features through similarity-based metrics, resulting in a transparent pipeline. This elegant design streamlines an explainable and transferable workflow, capable of tackling heterogeneous vision tasks (i.e., image classification, object detection, and image segmentation) with varying levels of clustering granularity (i.e., image-, box-, and pixel-level). Empirical results demonstrate that CLUSTERFORMER outperforms various well-known specialized architectures, achieving  $83.41\%$  top-1 acc. over ImageNet-1K for image classification,  $54.2\%$  and  $47.0\%$  mAP over MS COCO for object detection and instance segmentation,  $52.4\%$  mIoU over ADE20K for semantic segmentation, and  $55.8\%$  PQ over COCO Panoptic for panoptic segmentation. For its efficacy, we hope our work can catalyze a paradigm shift in universal models in computer vision.
Computer vision has seen the emergence of specialized solutions for different vision tasks (e.g., ResNet [34] for image classification, Faster RCNN [70] for object detection, and Mask RCNN [33] for instance segmenta
tion), aiming for superior performance. Nonetheless, neuroscience research [73, 65, 82, 5] has shown that the human perceptual system exhibits exceptional interpretive capabilities for complex visual stimuli, without task-specific constraints. This trait of human perceptual cognition diverges from current computer vision techniques [95, 44, 46], which often employ diverse architectural designs.
Human vision possesses a unique attention mechanism that selectively focuses on relevant parts of the visual field while disregarding irrelevant information [81, 40]. This can be likened to a clustering approach [2, 3, 89], in which individual pixel points are decomposed and reorganized into relevant concepts to address various tasks. This essentially is a hierarchical process that involves combining
Figure 1: CLUSTERFORMER is a clustering-based universal model, offering superior performance over various specialized architectures.
basic visual features, such as lines, shapes, and colors, to create higher-level abstractions of objects, scenes, and individuals [79, 59, 66, 27]. Inspired by the remarkable abilities of the human vision system, this work aims to develop a universal vision model that can replicate the unparalleled prowess.
To this end, we employ a clustering-based strategy that operates at varying levels of granularity for visual comprehension. By solving different vision tasks (i.e., image classification, object detection, and image segmentation), we take into account the specificity at which visual information is grouped (i.e., image-, box-, and pixel-level). We name our approach, CLUSTERFORMER (§3.2), as it utilizes a CLUSTERING mechanism integrated within the TransFORMER architecture to create a universal network. The method begins by embedding images into discrete tokens, representing essential features that are grouped into distinct clusters. The cluster centers are then recursively updated through a recurrent clustering cross-attention mechanism that considers associated feature representations along the center dimension. Once center assignments and updates are complete, features are dispatched based on updated cluster centers, and then both are fed into the task head for the target tasks.
CLUSTERFORMER enjoys a few attractive qualities. ① Flexibility: CLUSTERFORMER is a clustering-anchored approach that accommodates a broad array of visual tasks with superior performance (see Fig. 1) under one umbrella. The core epistemology is to handle various tasks with different levels of granularity (e.g., image-level classification, box-level detection, pixel-level segmentation, etc.), moving towards a universal visual solution. ② Transferability: The cluster centers generated by the CLUSTERFORMER encoder are directly employed by the task head as initial queries for clustering, allowing the entire architecture to transfer underlying representation for target-task predictions (see Table. 4). This elegant design facilitates the transferability of knowledge acquired from the upstream task (i.e., encoder trained on ImageNet [72]) to downstream tasks (e.g., decoder trained on instance segmentation on COCO [49]). ③ Explainability: Regardless of the target tasks, CLUSTERFORMER's decision-making process is characterized by a transparent pipeline that continuously updates cluster centers through similarity-based metrics. Since the reasoning process is naturally derivable, the model inference behavior is ad-hoc explainable (see §4.2). This differs CLUSTERFORMER from most existing unified models [17, 44, 95] that fail to elucidate precisely how a model works.
To effectively assess our method, we experimentally show: In §4.1.1, with the task of image classification, CLUSTERFORMER outperforms traditional counterparts, e.g.,  $0.13 \sim 0.39\%$  top-1 accuracy compared with Swin Transformoer [53] on ImageNet [72], by training from scratch. In §4.1.2, when using our ImageNet-pretrained, our method can be expanded to the task of object detection and greatly improve the performance compared to Dino [96] over Swin Transformer on COCO [49] ( $0.8 \sim 1.1\%$  mAP). In addition, our method can also adapt to more generic per-pixel tasks, a.k.a, semantic segmentation (see §4.1.3), instance segmentation (see §4.1.4), and panoptic segmentation (see §4.1.5). For instance, we achieve performance gains of  $0.6 \sim 1.3\%$  mIoU for semantic segmentation on ADE20K [101],  $1.0 \sim 1.4\%$  mAP for instance segmentation on MS COCO [49] and  $1.5 \sim 1.7\%$  PQ for panoptic segmentation on COCO Panoptic [42] compared with Mask2Former [17] over Swin Transformer. Our algorithm are extensively tested, and the efficacy for the core components is also demonstrated through a series of ablative studies outlined in §4.2,
Universal Vision Model. Transformers [81] have been instrumental in driving universal ambition, fostering models that are capable of tackling tasks of different specificity with the same architecture and embody the potential of these recent developments [23, 17, 16, 95, 96, 4, 80, 30, 57, 86] in the field. In the vision regime, mainstream research endeavors have been concentrating on the development of either encoders [53, 88] or decoders [44, 94]. The encoder is centered around the effort of developing foundation models [4, 53, 24, 22], trained on extensive data that can be adapted and fine-tuned to diverse downstream tasks. For instance, Swin Transformer [53] capably serves as a general-purpose backbone for computer vision by employing a hierarchical structure consisting of shifted windows; ViT-22B [22], parameterizes the architecture to 22 billion and achieves superior performance on a variety of vision tasks through learning large-scale data. Conversely, research on decoders [23, 17, 16, 95, 94, 44, 96, 87, 50, 20, 52, 19, 21, 51, 93, 76, 37, 99, 25, 48] is designed to tackle homogeneous target tasks, by using queries to depict visual patterns. For instance, Mask2Former [17] incorporates mask information into the Transformer architecture and unifies various segmentation tasks (e.g., semantic, instance, and panoptic segmentation); Mask-DINO [44] extends the decoding process from detection to segmentation by directly utilizing query embeddings
for target task predictions. Conceptually different, we streamline an elegant systemic workflow based on clustering and handle heterogeneous visual tasks (e.g., image classification, object detection, and image segmentation) at different clustering granularities.
Clustering in Vision. Traditional clustering algorithms in vision [39, 28, 29, 55, 91, 1, 10, 61, 6, 58] can be categorized into the hierarchical and partitional modes. The hierarchical methods [62, 38] involve the modeling of pixel hierarchy and the iterative partitioning and merging of pixel pairs into clusters until reaching a state of saturation. This approach obviates the necessity of a priori determination of cluster quantity and circumvents the predicaments arising from local optima. [98, 12]. However, it exclusively considers the adjacent pixels at each stage and lacks the capacity to assimilate prior information regarding the global configuration or dimensions of the clusters. [69, 64]. In contrast, partitional clustering algorithms [78, 36] directly generate a flat structure with a predetermined number of clusters and exclusively assign pixels to a single cluster. This design exhibits a dynamic nature, allowing pixels to transition between clusters [11, 63]. By employing suitable measures, this approach can effectively integrate complex knowledge within cluster centers. As a powerful system, human vision incorporates the advantages of both clustering modes [89, 83, 67]. We possess the capability of grouping analogous entities at different scales. Meanwhile, we can also effectively categorize objects purely based on their shape, color, or texture, without having the hierarchical information. Drawing on the above insights, we reformulate the attention mechanism (§3.2) in Transformer architectures [81] from the clustering's perspective to decipher the hierarchy of visual complexity.
Clustering. The objective of clustering is to partition a set of data points, denoted by  $X \in \mathbb{R}^{n \times d}$ , into  $C$  distinct clusters based on their intrinsic similarities while ensuring that each data point belongs to only one cluster. Achieving this requires optimizing the stratification of the data points, taking into account both their feature and positional information, to form coherent and meaningful groupings. Clustering methodologies typically employ advanced similarity metrics, such as cosine similarity, to measure the proximity between data points and cluster centroids. Additionally, they consider the spatial locality of the points to make more precise group assignments.
Cross-Attention for Generic Clustering. Drawing inspiration from the Transformer decoder architecture [81], contemporary end-to-end architecture [17, 9] utilize a query-based approach in which a set of  $K$  queries,  $C = [c_{1};\dots;c_{K}]\in \mathbb{R}^{K\times D}$ , are learned and updated by a series of crossattention blocks. In this context, we rethink the term "  $C$  " to associate queries with cluster centers at each layer. Specifically, cross-attention is employed at each layer to adaptively aggregate image features and subsequently update the queries:
$$
\boldsymbol {C} \leftarrow \boldsymbol {C} + \operatorname {s o f t m a x} _ {H W} \left(\boldsymbol {Q} ^ {C} \left(\boldsymbol {K} ^ {I}\right) ^ {\top}\right) \boldsymbol {V} ^ {I}, \tag {1}
$$
where  $Q^{C} \in \mathbb{R}^{K \times D}$ ,  $V^{I} \in \mathbb{R}^{HW \times D}$ ,  $K^{I} \in \mathbb{R}^{HW \times D}$  represent linearly projected features for query, key, and value, respectively. The superscripts “ $C$ ” and “ $I$ ” denote the features projected from the center and image features, respectively. Motivated by [95], we follow a reinterpretation of the cross-attention mechanism as a clustering solver by considering queries as cluster centers and applying the softmax function along the query dimension ( $K$ ) instead of the image resolution ( $HW$ ):
$$
\boldsymbol {C} \leftarrow \boldsymbol {C} + \operatorname {s o f t m a x} _ {K} \left(\boldsymbol {Q} ^ {C} \left(\boldsymbol {K} ^ {I}\right) ^ {\top}\right) \boldsymbol {V} ^ {I}. \tag {2}
$$
3.2 CLUSTERFORMER
In this subsection, we present CLUSTERFORMER (see Fig. 2(a)). The model has a serial of hierarchical stages that enables multi-scale representation learning for universal adaptation. At each stage, image patches are tokenized into feature embedding [81, 53, 24], which are grouped into distinct clusters via a unified pipeline — first recurrent cross-attention clustering and then feature dispatching.
Recurrent Cross-Attention Clustering. Considering the feature embeddings  $\mathbf{I} \in \mathbb{R}^{HW \times D}$  and initial centers  $C^{(0)}$ , we encapsulate the iterative Expectation-Maximization (EM) clustering process,
(c) Feature Dispatching (Eq. 6)
(b)
Figure 2: (a) Overall pipeline of CLUSTERFORMER. (b) Each Recurrent Cross-Attention Clustering layer carries out  $T$  iterations of cross-attention clustering ( $E$ -step) and center updating ( $M$ -step) (see Eq. 3). (c) The feature dispatching redistributes the feature embeddings on the top of updated cluster centers (see Eq. 6). consisting of  $T$  iterations, within a Recurrent EM Cross-Attention layer (see Fig. 2(b)):
$$
\begin{array}{l} E \text {- s t e p :} \quad \hat {\boldsymbol {M}} ^ {(t)} = \operatorname {s o f t m a x} _ {K} \left(\boldsymbol {Q} ^ {C ^ {(t)}} \left(\boldsymbol {K} ^ {I}\right) ^ {\top}\right), \\ M \text {- s t e p :} \boldsymbol {C} ^ {(t + 1)} = \hat {\boldsymbol {M}} ^ {(t)} \boldsymbol {V} ^ {I} \in \mathbb {R} ^ {K \times D}, \\ \end{array}
$$
where  $t \in \{1, \dots, T\}$  and  $\hat{\boldsymbol{M}} \in [0, 1]^{K \times HW}$  represents the "soft" cluster assignment matrix (i.e., probability maps of  $K$  clusters). As defined in Section 3.1,  $\boldsymbol{Q}^C \in \mathbb{R}^{K \times D}$  denotes the query vector projected from the center  $\boldsymbol{C}$ , and  $\boldsymbol{V}^I, \boldsymbol{K}^I \in \mathbb{R}^{HW \times D}$  correspond to the value and key vectors, respectively, projected from the image features  $\boldsymbol{I}$ . The Recurrent Cross-Attention approach iteratively updates cluster membership  $\hat{\boldsymbol{M}}$  (i.e.,  $E$ -step) and centers  $\boldsymbol{C}$  (i.e.,  $M$ -step). This dynamic updating strategy embodies the essence of partitional clustering. It enjoys a few appealing characteristics:
- Efficiency: While the vanilla self-attention mechanism has a time complexity of  $\mathcal{O}(H^2 W^2 D)$ , the Recurrent Cross-Attention approach exhibits a lower bound of  $\mathcal{O}(TKHWD)$ . This is primarily due to the fact that  $TK \ll HW$  (i.e., 4165 in Swin [53] vs. 1200 in ours). Specifically, considering the nature of the pyramid architecture [88, 53] during the encoding process,  $TK$  can indeed be much smaller than  $HW$ , especially in the earlier stages. It is important to note that during each iteration, merely the  $Q$  matrix requires an update, while the  $K$  and  $V$  matrices necessitate a single computation. Consequently, the whole model enjoys systemic efficiency (see Table 6c).
- Transparency: The transparency hinges on the unique role that cluster centers play in our Recurrent Cross-Attention mechanism. The cluster centers, derived through our clustering process, act as 'prototypes' for the features they cluster. These 'prototypes' serve as a form of a representative sample for each cluster, reflecting the most salient or characteristic features of the data points within that cluster. Moreover, the Recurrent Cross-Attention method adheres to the widely-established EM clustering algorithm, offering a lucid and transparent framework. This cluster center assignment behaves in a human-understandable manner (see Fig. 3) during representation learning and fosters ad-hoc explainability, allowing for a more intuitive understanding of the underlying relationships.
- Non-parametric fashion: The Recurrent Cross-Attention mechanism achieves a recursive nature by sharing the projection weights for query, key, and value across iterations. This approach effectively ensures recursiveness without the introduction of additional learnable parameters (see Table 6b).
Since the overall architecture is hierarchical, Recurrent Cross-Attention is able to thoroughly explore the representational granularity, which mirrors the process of hierarchical clustering:
$$
\boldsymbol {C} ^ {l} = \operatorname {R C A} ^ {l} \left(\boldsymbol {I} ^ {l}, \boldsymbol {C} _ {0} ^ {l}\right), \tag {4}
$$
where RCA stands for the recurrent cross-attention layer.  $I^l$  is the image feature map at different layers by standard pooling operation with  $H / 2^l\times W / 2^l$  resolution.  $C^l$  is the cluster center matrix for  $l^{th}$  layer and  $C_0^l$  is the initial centers at  $l^{th}$  layer. The parameters for Recurrent Cross-Attention at different layers, i.e.,  $\{\mathrm{RCA}^l\}_{l = 1}^L$ , are not shared. In addition, we initialize the centers from image grids:
$$
\left[ \boldsymbol {c} _ {1} ^ {(0)}; \dots ; \boldsymbol {c} _ {K} ^ {(0)} \right] = \operatorname {F F N} (\text {A d p t i v e} _ {-} \text {P o o l i n g} _ {K} (\boldsymbol {I})) \tag {5}
$$
where FFN stands for Position-wise Feedforward Network which is an integral part of the Transformer architecture. It comprises two fully connected layers along with an activation function used in the hidden layer. AdaptivePooling $_K(I)$  refers to select  $K$  feature centers from  $I$  using adaptive sampling, which calculates an appropriate window size to achieve a desired output size adaptively, offering more flexibility and precision compared to traditional pooling methods.
Feature Dispatching. After the cluster assignment, the proposed method employs an adaptive process that dispatches each patch within a cluster based on similarity (see Fig. 2(c)), leading to a more coherent and representative understanding of the overall structure and context within the cluster. For every patch embedding  $p_i \in I$ , the updated patch embedding  $p_i'$  is computed as:
$$
p _ {i} ^ {\prime} = p _ {i} + \operatorname {M L P} \left(\frac {1}{K} \sum_ {k = 0} ^ {K} \operatorname {s i m} \left(C _ {k}, p _ {i}\right) * C _ {k}\right) \tag {6}
$$
This equation represents the adaptive dispatching of feature embeddings by considering the similarity between the feature embedding and the cluster centers  $(C)$ , weighted by their respective similarities. By incorporating the intrinsic information from the cluster centers, the method refines the feature embeddings, enhancing the overall understanding of the image's underlying structure and context. All feature representations are utilized for handling the target tasks in the decoding process. In §3.3, we discuss more details about the implementation of the ending tasks.
3.3 Implementation Details
The implementation details and framework of CLUSTERFORMER are shown in (Fig. 2a). We followed the architecture and configuration of Swin Transformer [53]. The code will be available at here.
- Encoder: The encoding process is to generate presentation hierarchy, denoted as  $\{I^l\}$  with  $l = \{1,2,3,4\}$ , for a given image  $I$ . The pipeline begins with the feature embedding to convert the images into separate feature tokens. Subsequently, multi-head computing [81, 53] is employed to partition the embedded features among them. Center initialization (Eq. 5) is then adopted as a starting for initializing the cluster centers, and the recurrent cross-attention clustering (Eq. 3) is utilized to recursively update these centers. Once the centers have been updated, the features are dispatched based on their association with the updated centers (Eq. 6). The further decoding process leverage both the centers and the features, which guarantees well-rounded learning.
- Adaptation to Image Classification. The classification head is a single-layer Multilayer Perceptron (MLP) takes the cluster centers from the encoder for predictions.
- Adaptation to Detection and Segmentation. Downstream task head has six Transformer decoder layers with the core design of recurrent cross-attention clustering (Eq.4). Each layer has 3 iterations.
We evaluate our methods over five vision tasks viz. image classification, object detection, semantic segmentation, instance segmentation, and panoptic segmentation on four benchmarks.
ImageNet-1K for Image Classification. ImageNet-1K[72] includes high-resolution images spanning distinct categories (e.g., animals, plants, and vehicles). Following conventional procedures, the dataset is split into 1.2M/50K/100K images for train/validation/test splits.
MS COCO for Object Detection and Instance Segmentation. COCO [49] dataset features dense annotations for 80 common objects in daily contexts. Following standard practices [49], the dataset is split into 115K/5K/20K images for train2017/val2017/test-dev splits.
ADE20K for Semantic Segmentation. ADE20K [101] dataset offers an extensive collection of images with pixel-level annotations, containing 150 diverse object categories in both indoor and outdoor scenes. The dataset comprises 20K/2K/3K images for train/val/test splits.
COCO Panoptic for Panoptic Segmentation. The COCO Panoptic dataset [42] includes 80 "thing" categories and a carefully annotated set of 53 "stuff" categories. In line with standard practices [42], the COCO Panoptic dataset is split into  $115\mathrm{K} / 5\mathrm{K} / 20\mathrm{K}$  images for the train/val/test splits as well.
The ensuing section commences by presenting the main results of each task (§4.1), succeeded by a series of ablative studies (§4.2), which aim to confirm the efficacy of each modulating design.
4.1.1 Experiments on Image Classification
Training. We use mmclassification² as codebase and follow the default training settings. The default configuration for our model involves setting the number of centers to 100. To optimize the model's performance, we employ cross-entropy as the default loss function, which is widely used in classification tasks and helps in minimizing the difference between predicted probabilities and ground truth. For the training details, we run the model for 300 epochs, allowing sufficient time for the model to learn and converge. To manage the learning rate, we initialize it at 0.001 as default. The learning rate is then scheduled using a cosine annealing policy, which gradually decreases the learning rate over time. Due to limitations in our GPU capacity, we are constrained to set the total batch size at 1024. Models are trained from scratch on sixteen A100 GPUs.
Results on ImageNet. Table 1 illustrates our compelling results over different famous methods. CLUSTERFORMER exceeds the Swin Transformer [53] by  $0.13\%$  and  $0.39\%$  on Tiny-based and Small-based models with fewer parameters (i.e., 27.85M vs. 28.29M and 48.71M vs. 49.61M), respectively. On top-5 accuracy, our approach also outperforms the Swin-Tiny and Swin-Small with gains of  $0.71\%$  and  $0.84\%$ , respectively. In addition, our margins over the ResNet family [34] are  $3.44\% \sim 4.76\%$  on top-1 accuracy with on-par parameters (i.e., 27.85M vs. 25.56M and 48.71M vs. 44.55M).
Table 1: Classification top-1 and top-5 accuracy on ImageNet [72] val (see §4.1.1 for details).
Method	#Params	top-1	top-5
Context Cluster-Tiny[ICLR23][58]	5.3M	71.68%	90.49%
DeiT-Tiny[ICML21][80]	5.72M	74.50%	92.25%
PViG-Tiny[NeurIPS22][31]	9.46M	78.38%	94.38%
ResNet-50[CVPR2016][34]	25.56M	76.55%	93.06%
Swin-Tiny[ICCV2021][53]	28.29M	81.18%	95.61%
CLUSTERFORMER-Tiny	27.85M	81.31%	96.32%
Context Cluster-Small[ICLR23][58]	14.0M	77.42%	93.69%
DeiT-Small[ICML21][80]	22.05M	80.69%	95.06%
PViG-Small[NeurIPS22][31]	29.02M	82.00%	95.97%
ResNet-101[CVPR2016][34]	44.55M	77.97%	94.06%
Swin-Small[ICCV2021][53]	49.61M	83.02%	96.29%
CLUSTERFORMER-Small	48.71M	83.41%	97.13%
4.1.2 Experiments on Object Detection