SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ziyu Wang $^{1,2,3}$  Mike Zheng Shou $^{1,\dagger}$  Mengmi Zhang $^{2,3,\dagger}$
<sup>1</sup>Show Lab, National University of Singapore, Singapore
$^{2}$ Deep NeuroCognition Lab, CFAR and I2R, Agency for Science, Technology and Research, Singapore
†Corresponding authors; address correspondence to mengmi@i2r.a-star.edu.sg
Learning object-centric representations from complex natural environments enables both humans and machines with reasoning abilities from low-level perceptual features. To capture compositional entities of the scene, we proposed cyclic walks between perceptual features extracted from vision transformers and object entities. First, a slot-attention module interfaces with these perceptual features and produces a finite set of slot representations. These slots can bind to any object entities in the scene via inter-slot competitions for attention. Next, we establish entity-feature correspondence with cyclic walks along high transition probability based on the pairwise similarity between perceptual features (aka "parts") and slot-binded object representations (aka "whole"). The whole is greater than its parts and the parts constitute the whole. The part-whole interactions form cycle consistencies, as supervisory signals, to train the slot-attention module. Our rigorous experiments on seven image datasets in three unsupervised tasks demonstrate that the networks trained with our cyclic walks can disentangle foregrounds and backgrounds, discover objects, and segment semantic objects in complex scenes. In contrast to object-centric models attached with a decoder for the pixel-level or feature-level reconstructions, our cyclic walks provide strong learning signals, avoiding computation overheads and enhancing memory efficiency. Our source code and data are available at: link.
Object-centric representation learning refers to the ability to decompose the complex natural scene into multiple object entities and establish the relationships among these objects [29, 24, 14, 38, 39, 37, 22]. It is important in multiple applications, such as visual perception, scene understanding, reasoning, and human-object interaction [45, 2]. However, learning to extract object-centric representations from complex natural scenes in an unsupervised manner remains a challenge in machine vision.
Recent works attempt to overcome this challenge by relying on image or feature reconstructions from object-centric representations as supervision signals [29, 38, 22, 37] (Figure 1). However, these reconstructions have several caveats. First, these methods often require an additional decoder network, resulting in computation overheads and memory inefficiency. Moreover, these methods focus excessively on reconstructing unnecessary details at the pixel or feature levels and sometimes fail to capture object-centric representations from a holistic view.
To mitigate issues from reconstructions, some studies [47, 19] introduce mutual information maximization between predicted object-centric representations and feature maps. Building upon this idea, contrastive learning has become a powerful tool in unsupervised object-centric learning. The recent works [21, 5] propose to learn the spatial-temporal correspondences between a sequence of video frames with contrastive random walks. The objective is to maximize the likelihood of returning to
(a) SA, BO-QSA
(b) DINOSAUR
(c) SLATE
Figure 1: Comparison of the network structures trained with existing methods and our cyclic walks. Supervision signal is indicated by  $(\approx)$  in subplots a-c. From left to right, all the methods require some forms of reconstructions as supervisory signals except for our method (d): (a) image reconstruction for the Slot-Attention (SA) [29] and BO-QSA [22] (b) feature reconstruction for DINOSAUR [37](c) tokens obtained by a Discrete Variational Autoencoder [36] (DVAE) for SLATE [38] (d). Our cyclic walks do not require any image or feature reconstructions.
(d) Cyclic walks
the starting node by walking along the graph constructed from a palindrome of frames and repelling nodes with distinct features from adjacent frames in the graph.
Generalizing from contrastive random walks on video sequences, we proposed cyclic walks on static images. These walks are conducted between predicted object entities and feature maps extracted from vision transformers. The cycle consistency of these walks serves as supervision signals to guide object-centric representation learning. First, inspired by DETR [8] and Mask2Former [11], where these methods decompose the scene into multiple object representations or semantic masks with a span of a finite set of task-dependent semantically meaningful bases, we use a slot-attention module to interface with feature maps of an image and produce a set of object-centric representations out of the slot bases. These slot bases compete with one another to bind with feature maps at any spatial location via normalizing attention values over all the slot bases. All the features resembling a particular slot basis are aggregated into one object-centric representation explaining parts of the image.
Next, two types of cyclic walks, along high transition probability based on the pairwise similarity between aggregated object-centric representations and image feature maps, are established to learn entity-feature correspondence. The motivation of the cyclic walks draws a similar analogy with the part-whole theory stating that the whole (i.e. object-centric representations) is greater than its parts (i.e. feature maps extracted from natural images) and the parts constitute the whole. We use "parts" and "whole" for easy illustrations of the interaction between image feature maps and object-centric representations. On one hand, cyclic walks in the direction from the whole to the parts and back to the whole (short for "W-P-W walks") encourage diversity of the learnt slot bases. On the other hand, cyclic walks in the direction from the parts to the whole and back to the parts (short for "P-W-P walks") broaden the coverage of the slot bases pertaining to the image content. Both W-P-W walks and P-W-P walks form cycle consistencies and serve as supervision signals to enhance the specificity and diversity of learned object-centric representations. Our contributions are highlighted below:
- We introduce cyclic walks between parts and whole to regularize object-centric learning. Both W-P-W walks and P-W-P walks form virtuous cycles to enhance the specificity and diversity of the learned object-centric representations.
- We verify the effectiveness of our method over seven image datasets in three unsupervised vision tasks. Our method surpasses the state-of-the-art by a large margin.
- Compared with the previous object-centric methods relying on reconstruction losses, our method does not require additional decoders in the architectures, greatly reducing the computation overheads and improving memory efficiency.
One representative work [29] introduces the slot-attention module, where the slots compete to bind with certain regions of the image via a cross-attention mechanism and distill image features into object-centric representations. To train these slots, the networks are often attached to a decoder to
(a) Network architecture of our cyclic walks.
(b) W-P-W walks.
(c) P-W-P walks.
Figure 2: Overview of our proposed cyclic walks. The input image is divided into non-overlapped patches and encoded into feature vectors (aka "parts") through a frozen DINO ViT pre-trained on ImageNet with unsupervised learning methods [9]. The Slot-Attention module aggregates the patch features into object-centric representations (aka "whole") based on the learnable slot bases. To train the network, we perform the contrastive random walks in two directions between the parts and the whole: (b) any node from the whole should walk to the parts and back to itself (W-P-W walks) and (c) any node from the parts should walk to the whole and back to itself (P-W-P walks). See Fig A7 for applications of these learned object-centric representations in downstream tasks.
reconstruct either images or features based on object-centric representations (Figure 1). Subsequent Slot-Attention methods [6, 39, 24, 14, 4, 13] expand to video processing. To learn object-level representations, these methods often require either reconstructions from video frames or optical flows [4, 13]. Although the original slot attention was demonstrated to be effective in most cases, it fails to generalize to complex scenes [37, 22]. To mitigate this issue, DINOSAUR freezes the feature extractor while learning the slot bases. Its frozen feature extractor was pre-trained on ImageNet in unsupervised learning [9]. BO-QSA [22] proposes to initialize the slots with learnable queries and optimize these queries with bi-level optimization. Both works emphasize that good initialization of network parameters is essential for object-centric learning in complex scenes. Thus, following this line of research, we also use a frozen DINO [9] as the feature extractor. However, different from these works, we introduce cyclic walks in the part-whole hierarchy of the same image, without decoders for reconstructions.
Object-centric learning for unsupervised semantic segmentation. One of the key applications in object-centric representation learning is unsupervised semantic segmentation. Several notable works include STEGO [18], DeepCut [1], ACSeg [26], and FreeSOLO [42]. All these methods employ a fixed pre-trained transformer as a feature extractor, which has proven to be effective for unsupervised semantic segmentation. In contrast to these methods which directly perform optimization on pixel clusters or feature clusters, we introduce slot-attention and perform cyclic walks between slots and feature maps. Experimental results demonstrate the superior performance of our method.
Relation to contrastive learning. Our method connects object-centric representation learning with unsupervised contrastive learning. We highlight the differences between our method and the contrastive learning methods. Infoseg [19] reconstructs an image from a set of global vectors (aka "slot bases"). The reconstructed image is pulled closer to the original image and repelled away from a randomly selected image. SlotCon [44] introduces contrastive learning on two sets of slot bases extracted from two augmented views of the same image. In contrast to the two methods, our cyclic walks curate both positive and negative pairs along the part-whole hierarchy from the same image. Another method, SAN [47], performs contrastive learning between slot bases to encourage slot diversity, where every slot is repelled away from one another. Sharing similar motivations, our walks from the whole to the parts and back to the whole ("W-P-W" walks) strengthen the diversity of the slot bases. Moreover, our method also introduces "P-W-P" walks, which takes the holidistic view of the entire image, and broadens the slot coverage.
Here, we mathematically formulate Slot-Attention Module and Contrastive Random Walks.
3.1 Slot-Attention Module
The slot-attention module takes an input image and outputs object-centric representations. Given an input image, a pre-trained CNN or transformer extracts its feature vectors  $x \in R^{N \times D_{input}}$ , where  $N = W \times H$ .  $H$  and  $W$  denote the height and width of the feature maps. Taking  $x$  and a set of  $K$  slot bases  $s \in R^{K \times D_{slot}}$  as inputs, Slot-Attention binds  $s$  with  $x$ , and outputs a set of  $K$  object-centric feature vectors  $\tilde{s}$  of the same size as  $s$ . The Slot-Attention module employs the cross-attention mechanism [40], where  $x$  contributes to keys and values and  $s$  contributes to queries. Specifically, the inputs  $x$  are mapped to the dimension  $D$  with linear functions  $k(\cdot)$  and  $v(\cdot)$  and the slots  $s$  are mapped to the same dimension with linear function  $q(\cdot)$ . For each feature vector at every spatial location of the feature maps  $x$ , attention values are calculated with respect to all slot bases. To prevent parts of the image from being unattended, the attention matrix is normalized with the softmax function first over  $K$  slots and then over  $N$  locations:
$$
\operatorname {a t t n} _ {i, j} = \frac {e ^ {A _ {i , j}}}{\sum_ {K} e ^ {A _ {i , j}}} \text {w h e r e} A = \frac {k (x) q (s) ^ {\mathsf {T}}}{\sqrt {D}} \in R ^ {N \times K} \tag {1}
$$
$$
\tilde {s} = W ^ {T} v (x) \text {w h e r e} W _ {i, j} = \frac {\operatorname {a t t n} _ {i , j}}{\sum_ {N} \operatorname {a t t n} _ {i , j}} \tag {2}
$$
The cross-attention mechanism introduces competitions among slot bases for explaining parts of the image. Based on the attention matrix, an object-centric feature vector from  $\tilde{s}$  is distilled by the weighted sum of feature maps  $x$ . In the literature [8, 38, 29] and our method, the slot attention modules iteratively refine the output  $\tilde{s}$  from the slots in a recurrent manner with a Gated Recurrent Unit (GRU) for better performances. We represent the initial slot bases as  $s^0$ . The parameters of  $s^0$  are initialized by randomly sampling from Gaussian distribution with learnable mean  $\mu$  and variance  $\sigma$ . The output  $\tilde{s}$  at time step  $t$  acts as the slot bases and feeds back to the GRU at the next recurrent step  $t + 1$ . Thus, the final object-centric representations  $\hat{s}$  after  $T$  iterations can be formulated as:
$$
\hat {s} = s ^ {T}, \text {w h e r e} s ^ {t + 1} = G R U (s ^ {t}, \tilde {s} ^ {t}) \tag {3}
$$
After obtaining the object-centric representation  $\hat{s}$ , the traditional object-centric learning models decode the images from the slots with a mixture-based decoder or a transformer-based decoder [8]. The training objective of the models is to minimize the Mean Square Error loss between the output of the decoder and the original image at the feature or pixel levels.
3.2 Contrastive Random Walks
A random walk describes a random process where an independent path consists of a series of hops between nodes on a directed graph in a latent space. Without loss of generality, given any pair of feature sets  $a \in R^{m \times d}$  and  $b \in R^{n \times d}$ , where  $m$  and  $n$  are the numbers of nodes in the feature sets and  $d$  is the feature dimension, the adjacency matrix  $M_{a,b}$  between  $a$  and  $b$  can be calculated as their normalized pairwise feature similarities:
$$
M _ {a, b} = \frac {e ^ {f (a) f (b) ^ {T} / \tau}}{\sum_ {n} e ^ {f (a) f (b) ^ {T} / \tau}} \in R ^ {m \times n}, \tag {4}
$$
where  $f(\cdot)$  is the  $l_{2}$ -normalization and  $\tau$  is the temperature controlling the sharpness of distribution with its smaller values indicating sharper distribution.
In the original work [21], random walks serve as supervisory signals to train the object-centric learning models to capture the space-time correspondences from raw videos. The contrastive random walks are formulated as a graph.  $a$  and  $b$  are the feature maps extracted from video frames with either a CNN or a transformer.  $m$  and  $n$  are the numbers of spatial locations on the feature maps and they are often the same across video frames. On the graph, only nodes from adjacent video frames  $F_{t}$  and  $F_{t + 1}$  at time  $t$  and  $t + 1$  share a directed edge. Their edge strengths, indicating the transition probabilities of a random walk between frames, are the adjacency matrix  $M_{F_t,F_{t + 1}}$ . The objective of the contrastive random walks is to maximize the likelihood of a random walk returning to the starting node along the graph constructed from a palindrome video sequence  $\{F_t,F_{t + 1},\dots ,F_T,F_{T - 1},\dots ,F_t\}$ .
4.1 Object-centric Feature Encoding
Following [18, 37], we use a self-supervised vision transformer trained with DINO on ImageNet, as the image feature extractor (Figure 2a). In Section 5.4, we verify that our method is also agnostic to other self-supervised feature learning backbones. Given an input image, DINO parses it into non-overlapped patches and each patch is projected into a feature token. We keep all patch tokens except for the classification token in the last block of DINO. As in Section 3.1, we use the same notation  $x \in R^{N \times D_{input}}$  to denote feature maps extracted from a static image. The Slot-Attention module takes the feature vectors  $x \in R^{N \times D_{input}}$  and a set of K learnable slot bases  $s \in R^{K \times D_{slot}}$  as inputs and produces the object-centric representations  $\hat{s} \in R^{K \times D_{slot}}$  (Equations 1 and 2).
4.2 Whole-Parts-Whole Cyclic Walks
Features of the image patches constitute the objects. The interactions between parts and the whole provide a mechanism for clustering and discovering objects in the scene. Motivated by this, we introduce cyclic walks in two directions: (a) from the whole to the parts and back to the whole (W-P-W walks) and (b) from the parts to the whole and back to the parts (P-W-P walks). Both serve as supervisory signals for learning slot bases  $s$ .
Given feature vectors  $x \in R^{N \times D_{input}}$  of all image patches (aka "parts") and the object-centric representations  $\hat{s} \in R^{K \times D_{slot}}$  (aka "whole"), we apply a linear layer to map both  $x$  and  $\hat{s}$  to be of the same dimension  $D$ . In practice,  $K$  is much smaller than  $N$ . Following the formulations of Contrastive Random Walks [21] in Section 3.2, random walks are conducted along high transition probability based on pairwise similarity from  $\hat{s}$  to  $x$  and from  $x$  to  $\hat{s}$  using Equation 4:
$$
M _ {w p w} = M _ {\hat {s}, x} M _ {x, \hat {s}} \in R ^ {K \times K}, \text {w h e r e} M _ {\hat {s}, x} \in R ^ {K \times N} \text {a n d} M _ {x, \hat {s}} \in R ^ {N \times K} \tag {5}
$$
$M$  is defined as an adjacency matrix. Different from the original work [21] performing contrastive random walks on a palindrome video sequence, here, we enforce a palindrome walk in a part-whole hierarchy of the image. Ideally, if the W-P-W walks are successful, all the bases in the slot-attention module have to establish one-to-one mapping with certain parts of the image. This encourages the network to learn diverse object-centric representations so that each slot basis can explain certain parts of the image and every part of the image can correspond to a unique slot basis. To achieve this in W-P-W walks, we enforce that the probability of walking from  $\hat{s}$  to  $x$  belonging to the object class and back to  $\hat{s}$  itself should be an identity matrix  $I$  of size  $K\times K$ . The first loss term is defined as:  $L_{wpw} = CE(M_{wpw},I)$ , where  $CE(\cdot ,\cdot)$  refers to cross entropy loss.
4.3 Parts-Whole-Parts Cyclic Walks
Though the W-P-W walks enhance the diversity of the slot bases, there is an ill-posed situation, when there exists a limited set of slot bases failing to cover all the semantic content of an image but every trivial W-P-W walk is always successful. For example, given two slot bases and an image consisting of a triangle, a circle, and a square on a background, one slot could prefer "triangles", while the other prefers "circles". In this case, the W-P-W walks are always successful; however, the two slots fail to represent squares and the background, defeating the original intention of foreground and background extraction. To mitigate this problem and bind all the slots with all the semantic content in the entire scene, we introduce additional walks from the parts to the whole and back to the parts (P-W-P walks), complementary to W-P-W walks:  $M_{pwp} = M_{x,\hat{s}}M_{\hat{s},x} \in R^{N \times N}$ . As N is typically far larger than K and features  $x$  at nearby locations tend to be similar, the probabilities of random walks beginning from one feature vector of  $x$ , passing through  $\hat{s}$ , and returning to itself could no longer be 1. Thus, we use the feature-feature correspondence  $S$  as the supervisory signal to regularize  $\hat{s}$ :
$$
S _ {i, j} = \frac {e ^ {W _ {i , j}}}{\sum_ {N} e ^ {W _ {i , j}}} \text {w h e r e} W _ {i, j} = \left\{ \begin{array}{l l} - \infty , & \text {i f} F _ {i, j} <   = \gamma \\ F _ {i, j}, & \text {i f} F _ {i, j} > \gamma \end{array} \right., \text {a n d} F = f (x) f (x) ^ {T} \in R ^ {N \times N} \tag {6}
$$
Hyper-parameter  $\gamma$  is the similarity threshold for preventing the random walks from returning to locations where their features are not as similar as the starting point. Thus, the overall loss of our
Model	Birds	Dogs	Cars	Flowers
mIoU	Dice	mIoU	Dice	mIoU	Dice	mIoU	Dice
Slot-Attention [29]	35.6	51.5	39.6	55.3	41.3	58.3	30.8	45.9
SLATE [38]	36.1	51.0	62.3	76.3	75.5	85.9	68.1	79.1
DINOSAUR [37]	67.2	78.4	73.7	84.6	80.1	87.6	72.9	82.4
BO-QSA [22]	71.0	82.6	82.5	90.3	87.5	93.2	78.4	86.1
Cyclic walks (ours)	72.4	83.6	86.2	92.4	90.2	94.7	75.1	83.9
Table 1: Results in the unsupervised foreground extraction task. We report the results in mIoU and Dice on CUB200 Birds (Birds), Stanford Dogs (Dogs), Stanford Cars (Cars), and Flowers (Flowers) datasets. Numbers in bold are the best.
(a) Foreground extraction
(b) Object discovery
Figure 3: Visualization of attention maps predicted by slot bases. In (a) unsupervised foreground extraction, we present two examples each from Birds, and Flowers datasets. In each example, for all the slots, we provide their attention maps (Col.2-3) and their corresponding attended image regions (Col.4-5). Ground truth foreground and background masks are shown in the last Col. In (b) unsupervised object discovery, we provide 3 examples on the PASCAL VOC dataset. In each example, for all the slots, we provide their attention maps (att., Col.2-5) and their corresponding attended image regions (reg., Col. 6-9). Their combined object discovery masks from all the slots are shown in Col.10 (pred.). Each randomly assigned color denotes an image region activated by one slot. Ground truth object masks are presented in the last column.
cyclic walks can be calculated below, where  $\alpha$  and  $\beta$  are the coefficients balancing the two losses.
$$
L = \alpha L _ {w p w} + \beta L _ {p w p}, \text {w h e r e} L _ {p w p} = C E (M _ {p w p}, S) \tag {7}
$$
Training Configuration. We use ViT-Small pre-trained with DINO as our feature extractor and freeze the parameters of DINO throughout the entire training process. There are multiple reasons for freezing the feature extractor. First, pre-trained unsupervised feature learning frameworks have already generated semantically consistent content in nearby locations [18, 1, 26]. Object-centric learning addresses its follow-up problem of capturing compositional entities out of the semantic content. Second, recent object-centric learning works [37, 22] have emphasized the importance of freezing the pre-trained weights of feature extractors (see Section 2). For a fair comparison with these methods, we keep the parameters of the feature extractor frozen. Third, we provided empirical evidence that our method benefits from freezing DINO, and freezing DINO is not a limitation of our method in Appendix A1. We use a similarity threshold of 0.7 and a temperature of 0.1 for all the experiments over all the datasets. In general, the variations of these two hyperparameters lead to moderate changes in performances (see Sec 5.4). See Appendix A2 for more training and implementation details. All models are trained on 4 Nvidia RTX A5000 GPUs with a total batch size of 128. We report the mean  $\pm$  standard deviation of 5 runs with 5 random seeds for all our experiments.
5.1 Unsupervised Foreground Extraction
Task Setting, Metrics, Baselines and Datasets. In the task, all models learn to output binary masks separating foregrounds and backgrounds in an unsupervised manner. See Appendix A3.1 for implementations of foreground and background mask predictions during the inference. We evaluate the quality of the predicted foreground and background masks with mean Intersection over Union