SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mengcheng Lan $^{1*}$  Xinjiang Wang $^{3}$  Yiping Ke $^{2}$  Jiaxing Xu $^{2}$  Litong Feng $^{3}$  Wayne Zhang $^{3}$
$^{1}$  S-Lab, Nanyang Technological University
$^{2}$  SCSE, Nanyang Technological University
$^{3}$  SenseTime Research
{lanm0002, jiaxing003}@e.ntu.edu.sg ypke@ntu.edu.sg
{wangxinjiang, fenglitong, wayne.zhang}@sensetime.com
https://github.com/mc-lan/SmooSeg
Unsupervised semantic segmentation is a challenging task that segments images into semantic groups without manual annotation. Prior works have primarily focused on leveraging prior knowledge of semantic consistency or priori concepts from self-supervised learning methods, which often overlook the coherence property of image segments. In this paper, we demonstrate that the smoothness prior, asserting that close features in a metric space share the same semantics, can significantly simplify segmentation by casting unsupervised semantic segmentation as an energy minimization problem. Under this paradigm, we propose a novel approach called SmooSeg that harnesses self-supervised learning methods to model the closeness relationships among observations as smoothness signals. To effectively discover coherent semantic segments, we introduce a novel smoothness loss that promotes piecewise smoothness within segments while preserving discontinuities across different segments. Additionally, to further enhance segmentation quality, we design an asymmetric teacher-student style predictor that generates smoothly updated pseudo labels, facilitating an optimal fit between observations and labeling outputs. Thanks to the rich supervision cues of the smoothness prior, our SmooSeg significantly outperforms STEGO in terms of pixel accuracy on three datasets: COCOStuff  $(+14.9\%)$ , Cityscapes  $(+13.0\%)$ , and Potsdam-3  $(+5.7\%)$ .
Semantic segmentation is a crucial task in computer vision that allows for a better understanding of the visual content and has numerous applications, including autonomous driving [1] and remote sensing imagery [2]. Despite advancements in the field, most traditional semantic segmentation models heavily rely on vast amounts of annotated data, which can be both arduous and costly to acquire. Consequently, unsupervised semantic segmentation [3; 4; 5; 6; 7; 8] has emerged as a promising alternative. Prior knowledge is fundamental to the success of unsupervised semantic segmentation models. One key prior knowledge is the principle of semantic consistency, which stipulates that an object's semantic label should remain consistent despite photometric or geometric transformations. Recent advances [3; 9; 4; 10] use contrastive learning to achieve consistent features or class assignments. Another essential prior knowledge is the priori concepts implicitly provided by self-supervised learning techniques, e.g., DINO [11] and precedent arts [12; 8; 6] whose learned features can be employed to partition each image into different segments. Despite their effectiveness,
SmooSeg
Figure 1: A case study of our SmooSeg with two state-of-the-arts, STEGO [8] and TransFGU [6], on the COCOStuff dataset. Our observations reveal that the segmentation maps generated by STEGO and TransFGU for regions such as the sand beach (first row) and the grassland (second row) are incomplete and lack smoothness and coherence. In contrast, our SmooSeg exhibits improved segmentation results for all these regions by considering the smoothness prior.
STEGO
TransFGU
these methods often overlook the coherence property of image segments, resulting in predicted segments that are incomplete and lacking in coherence, as shown in Fig. 1.
Real-world images often demonstrate a natural tendency towards piecewise coherence regarding semantics, texture, or color. Observations close to each other, either in the form of adjacent pixels in the coordinate space or close features in a metric space, are expected to share similar semantic labels, and vice versa. This essential property, known as the smoothness prior, plays a crucial role in various computer vision tasks [13; 14; 15]. Surprisingly, it is still under-explored in the field of unsupervised semantic segmentation.
In this paper, we attempt to tackle unsupervised semantic segmentation from the perspective of smoothness prior. As a dense prediction task, semantic segmentation aims at finding a labeling  $f \in \mathcal{F}$  that assigns each observation (pixel, patch, features)  $p \in \mathcal{P}$  a semantic category  $f(p)$ , which could be formulated within an energy minimization framework [16]:  $E(f) = E_{\mathrm{smooth}}(f) + E_{\mathrm{data}}(f)$ .  $E_{\mathrm{smooth}}$  is a pairwise smoothness term that promotes the coherence between observations, and  $E_{\mathrm{data}}$  represents a pointwise data term that measures how well  $f(p)$  fits the observation  $p$ . However, directly applying smoothness prior to unsupervised semantic segmentation faces several obstacles. 1) Due to the large intra-class variations in appearances within an image, it is difficult to define a well-suited similarity (dissimilarity) relationship among low-level observations. This makes it challenging to discover groups of complex observations as coherent segments. 2)  $E_{\mathrm{smooth}}$  can lead to a trivial solution where  $f$  becomes smooth everywhere, a phenomenon known as model collapse. 3) Optimizing  $E_{\mathrm{data}}$  without any observed label can be challenging.
In this study, we propose a novel approach called SmooSeg for unsupervised semantic segmentation to address the aforementioned challenges. By leveraging the advantages of self-supervised representation learning in generating dense discriminate representations for images, we propose to model the closeness relationships among observations by using high-level features extracted from a frozen pre-trained model. This helps capture the underlying smoothness signals among observations. Furthermore, we implement a novel pairwise smoothness loss that encourages piecewise smoothness within segments while preserving discontinuities across image segments to effectively discover various semantic groups. Finally, we design an asymmetric teacher-student style predictor, where the teacher predictor generates smooth pseudo labels to optimize the data term, facilitating a good fit between the observations and labeling outputs.
Specifically, our model comprises a frozen feature extractor, a lightweight projector, and a predictor. The projector serves to project the high-dimensional features onto a more compact, low-dimensional embedding space, and the predictor employs two sets of learnable prototypes to generate the final segmentation results. We optimize our model using a novel energy minimization objective function. Despite its simplicity, our method has demonstrated remarkable improvements over state-of-the-art approaches. In particular, our method significantly outperforms STEGO [8] in terms of pixel accuracy on three widely used segmentation benchmarks: COCOStuff  $(+14.9\%)$ , Cityscapes  $(+13.0\%)$ , and Potsdam-3  $(+5.7\%)$ .
Unsupervised semantic segmentation has gained increasing attention for automatically partitioning images into semantically meaningful regions without any annotated data. Early CRF models [17; 18] incorporate smoothness terms that maximize label agreement between similar pixels. They define adjacency for a given pixel in the coordinate space, e.g., using 4-connected or 8 connected grid, which relies heavily on the low-level appearance information and falls short in capturing high-level semantic information in images. Recently, many methods [3; 4; 10] have attempted to learn semantic relationships at the pixel level with semantic consistency as a supervision signal. For example, IIC [3] is a clustering method that discovers clusters by maximizing mutual information between the class assignments of each pair of images. PiCIE [4] enforces semantic consistency between an image and its photometric and geometric augmented versions. HSG [10] achieves semantic and spatial consistency of grouping among multiple views of an image and from multiple levels of granularity. Recent advances [19; 12; 6; 8; 20] have benefited from self-supervised learning techniques, which provide priori concepts as supervision cues. For instance, InfoSeg [19] segments images by maximizing the mutual information between local pixel features and high-level class features obtained from a self-supervised learning model. The work in [12] directly employs spectral clustering on an affinity matrix constructed from the pre-trained features. TransFGU [6] generates pixel-wise pseudo labels by leveraging high-level semantic concepts discovered from DINO [11]. Additionally, STEGO [8] utilizes knowledge distillation to learn a compact representation from the features extracted from DINO based on a correspondence distillation loss, which also implies a smoothness regularization through the dimension reduction process. However, the utilization of smoothness prior in STEGO is implicit and entails separate post-process, such as min-batch K-Means, for the final semantic clustering. Besides, MaskContrast [21] and FreeSOLO [22] leverage mask priors and primarily focus on foreground object segmentation. In contrast, we propose to leverage the smoothness prior as a supervision cue to directly optimize the generated semantic map, achieving more coherent and accurate segmentation results.
Self-supervised representation learning (SSL) aims to learn general representations for images without additional labels, which has offered significant benefits to various downstream tasks, including detection and segmentation [6; 8]. One main paradigm of SSL is based on contrastive learning [23; 24; 25; 26; 27; 11; 28; 29], which seeks to maximize the feature similarity between an image and its augmented pairs while minimizing similarity between negative pairs. For example, MoCo [24] trains a contrastive model by using a memory bank that stores and updates negative samples in a queue-based fashion. SimCLR [23] proposes to learn a nonlinear transformation, i.e., a projection head, before the contrastive loss, to improve performance. Notably, DINO [11], built upon Vision Transformer (ViT) [30], has a nice property of focusing on the semantic structure of images, such as scene layout and object boundaries. Features extracted by DINO exhibit strong semantic consistency and have demonstrated significant benefits for downstream tasks [12; 6; 8]. Another mainstream belongs to the generative learning approach [31; 32; 33]. MAE [32] and SimMIM [31] propose to predict the raw masked patches, while MaskFeat [32] proposes to predict the masked features of images. Our work also leverages recent progress in SSL for unsupervised semantic segmentation.
Problem setting. Given a set of unannotated images  $I = [I_1, \dots, I_B] \in \mathbb{R}^{B \times 3 \times H \times W}$ , where  $B$  denotes the number of images, and  $3, H, W$  represent the channel, height, and width dimensions respectively, the objective of unsupervised semantic segmentation is to learn a labeling function  $f \in \mathcal{F}$  that predicts the semantic label for each pixel in each image. We represent the predicted semantic maps as  $Y = [Y_1, \dots, Y_B] \in \{1, \dots, K\}^{B \times H \times W}$ , where  $K$  refers to the number of predefined categories.
Architecture. To achieve this goal, we introduce the SmooSeg approach, which capitalizes on self-supervised representation learning and smoothness prior within an energy minimization framework, as illustrated in Fig. 2. SmooSeg comprises three primary components: a feature extractor  $f_{\theta}$ , a projector  $h_{\theta}$ , and a predictor  $g_{\theta}$ . Initially, for each image  $I_{i}$ , we employ a pre-trained backbone network, such as a frozen version of DINO, to acquire feature representations  $X_{i} = f_{\theta}(I_{i}) \in \mathbb{R}^{C \times N}$  where  $C$  and  $N$  denote the number of feature channels and image patches, respectively. Subsequently, the projector  $h_{\theta}$  maps these features onto a low-dimensional embedding space, resulting in a set
Figure 2: Overview of our SmooSeg framework, showing the application of the smoothness prior within image  $I_{i}$  and across images  $I_{i}$  and  $I_{i'}$ . sg denotes the stop-gradient operation.
of compact features  $Z_{i} = h_{\theta}(X_{i}) \in \mathbb{R}^{D \times N}$ , where  $D$  denotes the reduced feature dimensionality. Finally, the predictor  $g_{\theta}$  generates the label assignments  $A_{i}^{\{s,t\}} \in \mathbb{R}^{K \times N}$  by computing the similarity scores between the compact features  $Z_{i}$  and the prototypes  $P^{\{s,t\}}$ . Here,  $P^{s}$  and  $P^{t}$  represent student and teacher prototypes, respectively. The semantic map  $Y_{i}$  for image  $I_{i}$  can be obtained by reshaping the output  $Y_{i}^{t}$  of the teacher branch.
3.1 Smoothness Prior
Real-world images typically exhibit inherent continuity and coherence in terms of semantics, texture, and color. Within a single object, semantic labels tend to demonstrate smoothness and consistency, ensuring a cohesive representation of the object. In contrast, labels between distinct objects manifest discontinuity and divergence, facilitating the separation of different object instances. This essential property, known as the smoothness prior, is expected to play a critical role in guiding unsupervised semantic segmentation tasks toward more accurate and meaningful segmentation results. We therefore consider the following pairwise smoothness term:
$$
E _ {\text {s m o o t h}} = \sum_ {i = 1} ^ {B} \sum_ {p, q = 1} ^ {N} W _ {p q} ^ {i i} \cdot \delta \left(Y _ {i, p}, Y _ {i, q}\right), \tag {1}
$$
where  $W^{ii} \in \mathbb{R}^{N \times N}$  is the closeness matrix of image  $I_i$ .  $\delta(Y_{i,p}, Y_{i,q})$  is the penalty that takes the value of 1 if  $Y_{i,p} \neq Y_{i,q}$ , and 0 otherwise. By minimizing this smoothness term, two close patches with different labels will be penalized. In other words, the segmentation model is encouraged to assign similar labels to close patches, thereby promoting the coherence within objects.
Closeness matrix. It is worth noting that the large intra-class variation in appearances within the raw pixel space renders the discovery of well-suited closeness relationships among low-level observations challenging. We therefore propose to model the closeness relationships by the cosine distance in the high-level feature space. Specifically,  $W^{ii}$  can be calculated by:
$$
W _ {p q} ^ {i i} = \frac {X _ {i , p} \cdot X _ {i , q}}{\| X _ {i , p} \| \| X _ {i , q} \|}, \tag {2}
$$
where  $X_{i,p}$  and  $X_{i,q}$  represent the feature vectors for patches  $p$  and  $q$  of image  $I_i$ , respectively. Theoretically, a large element value in the closeness matrix, i.e., a high cosine similarity, suggests a high possibility of a close patch pair, and vice versa. We apply a zero-mean normalization to this matrix:  $\bar{W}_p^{ii} = W_p^{ii} - \frac{1}{N}\sum_qW_{pq}^{ii}$ . This normalization balances the negative and positive forces during optimization, which prevents excessive influence from either the negative or positive components of the closeness matrix and ensures that the optimization process is more stable.
Label penalty. Directly minimizing Eq. 1 to optimize our segmentation model is not feasible due to the non-differentiable property of  $\delta (\cdot ,\cdot)$  and the hard label assignment  $Y$ . As a result, we have to
resort to another form of penalty cost. Suppose we have the soft label assignment  $A_{i}^{t}\in \mathbb{R}^{K\times N}$  of image  $I_{i}$  (which will be introduced later), by which we can redefine the penalty cost function as:
$$
\delta \left(A _ {i, p} ^ {t}, A _ {i, q} ^ {t}\right) = 1 - \frac {A _ {i , p} ^ {t} \cdot A _ {i , q} ^ {t}}{\| A _ {i , p} ^ {t} \| \| A _ {i , q} ^ {t} \|\}. \tag {3}
$$
Because the non-negative property of the softmax output, i.e.,  $0 \leq A^t$ ,  $0 \leq \delta(\cdot, \cdot) \leq 1$  always holds. A larger value of  $\delta(\cdot, \cdot)$  denotes a greater dissimilarity between two labels, thereby indicating a higher penalty cost, and vice versa.
Smoothness prior within and across images. To prevent the model from converging to a trivial solution where the labeling function becomes smooth everywhere, we also apply the smoothness prior across images, acting as a strong negative force, by introducing another image  $I_{i'}$  that is randomly selected from the current batch. We then obtain the final smoothness term:
$$
E _ {\text {s m o o t h}} = E _ {\text {s m o o t h}} ^ {\text {w i t h i n}} + E _ {\text {s m o o t h}} ^ {\text {a c r o s s}} = \sum_ {i = 1} ^ {B} \sum_ {p, q = 1} ^ {N} \left\{\left(\bar {W} _ {p q} ^ {i i} - b _ {1}\right) \cdot \delta \left(A _ {i, p} ^ {t}, A _ {i, q} ^ {t}\right) + \left(\bar {W} _ {p q} ^ {i i ^ {\prime}} - b _ {2}\right) \cdot \delta \left(A _ {i, p} ^ {t}, A _ {i ^ {\prime}, q} ^ {t}\right) \right\}. \tag {4}
$$
Here, we introduce a scalar  $b_{1}$  to adjust the threshold for applying the penalty. That is, when  $\bar{W}_{pq}^{ii} - b_{1} > 0$ , indicating that two patches  $p, q$  with a high closeness degree are nearby patches in the embedding space, patches  $p, q$  with different labels will be penalized, encouraging the piecewise smoothness within segments; otherwise, they are rewarded to assign different labels, leading to the discontinuities across segments. By doing so, SmooSeg is capable of finding globally coherent semantic segmentation maps.
Discussion with CRF and STEGO. CRF methods [17; 18] model the closeness relationship of pixels using their spatial coordinates, emphasizing the local smoothness within each image. On the contrary, our SmooSeg encodes the global closeness relationship of image patches based on the cosine distance in the feature space, which can discover the high-level semantic groups of images. Our smoothness term appears to be similar to the correlation loss in STEGO:  $\mathcal{L}_{\mathrm{corr}} = -\sum (F - b)\max (S,0)$ , but essentially the two losses model different things. In STEGO,  $S$  denotes the feature correlation, by which STEGO aims to learn low-dimensional compact representations for images through a learnable projection head. A separate clustering algorithm, e.g., k-means, is required to obtain the final segmentation maps. However, even with the learned compact representations, the coherence of image segments is not guaranteed in STEGO as slight differences in features may lead to inconsistent labels in the clustering stage. In contrast, our SmooSeg aims to directly learn a labeling function (projector + predictor) based on the smoothness prior, which encourages piecewise smoothness within segments and preserves disparities across segments, leading to more coherent and semantically meaningful segmentation maps. Additionally, the negative part of  $S$  contradicts the learning intention of STEGO and therefore requires a 0-clamp via  $\max (S,0)$ , which however, represents discontinuities between image patches and should be preserved. In contrast, our label penalty  $0\leq \delta (\cdot ,\cdot)\leq 1$  has a desirable property compared to  $S$ .
3.2 Asymmetric Predictor
A desirable labeling function learnt through energy minimization should on the one hand produce piecewise smooth results, and on the other hand be well fit between the observations and labeling outputs. For semantic segmentation, we expect the labeling output of an image to align well with its semantic map. In other words, the labeling output should accurately predict a category for each individual pixel with high confidence or low entropy. However, this goal is a nutshell in unsupervised semantic segmentation as there is no observed semantic map.
Self-training [34; 29] emerges as a promising solution for tasks involving unlabeled data. To address the above challenge, we design an asymmetric student-teacher style predictor to learn the labeling function through a stable self-training strategy. The student branch employs a set of  $K$  learnable prototypes (class centers)  $P^{s} = [p_{1}^{s},\dots ,p_{K}^{s}] \in \mathbb{R}^{K\times D}$  to predict the semantic maps of images. The teacher branch holds the same number of prototypes  $P^{t}$  as the student, and  $P^{t}$  is updated as an exponential moving average of  $P^{s}$ . We then compute the soft assignment  $A_{i}^{\{s,t\}}$  of the embeddings  $Z_{i}$  with the prototypes  $P^{\{s,t\}}$  by computing their cosine similarity. With  $\ell_2$ -normalized embeddings
Algorithm 1 SmooSeg: PyTorch-like Pseudocode
f, h: extractor, projector   
# Ps, Pt, tau, alpha: student prototypes, teacher prototypes, temperature, momentum   
for img in dataloader:   
 $\mathbf{x} = \mathbf{f}(\mathbf{img})$  # feature extraction   
 $\mathbf{z} = \mathbf{h}(\mathbf{x})$  # feature projection   
 $\{\mathbf{x},\mathbf{z},\mathbf{Ps},\mathbf{Pt}\} = \mathbf{F}$  .normalize({x,z,Ps,Pt},dim=1)   
As  $=$  matmul(z.Tdetach(),Ps) # student label assignment   
At  $=$  softmax(matmul(z.T, Ptdetach())/tau, dim=1) # teacher label assignment   
loss_data  $=$  F-cross_entropy(As, argmax(A))   
At  $=$  F.normalize(At, dim=1)   
{x_p, At_p} shuffle({x, At}) # shuffle along the batch dimension   
W  $=$  matmul(x.T,x)-matmul(x.T,x).mean(dim=-1)   
W_p  $=$  matmul(x.T,x_p)-matmul(x.T,x_p).mean(dim=-1)   
delta  $= 1$  -matmul(At.T, At)   
delta_p  $= 1$  -matmul(At.T, At_p)   
loss(smooth  $=$  sum((W-b1)*delta+ (W_p-b2)*delta_p)   
(loss(smooth + loss_data).backward()   
update(h, Ps) # update projector and student prototypes   
Pt  $=$  alpha \* Pt  $+$  (1 - alpha)\* Ps # momentum update teacher prototypes
$\bar{Z}_i = Z_i / \| Z_i\|$  and prototypes  $\bar{P}^{\{s,t\}} = P^{\{s,t\}} / \| P^{\{s,t\}}\|$ , we have
$$
A _ {i} ^ {s} = \operatorname {s o f t m a x} \left(\bar {P} ^ {s} \cdot \operatorname {s g} \left(\bar {Z} _ {i}\right)\right), \quad A _ {i} ^ {t} = \operatorname {s o f t m a x} \left(\left(\operatorname {s g} \left(\bar {P} ^ {t}\right) \cdot \bar {Z} _ {i}\right) / \tau\right) \in \mathbb {R} ^ {K \times N}, \tag {5}
$$
where temperature parameter  $\tau > 0$  controls the sharpness of the output distribution of the teacher branch. The teacher branch is responsible for generating smoothly updated pseudo labels to supervise the student prototypes' learning. By using a patch-wise cross-entropy loss, we have the data term as
$$
E _ {\text {d a t a}} = - \sum_ {i = 1} ^ {B} \sum_ {p = 1} ^ {N} \sum_ {k = 1} ^ {K} \mathbb {I} _ {Y _ {i, p} ^ {t} = k} \log A _ {i, p, k} ^ {s}, \tag {6}
$$
where  $\mathbb{I}$  is an indicator that outputs 1 if the argument is true, and 0 otherwise.  $Y_{i}^{t} = \mathrm{argmax}A_{i}^{t}$  is the hard pseudo label for patch  $p$  of image  $I_{i}$ . By minimizing  $E_{\mathrm{data}}$ , the segmentation model is expected to generate label assignments for each patch with high confidence, thus ensuring a better fit between the observations and their predicted labels.
3.3 Overall Optimization Objective
Our final optimization objective function for training SmooSeg is obtained by incorporating the smoothness term and the data term as follows:
$$
\begin{array}{l} \mathcal {L} = \sum_ {i = 1} ^ {B} \sum_ {p, q = 1} ^ {N} \left\{\left(\bar {W} _ {p q} ^ {i i} - b _ {1}\right) \cdot \delta \left(A _ {i, p} ^ {t}, A _ {i, q} ^ {t}\right) + \left(\bar {W} _ {p q} ^ {i i ^ {\prime}} - b _ {2}\right) \cdot \delta \left(A _ {i, p} ^ {t}, A _ {i ^ {\prime}, q} ^ {t}\right) \right\} \tag {7} \\ - \sum_ {i = 1} ^ {B} \sum_ {p = 1} ^ {N} \sum_ {k = 1} ^ {K} \mathbb {I} _ {Y _ {i, p} ^ {t} = k} \log A _ {i, p, k} ^ {s}. \\ \end{array}
$$
In practice,  $\mathcal{L}$  could be approximately minimized using Stochastic Gradient Descent (SGD). During each training iteration, the projector is optimized using gradients from the smoothness loss, while the student prototypes are optimized using gradients from the data loss. The teacher prototypes are updated as an exponential moving average of the student prototypes:  $P^t = \alpha P^t + (1 - \alpha)P^s$ , with  $\alpha$  denoting the momentum value. After training, we use the output from the teacher branch as the segmentation results. The overall procedure in pytorch-like pseudocode of SmooSeg is summarized in Algorithm 1.