SlowGuess/NeurIPS_2023_OCR huggingface dataset:
By rewriting  $1 - \delta_{k(z^{+})k(z)}$  in Eq. [10] as  $\sum_{k=1}^{K} \delta_{kk(z^{+})} - \delta_{kk(z^{+})} \delta_{kk(z)}$ , and with the relaxed hard cluster assignment via a classification head, the objective for the inter-view clustering can be
expressed by
$$
\min  _ {\mathcal {M}} \frac {1}{N ^ {\prime}} \sum_ {\hat {z} \in \hat {\mathcal {Z}}} \mathrm {H} (\boldsymbol {p} (\hat {z} ^ {+}), \boldsymbol {p} (\hat {z})) + D _ {\mathrm {K L}} (\bar {\boldsymbol {p}} \| \boldsymbol {\pi}), \tag {11}
$$
where  $\pmb{p}(\pmb{x}) = \mathrm{softmax}\left(\tau' \pmb{W}_C^\top \pmb{x}\right)$ ,  $\tau' \ll 1$  with  $\pmb{W}_C$  defined as a matrix of  $K$  orderly concatenated centroids, and  $\mathrm{H}(x, y) = -x \log y$  (cf. Appendix C).
Positive sample retrieval Unlike the common instance-level SSL, the positive samples in FLSL are amorphous clusters of features,  $(\tilde{z}^{+},\tilde{z})$ , corresponding to the same semantic concept in two views. In contrast to previous works assigning the best-matching patch [38, 63] or thresholded vicinity [73], we leverage the cluster assignment mechanism inherent in mean-shift, where a query  $z$  is automatically assigned to a cluster represented by the return  $\hat{z}$ . For query from another view, the mean-shift naturally manifests as a cross-attention (CA),
$$
\hat {\boldsymbol {z}} ^ {+} = \boldsymbol {Z} ^ {+} \operatorname {s o f t m a x} \left(\tau \boldsymbol {z} ^ {\top} \boldsymbol {Z} ^ {+}\right), \tag {12}
$$
With representations semantically coherent on local and global levels, the returned  $\hat{z}^{+}$  from the augmented view  $Z^{+}$  by query  $z$  should agree with the returned  $\hat{z}$  from the original view. To help establish this semantic constraint, representations at the projected positions from the augmented view can be used as positive samples at the early stage of training. This process can be viewed as data retrieval in dense associative memory recognized in [48].
4.3 FLSL Objective
By combining the objectives from the two clustering levels, we arrive at the objective of FLSL:
$$
\min  \frac {1}{N ^ {\prime}} \sum_ {\boldsymbol {Z} \in \mathcal {Z}} \sum_ {\boldsymbol {z} \in \boldsymbol {Z}} v \| \boldsymbol {z} - \hat {\boldsymbol {z}} \| _ {2} ^ {2} + \eta \sum_ {\boldsymbol {z} \in \boldsymbol {Z}} \mathrm {H} (\boldsymbol {p} (\hat {\boldsymbol {z}} ^ {+}), \boldsymbol {p} (\hat {\boldsymbol {z}})) + \gamma D _ {\mathrm {K L}} (\bar {\boldsymbol {p}} \| \boldsymbol {\pi}), \tag {13}
$$
$$
\text {w i t h} \hat {z} = \operatorname {S A} (z, Z, Z), \hat {z} ^ {+} = \operatorname {C A} (z, Z ^ {+}, Z ^ {+}),
$$
where  $\upsilon, \eta$  and  $\gamma$  are the hyperparameters controlling the importance of each term, and the SA and CA above are non-parametric.
Figure 2 illustrates the FLSL framework. We follow the common joint-embedding strategy of SSL, except that we simultaneously maximize the agreement between positive cluster representatives  $(\pmb{p}(\hat{\pmb{z}}^{+}),\pmb{p}(\hat{\pmb{z}}))$  and the agreement between an in-cluster point and its cluster representative  $(z,\hat{z})$ . The KL-divergence term in Eq. 13 serves as a volume maximization regularizer. Experiments show that the FLSL objective effectively promote locally and globally semantic representations, resulting in significantly improved transferability of learned features to object detection and segmentation. Note that FLSL does not involve a class token in its objective (Eq. 13).
In this section, we evaluate the performance of FLSL by conducting extensive experiments. Specifically, we compare FLSL to existing SSL approaches on multiple dense prediction benchmarks: (i) MS-COCO [42] object detection and instance segmentation, (ii) UAVDT [23] object detection from UAV platforms, and (iii) DAVIS video instance segmentation [46]. Moreover, we investigate the properties of FLSL features in terms of semantic alignment and feature separability in the embedding space. Detailed experimental setups are provided in the respective subsections and supplementary materials. All our experiments are performed on Nvidia RTX A6000.
Implementation details The implementation of ViT in our experiments mostly follows DeiT [54] excluding the [class] token. The configuration of the ViT variants utilized in this paper is summarized in Appendix E.3 The coefficients of Eq. [13] in our experiments are  $v = .03$ ,  $\eta = 1$  and  $\gamma = 5$  unless stated otherwise. We assume a uniform prior, i.e.,  $\pi_k = 1 / K$ ,  $\forall k$ . Models are pretrained on ImageNet-1k [52] dataset using AdamW optimizer [45] with a batch size of 512. We follow the data augmentation from BYOL [27] (e.g., color jittering of brightness, contrast, saturation and hue, Gaussian blur and solarization) with preceding random crops and resizing (to  $224 \times 224$ ) and make them asymmetric. Computation among dense features can be expensive. Therefore, we apply a grid random sampling to the queries. All ViT models are pretrained for 300 epochs as in most baselines for a fair comparison. Pseudo-code, training details, and settings of augmentation pipeline are provided in Appendix E.
Pretrain	Backbone	Epoch	#Params	\(AP^{bbox}\)	\(AP^{bbox}\)	\(AP^{bbox}\)	\(AP^{mk}\)	\(AP^{mk}\)	\(AP^{mk}\)
MoCo-v2	RN50	200	23M	38.9	59.2	42.4	35.5	56.2	37.8
DetCo	RN50	200	23M	40.1	61.0	43.9	36.4	58.0	38.9
DenseCL	RN50	200	23M	40.3	59.9	44.3	36.4	57.0	39.2
BYOL	RN50	1000	23M	40.4	61.6	44.1	37.2	58.8	39.8
SCRL	RN50	1000	23M	41.3	62.4	45.0	37.7	59.6	40.7
MOCO-v3	ViT-S/16	300	21M	39.8	62.6	43.1	37.1	59.6	39.2
MoBY	ViT-S/16	300	21M	41.1	63.7	44.8	37.6	60.3	39.8
DINO	ViT-S/16	300	21M	40.8	63.4	44.2	37.3	59.9	39.5
DINO+SelfPatch	ViT-S/16	200	21M	42.1	64.9	46.1	38.5	61.3	40.8
ADCLR	ViT-S/16	300	21M	44.3	65.4	47.6	39.7	62.1	41.5
FLSL	ViT-S/16	300	21M	44.9	66.1	48.1	40.8	64.7	44.2
FLSL	ViT-S/8	300	21M	46.5	69.0	51.3	42.1	65.3	45.0
Table 1: MASK R-CNN ON COCO
Pretrain	\( AP^{bbox} \)	\( AP_{s}^{bbox} \)	\( AP_{m}^{bbox} \)	\( AP_{l}^{bbox} \)	\( AP^{mk} \)
None	48.1	-	-	-	42.6
IN-1k Supv.	47.6	-	-	-	42.4
IN-21k Supv.	47.8	-	-	-	42.6
IN-1k DINO	48.9	32.9	52.2	62.4	43.7
IN-1k MAE	51.2	34.9	54.7	66.0	45.5
IN-1k FLSL	53.1	36.9	56.2	67.4	47.0
Table 2: VITDET-B/16 WITH MASK R-CNN ON COCO
Pretrain	Backbone	APVOC
IN-1k DINO	ViT-S/16	48.9
IN-1k DINO	ViT-B/16	49.1
IN-1k DINO	ViT-S/8	51.1
IN-1k FLSL	ViT-S/16	53.1
IN-1k FLSL	ViT-B/16	53.5
IN-1k FLSL	ViT-S/8	55.2
Table 3: FASTER R-CNN FPN ON UAVDT
Baselines We compare FLSL with various existing SSL approaches that are based on the ResNet [31] and ViT [22] architectures: (a) self-supervised ResNet: MoCo-v2 [15], DetCo [69], DenseCL [63], BYOL [27], and SCRL [50]; and (b) self-supervised ViT: MoCo-v3 [16], MoBY [72], DINO [10], MAE [28], SelfPatch [75], and ADCLR [76].
Protocol for hyperparameter tuning Standard instance-level SSL evaluation protocols typically utilize one of the two approaches: employing a  $k$ -NN classifier or training a linear classifier on fixed features. Since FLSL learns dense semantic representations rather than a single instance-level representation, both standard evaluation protocols are not suitable for evaluating FLSL in training. Moreover, fine-tuning on a downstream dense prediction tasks can be computationally expensive due to complex prediction heads, and may introduce task-specific biases during hyperparameter tuning. Therefore, we design a bbox-aligned  $k$ -NN classifier modified from [67] to evaluate the feature quality directly without additional network tuning. Here is an overview of the method. Features of the training data are first extracted with a fixed model. These features are then aligned with their corresponding bounding boxes provided by ILSVRC [51]. For each image, a certain number of representative features  $\hat{z}$ s (e.g., 9) are selected by a partition criterion and stored in memory. The  $k$ -NN classifier matches each selected features to its  $k$ -nearest stored features, which collectively vote for its label. A feature is considered successfully classified if any of the representative features match its class. This protocol is employed for hyperparameter tuning and ablation study of the FLSL pipeline. Appendix E provides further details on the choice of  $k$ , implementation specifics and evaluation results.
5.1 MS-COCO Object Detection & Segmentation
We adopt Mask R-CNN detection framework by incorporating three variants of ViT: (i) ViT-S/16 with FPN [41], (ii) ViT-S/8 with FPN, and (iii) ViT-B/16 with simple feature pyramid (ViTDet) [40]. Models of (i) and (ii) are fine-tuned following the multi-scale training [66, 6] under the standard  $1 \times$  schedule for a fair comparison. For the model of (iii), we follow the training recipe of [40] and fine-tune the model for 100 epochs.
Results. Table 1 reports the detection and segmentation performance of ViT-S/16 and ViT-S/8 with Mask R-CNN [30] on COCO. Specifically, FLSL with ViT-S/16 outperforms ADCLR [76] by  $+0.6\%$  and  $+1.1\%$ , and substantially outperforms DINO+SelfPatch [75] by  $+2.8\%$  and  $+2.4\%$  on detection  $(\mathrm{AP}^{\mathrm{bbox}})$  and segmentation  $(\mathrm{AP}^{\mathrm{mk}})$ , respectively. Both baseline methods feature patch-level contrastive learning. Unlike SelfPatch contrasting between patches within the adjacent neighborhood and ADCLR contrasting via learned queries of random crops, FLSL contrasts the representatives (modes) of feature clusters, which aligns closer with the downstream tasks and thus leads to superior performance. Notably, FLSL with ViT-S/8 further improves the performance by a large margin of  $+4.4\%$  in  $\mathrm{AP}^{\mathrm{bbox}}$  and  $+3.6\%$ $\mathrm{AP}^{\mathrm{mk}}$  over SelfPatch. Table 2 summarizes the results of ViTDet. FLSL shows large performance gains over the DINO baseline by  $+4.2\%$ $\mathrm{AP}^{\mathrm{bbox}}$  and  $+3.3\%$ $\mathrm{AP}^{\mathrm{mk}}$ . FLSL also outperforms the SOTA generative approach, MAE, by  $+1.7\%$  and  $+1.4\%$  in the two tasks, respectively.
Figure 3: Visualization of the maps of the aggregated attention score (ASS) from different layers of ViT-S/16.  $l = 0$  denotes the projection layer. As layer goes deeper, the map becomes more partitioned with brightness aligned with the area of the underlying semantic region, e.g., objects or stuff.
Pretrain	Arch. (J&F)m	Jm	Fm
IN-1k supv.	ViT-S/8	66.0	63.9	68.1
VLOG CT	RN50	48.7	46.4	50.0
YT-VOS MAST	RN18	65.5	63.3	67.6
IN-1k DINO	ViT-S/16	61.8	60.2	63.4
IN-1k DINO	ViT-B/16	62.3	60.7	63.9
IN-1k DINO	ViT-S/8	69.9	66.6	73.1
IN-1k FLSL	ViT-S/16	65.6	62.4	69.4
IN-1k FLSL	ViT-B/16	66.1	62.9	70.0
IN-1k FLSL	ViT-S/8	73.5	69.9	78.1
Table 4: DAVIS 2017 VIDEO INSTANCE SEGMENTATION. We evaluate the quality of frozen features on video instance tracking. We report mean region similarity  ${\mathcal{I}}_{m}$  and mean contour-based accuracy  ${\mathcal{F}}_{m}$  .
5.2 Small Object Detection: UAVDT
To assess the transferability of FLSL beyond the datasets of common images like COCO, we further investigate its performance on a UAV benchmark, UAVDT [23], which exhibits significant domain shifts from common images (i.e., images captured by ground-level cameras). We utilize Faster R-CNN framework [49] with the same ViT variants used in the COCO experiments and follow the training settings outlined in ClusDet [74]. All ViT-backboned models are trained with  $1 \times$  schedule.
Result Table presents the performance of ViT-S/16, ViT-S/8, and ViT-B/16 with Faster R-CNN for detection tasks on UAVDT under different pretrain schemes. We utilize the official evaluation method in [23], which calculates the class-agnostic VOC AP exclusive of the predictions that falls in the ignored areas. FLSL consistently outperforms DINO (a typical instance-level SSL for ViT) across all three ViT variants by a significant margin. With smaller objects and an imbalanced foreground-background ratio, the significance of local semantics becomes evident. Models require local context to discover small objects and make accurate predictions rather than relying solely on the global semantics of the entire image. This situation aligns well with the strengths of FLSL.
5.3 DAVIS Segmentation
To further assess the quality of frozen features learned by FLSL, we evaluate FLSL-pretrained ViT models on DAVIS2017 [46], following the evaluation protocol in [36, 10] that requires fixed representations with no extra training.
Results Table 4 shows that FLSL consistently outperforms DINO across all ViT variants in our experiments. The protocol evaluates the quality of learned dense features via segmenting scenes with  $k$ -nearest neighbors ( $k = 5$ ) within a fixed window ( $12 \times 12$ ) between consecutive frames. This requires dense features to be locally semantic, i.e., features corresponding to the same semantics should be more correlated. Therefore, the improved performance confirms that FLSL encourages model to extract locally semantic representations.
5.4 Alignment with Image Semantics
To qualitatively show that FLSL is better aligned with the semantic layout of an image than the common SSL methods, Figure 4(a) compares the self-attention probing maps for features learned via FLSL
and DINO. Features from the last layer are used for evaluation. The visualizations are obtained with  $224^{2}$  images. Positions of the query tokens are marked out in green circle in the top row. As shown in the middle and bottom rows of the figure, DINO promotes more correlated attention (i.e., less separation between tokens of query-related area and that of the rest image), while FLSL encourages
Figure 4: (a) visualization of attention probing by query patches (marked out in green circle in the top row) from the last layer of ViT-S/16 pretrained with FLSL and with DINO. FLSL encourages the model to learn semantic correlations among patches; (b) visualization of separability of the dense representations throughout the transformer (ViT-S/16).
Sinkhorn	η	γ	v = 0.0	v = .01	v = .02	v = .03	~	v = .1
✓	1.0	1.0	0.1	68.7	70.7	71.2	~	65.1
×	1.0	1.0	-	-	-	66.6	-	-
×	1.0	5.0	-	-	-	72.4	-	-
Table 5: IMPACT OF COEFFICIENTS IN THE FLSL OBJECTIVE.
K	1024	2048	4096	8192	16384
k-NN top-1	68.1	72.1	72.4	72.5	72.1
Table 6: IMPACT OF NUMBER OF CENTROIDS  $K$
attention to the regions of high semantic relevance with the query tokens and results in clearer maps consistent with the underlying objects/stuff.
5.5 Feature Distribution and Separability
We demonstrate the qualitative results by visualizing the Aggregated Similarity Score (ASS) and the feature distribution in the embedding space using t-sne [57] in Figure 3 and Figure 4(b), respectively. To generate the map of ASS, we sum up the cosine-similarity maps of all tokens, normalize the resulting map with its maximum score and visualize it as a thermal image, i.e., the brighter the pixel, the higher the score. For a semantically well-separated image, each patch only attends to the patches of its own semantic region, e.g., a patch of an object has high similarity scores only with the patches of that object and low scores with the rest. This results in an image with partitions of different brightness proportional to the area of that region, i.e., ideally the larger the size of an object/stuff, the brighter the color. As shown in Figure 3 as the layer goes deeper, the brightness partition of the ASS is more consistent with the underlying objects and stuff in the images (e.g., person, vehicles, horse, switches, wall, and ground, etc.), which indicates the desired separation of the learned features. This is also reflected in the t-sne visualization of the embeddings in Figure 4(b), where the representations become more clustered and separated as the attention layer goes deeper.
5.6 Ablation Study
Due to limited space, we present two major ablation studies in this section to help understand the effectiveness of FLSL. The model considered for this entire study is ViT-S trained with 100 epochs. We refer the reader to Appendix [1] for the complete work.
Impact of coefficients in the FLSL objective The FLSL objective (Eq. [13] contains three components: (1) similarity between  $\ell_2$ -normalized  $z$  (features) and  $\hat{z}$  (modes), (2) cross-entropy of the probabilities of an augmented pair  $H(p(\hat{z}^+), p(\hat{z}))$ , and (3) the volume maximization regularizer  $\bar{D}_{\mathrm{KL}}(\bar{p} \| \pi)$ . It is computationally expensive to optimally determine the values of more than two coefficients by performing grid search, especially when the ratios among them are large. We tackle this problem by first fixing  $\eta = 1$  and setting  $\gamma = 1$  along with Sinkhorn normalization [19] to perform a grid search on the value of  $v$  with the empirical base condition  $v \leq 1$  and  $\gamma \geq 1$  [75]. With the fixed  $v$ , we then perform another grid search on  $\gamma$  without Sinkhorn normalization. We implement Sinkhorn normalization as the softmax operation along the batch dimension. Table 5 summarizes the score of bbox-aligned  $k$ -NN evaluation using different coefficient settings.
Impact of number of centroids  $K$  FLSL is formulated as an explicit clustering problem, with the output dimension of the last fully-connected layer equal to the number of centroids  $K$ . Compared to its instance-level counterpart DINO [10], FLSL enjoys a smaller output dimension (shown in Table 6). This is because images have higher feature variance compared to feature clusters. For example, an image in ImageNet may contain diverse content from different categories, requiring a large number of centroids to cover the distribution. In contrast, a semantic cluster contains highly correlated features, such as similar textures or objects from the same category, thus requiring fewer centroids. Experimentally, we find that a large number of centroids benefits performance, but is detrimental and costly when being too large. We pick  $K = 4$ , 096 for all our experiments as it strikes a good balance between performance and cost-effectiveness.
More experiment results on semantic segmentation and ablations including the impact of batch size and random pooling window size are relegated to Appendix.
This paper proposes FLSL, a feature-level self-supervised learning method that bridges the gap between the current SSL methods and downstream dense prediction tasks. We demonstrate for the first time the underlying mean-shift clustering process of ViT, which aligns well with natural image semantics. Facilitated by ViT for joint embedding and feature clustering, FLSL performs a bi-level clustering: (i) intra-view clustering to extract the representatives for clusters of features within an image, and (ii) inter-view clustering to encourage the representatives to be globally semantic over
the entire dataset. FLSL achieves a significant improvement over the SOTAs in the dense prediction tasks, including object detection and instance segmentation.
Limitations and broader impacts FLSL does not have any significant limitations other than the method is more complex (due to its bi-level clustering) than other SSL methods, and it currently only fits for ViT-based models on dense prediction tasks. Exploring ways to extend FLSL for tasks that necessitate a global representation while retaining its existing properties could be a potential future work. As far as we can foresee, there is no negative societal impact.
This research was sponsored by the Army Research Laboratory under Cooperative Agreement #W911NF-22-2-0025. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the Army Research Laboratory or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.
[1] Mahmoud Assran, Mathilde Caron, Ishan Misra, Piotr Bojanowski, Florian Bordes, Pascal Vincent, Armand Joulin, Mike Rabbat, and Nicolas Ballas. Masked siamese networks for label-efficient learning. In European Conference on Computer Vision, pages 456-473. Springer, 2022.
[2] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473, 2014.
[3] Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Variance-invariance-covariance regularization for self-supervised learning. arXiv preprint arXiv:2105.04906, 2021.
[4] Paul S Bradley, Kristin P Bennett, and Ayhan Demiriz. Constrained k-means clustering. Microsoft Research, Redmond, 20(0):0, 2000.
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[6] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European conference on computer vision, pages 213-229. Springer, 2020.
[7] Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze. Deep clustering for unsupervised learning of visual features. In ECCV, 2018.
[8] Mathilde Caron, Piotr Bojanowski, Julien Mairal, and Armand Joulin. Unsupervised pre-training of image features on non-curated data. In ICCV, 2019.
[9] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Advances in Neural Information Processing Systems, 33:9912-9924, 2020.
[10] Mathilde Caron, Hugo Touvron, Ishan Misra, Herve Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9650-9660, 2021.
[11] Miguel Carreira-Perpínán. Reconstruction of sequential data with probabilistic models and continuity constraints. Advances in neural information processing systems, 12, 1999.
[12] Frédéric Chazal, Leonidas J Guibas, Steve Y Oudot, and Primoz Skraba. Persistence-based clustering in riemannian manifolds. Journal of the ACM (JACM), 60(6):1-38, 2013.