SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Fei Zhang $^{1}$  Tianfei Zhou $^{3}$  Boyang Li $^{4}$  Hao He $^{5}$  Chaofan Ma $^{1}$  Tianjiao Zhang $^{1}$  Jiangchao Yao $^{1,2\dagger}$  Ya Zhang $^{1,2}$  Yanfeng Wang $^{1,2}$
$^{1}$ CMIC, Shanghai Jiao Tong University  $^{2}$ Shanghai AI Laboratory  $^{3}$ Beijing Institute of Technology  $^{4}$ National University of Defense Technology  $^{5}$ CUHK ferenas@sjtu.edu.cn, ztfei(debug@gmail.com {chaofanma, xiaoyuztj, Sunarker, ya_zhang, wangyanfeng622@} @sjtu.edu.cn
This paper studies the problem of weakly open-vocabulary semantic segmentation (WOVSS), which learns to segment objects of arbitrary classes using mere image-text pairs. Existing works turn to enhance the vanilla vision transformer by introducing explicit grouping recognition, i.e., employing several group tokens/centroids to cluster the image tokens and perform the group-text alignment. Nevertheless, these methods suffer from a granularity inconsistency regarding the usage of group tokens, which are aligned in the all-to-one v.s. one-to-one manners during the training and inference phases, respectively. We argue that this discrepancy arises from the lack of elaborate supervision for each group token. To bridge this granularity gap, this paper explores explicit supervision for the group tokens from the prototypical knowledge. To this end, this paper proposes the non-learnable prototypical regularization (NPR) where non-learnable prototypes are estimated from source features to serve as supervision and enable contrastive matching of the group tokens. This regularization encourages the group tokens to segment objects with less redundancy and capture more comprehensive semantic regions, leading to increased compactness and richness. Based on NPR, we propose the prototypical guidance segmentation network (PGSeg) that incorporates multi-modal regularization by leveraging prototypical sources from both images and texts at different levels, progressively enhancing the segmentation capability with diverse prototypical patterns. Experimental results show that our proposed method achieves state-of-the-art performance on several benchmark datasets. The source code is available at https://github.com/Ferenas/PGSeg.
Recently, the remarkable success of vision-language pre-training (VLP) methods [39, 29, 2] has invigorated the field of semantic segmentation, one of the prominent computer vision tasks. This advancement has led to the emergence of an intriguing task known as open-vocabulary semantic segmentation (OVSS), which aims to segment object pixels belonging to arbitrary classes beyond pre-defined categories. To address this challenge, most works [22, 19, 34, 37] have turned to a large quantity of image-text pairs equipped with precisely-annotated masks. To liberate OVSS from exhaustive pixel-level ground truth, we in this paper excavate weakly open-vocabulary semantic segmentation (Wovss), a more arduous setting that achieves OVSS with mere image-text pairs.
$\dagger$  denotes the corresponding author
Figure 1: Illustration of our motivation. Our method explores the prototypical knowledge to provide explicit supervision for the group tokens, which improves the segmentation results with the increased richness and compactness. The former improves the feature representation of group tokens, resulting in enlarged semantic regions, and the latter reduces cluster redundancy and noise.
To learn from vast image-text data, vision transformer (ViT) [14] has exhibited impressive advances in acquiring powerful visual representations from the text [39, 29]. However, the vanilla ViT lacks an explicit grouping component, rendering it inadequate for achieving comparable fine-grained segmentation with only text supervision. To imbue ViT with the potential for segmenting ability, most WOVSS approaches [33, 50, 41, 51, 35] proposed to cluster the patch-level visual features into several learnable group tokens/centroids, and process the group-text alignment to generate the corresponding categories. Though effective, these methods inevitably are plagued with a granularity inconsistency concerning the group tokens. During the training stage, these learnable group tokens are averaged to facilitate all-to-one group-text alignment, while a one-to-one alignment strategy is employed during zero-shot inference (please refer to Figure 2 for more details). This inconsistency arises due to the weak supervision inherent in WOVSS, otherwise, they could be regularized with, e.g. pixel-level ground truth, to perform promising segmentation results as normal OVSS methods do [9, 28]. To break this ill-regularized learning pattern, this paper aims to bridge the granularity gap by exploring explicit supervision for the group tokens, remedying the flawed supervision in WOVSS.
Before delving into the proper guidance for group tokens, let us revisit an interesting question: What constitutes a good cluster? Such a question drives us to put forward two properties that a reliable group centroid should possess. 1) Compactness ensures that a cluster centroid and its clustered items are closely located in the feature space, forming a compact cluster with minimal noise and redundancy [31, 24, 30]. 2) Richness refers to the capability of a centroid to capture diverse and accurate patterns, thereby enhancing zero-shot generalization capability [3, 4, 56]. These two properties motivate us to find the supervision by exploiting the prototypical knowledge [4, 31, 30] from an expectation-maximization (EM) estimated data density. To this end, we propose the nonlearnable prototypical regularization (NPR) that adopts the Gaussian mixture models (GMM) [42], one of the soft clustering models, to generate the supervision from the source features for each group token. Specifically, we treat the learned Gaussian distributions from the source features as prototypes, which are then used to align with the group tokens in a contrastive manner. Notably, each nonlearnable prototype (uninvolved in gradient backpropagation) is able to regularize the corresponding group token, enabling it to segment compactly and richly. As shown in Figure 1, the group tokens could benefit from this prototypical supervision to segment the objects with less redundancy and more accurate semantic patterns, effectively alleviating the under- and over-segmentation problems.
To instantiate the prototypical patterns in NPR, this paper introduces a novel investigation into using multi-modal information as the source features. The low-level image features, with detailed texture information, could be an intuitive choice to refine the segmentation results [44, 49, 1]. Beyond such simple uni-modality, we further mine out the prototypes from the text to regularize the group token with textual information. Hence we propose two strategies, namely image-level NPR (I-NPR) and text-level NPR (T-NPR), to provide multi-modality regularization to the group tokens. Based on this, we propose the prototypical guidance segmentation network (PGSeg), a hierarchical segmentation model that incorporates I-NPR and T-NPR into the group tokens at different levels, progressively improving the segmenting ability of the group tokens. Overall, we make the following contributions:
- We propose NPR that explores and exploits the prototypical knowledge to serve as a valid supervision of the group tokens in segmenting objects. This explicit regularization is encouraged to bring compact and rich feature representation to the group tokens.
- We propose PGSeg, a simple yet effective segmentation architecture that extracts the prototypical knowledge from both the image and text to regularize the group tokens at different levels, progressively guiding the group tokens to segment in an explicit way.
- Extensive results on several benchmarks demonstrate the superiority and effectiveness of our method. Particularly, our method yields new state-of-the-art performance by reaching  $53.2\%$  mIoU and  $28.7\%$  mIoU on PASCAL VOC12 [16] and COCO [32], respectively. It is worth highlighting that our PGSeg model, trained solely on the CC12M dataset [7], surpasses some state-of-the-art methods utilizing large foundation models such as CLIP [39] and BERT [13] by up to  $14.5\%$  and  $5.2\%$  in terms of mIoU on PASCAL VOC12 and COCO, respectively.
Weakly Open-Vocabulary Semantic Segmentation. Most existing works addressing WOVSS can be categorized into two groups based on whether CLIP [39] or Diffusion Models [43] is employed as the fundamental model. The first category focuses on extracting coarse localization features from CLIP or Stable Diffusion Models and subsequently refining them to achieve fine-grained segmentation results [55, 6, 36]. The second category of approaches, distinct from those focused on CLIP, centers around enhancing the plain ViT by incorporating grouping recognition, resulting in a foundational segmentation model [33, 50, 41, 51]. In these methods, several learnable group tokens/centroids are introduced to extract visual concepts from the image features. [50] proposed GroupViT that assigns these tokens to input patch tokens, enabling a learnable clustering process during training. [35] also presented a grouping-based approach, and introduced a reconstruction loss and a superpixel-based regularization loss to improve the inner clustering results. Our work aligns with the second category of approaches. Note that the setting of WOVSS is extremely similar to weakly supervised semantic segmentation (WSSS), where a segmentation model is obtained with simply image-level labels. Most works addressing WSSS aim to use the low-level of the image information to iteratively refine the segmentation mask [44, 1, 54], requiring massive additional training or inference stages on the target dataset. Therefore, this paper explores an end-to-end mechanism that efficiently incorporates low-level information on the segmentation mask.
Prototypes for Deep Representation Learning. The prototypes typically refer to the centroids from conventional clustering methods [15]. Based on the expectation-maximization (EM) algorithm, prototypes are learned by estimating the data features through a mixture of prior distributions. As a result, prototypes are often regarded as "non-learnable" since they deviate from the usual gradient-based learning in deep neural networks [31, 56]. The inclusion of prototypical patterns has been extensively explored to enhance feature representation learning in contrastive learning (CL) [3, 4, 5, 31, 57], which aims to match the feature embeddings of a pair of aligned samples. The success of these approaches highlights two important benefits that prototypes bring to feature alignment. The first goes to the compactness [31, 30, 24], where they find that the prototypes could reformulate features into a more compact representation, reducing redundancy and noise in feature alignment. This leads to more reliable feature representations. Another benefit is to enhance the richness of the feature representation by capturing more learning patterns. CL often suffers from the dimensional collapse, where embedding vectors occupy a lower-dimensional subspace than their original dimension, resulting in limited diversity in feature representation. To address it, a line of work has leveraged the prototypes to serve as a constraint on the feature alignment, which is validated to effectively enrich the feature representation [3, 4, 5, 53]. This work explores the use of prototypical knowledge with the expectation of providing the aforementioned advantages to segmentation clusters.
3 Rethinking the Semantic Grouping Mechanism in WOVSS
To effectively tackle WOVSS, recent works [50, 41, 35, 51] have placed significant emphasis on incorporating explicit grouping recognition into the plain model. To this end, these methods developed a semantic grouping mechanism (SGM) based on ViT [14]. Formally, given  $m$  input patch tokens  $S = [s_1, s_2,.., s_m] \in \mathbb{R}^{m \times d}$  and extra  $q$  learnable group tokens  $G =$
$[g_1, g_2, \dots, g_q] \in \mathbb{R}^{q \times d}$ , where  $d$  is the dimension of data and  $q < m$ . SGM clusters  $\pmb{S}$  and outputs new clustered tokens  $\hat{\pmb{S}} \in \mathbb{R}^{q \times d}$  by  $\hat{\pmb{S}} = \text{Gumbel-Softmax}(\mathcal{Q}(\pmb{G})\mathcal{K}(\pmb{S})^\top)\mathcal{V}(\pmb{S}) + \pmb{G}$ , where  $\mathcal{Q}: \mathbb{R}^{q \times d} \to \mathbb{R}^{q \times d}, \mathcal{K}(\mathcal{V}): \mathbb{R}^{m \times d} \to \mathbb{R}^{m \times d}$  represent the Query, Key (Value) mapping function. Figure 2 clearly demonstrates this cross-attention-based clustering process. Here each patch token is assigned to a corresponding group token by the Straight-Through Gumbel-Softmax [23], making this process end-to-end trainable. We formulate this patch-group assignment as  $\pmb{A} = \mathcal{Q}(\pmb{G})\mathcal{K}(\pmb{S})^\top \in \mathbb{R}^{q \times m}$ . By plugging the SGM into ViT, the vanilla image encoder could be empowered with potential segmenting ability. However, it could be observed that this mechanism presents a granularity inconsistency for the group tokens between the training and inference stages (as shown in Figure 2). More specifically, during the training phase, all group tokens are globally averaged to match the corresponding text embedding for the final group-text alignment,
while during the inference phase, each group token necessitates a separate comparison with each class embedding in order to acquire the semantic label, which is then used to annotate the corresponding image regions based on the patch-group assignment. Consequently, the group tokens used in the one-to-one group-text alignment do not receive explicit supervision
Figure 2: The granularity inconsistency of SGM.
as they are subjected to all-to-one text-based regularization during the training stage. This discrepancy in supervision may contribute to the performance gap observed between OVSS and WOVSS. In OVSS, each learnable group token can be treated as a query embedding, generating a dense embedding for object mask prediction. This dense embedding can be further regularized by the patterns extracted from pixel-level ground truth annotations [9, 28, 58]. Therefore, such granularity discrepancy essentially is derived from the weak supervision of WOVSS. Despite the challenges posed, we are motivated to explore customized regularization techniques for each group token, aiming to compensate for the absence of pixel-level annotations. By explicitly addressing the granularity gap, we aim to enhance the segmentation performance in WOVSS.
4.1 Exploring Prototypical Knowledge for Explicit Supervision
To find explicit and reliable supervision, we turn to the prototypical knowledge to form a regularized basis that can bring certain benefits to the group token in segmentation. Gaussian mixture models (GMM) [42] has been experimentally validated to form a basis that could reduce the redundancy of the features [31, 24]. Inspired by this, we propose the non-learnable prototypical regularization (NPR) that uses GMM to extract the prototypes from the prototypical source (like a way of data mining), and then aligns such prototypes with the group centroids in a contrastive way.
Prototype Generation. The first stage of NPR is to generate the supervision by GMM. GMM is based on a mixture of Gaussian distributions, where each component of the mixture corresponds to a cluster in the data. Formally, given the prototypical source features  $\mathbf{V} = [\mathbf{v}_1,\dots,\mathbf{v}_m]\in \mathbb{R}^{m\times d}$ , and extra  $q$  randomly-initialized prototypes  $\pmb {P} = [p_1,\dots,p_q]\in \mathbb{R}^{q\times d}$ , where  $m$  and  $d$  represents the number and the dimension of the prototypical source. In this way, the distribution of  $\mathbf{V}$  could be expressed as  $p(\mathbf{V}) = \sum_{i = 1}^{q}\pi_{i}\mathcal{N}(\mathbf{V}|p_{i},\pmb{\Sigma}_{i})$ , where  $\pi_i\in \mathbb{R}^1$ ,  $p_i\in \mathbb{R}^d$  and the  $\pmb{\Sigma}_i\in \mathbb{R}^{d\times d}$  are the weight, mean and covariance of the  $i$ -th Gaussian component. Here we take the means as the prototypes. To work out  $(\pmb {P},\pmb {\Sigma},\pi)$ , the log-likelihood of  $p(\mathbf{V})$  is maximized through the expectation-maximization (EM) algorithm, which optimizes the model, until convergence is reached, by alternating between an expectation (E) step and a maximization (M) step. In the E step, the probability of  $j$ -th source feature
Algorithm 1 Non-learnable Prototypical Regularization (NPR)
Require: Group tokens  $\pmb{G} \in \mathbb{R}^{q \times d}$ , prototypes  $\pmb{P} \in \mathbb{R}^{q \times d}$ , prototypical source features  $\pmb{V} \in \mathbb{R}^{m \times d}$ , iterations  $T$  ( $T = 10$  in our setting), selecting threshold  $\phi$ .
1:  $\triangleright$  Prototype Generation
2: for iteration  $t = 1$  to  $T$  do
3: (E-step) Calculate the probability of  $V$  belonging to  $P$  in Eq. (1)
4: (M-step) Update the prototypes  $P$  by using Eq. (2)
5: end for
6:  $\triangleright$  Prototype Supervision
7: Generate the matched prototypes  $P^{\mathrm{h}}$  by using the Hungarian matching between  $P$  and  $G$
8: Select the matched pairs  $(P^{\mathrm{h}}, G)$  whose similarity scores are beyond  $\phi$
9: Regularize the selected  $G$  with the matched  $P^{\mathrm{h}}$  by using  $\mathcal{L}_{\mathrm{PG}}$  in Eq. (3)
$\pmb{v}_j$  belonging to the  $i$ -th Gaussian prototype could be expressed as
$$
y _ {i j} = \frac {\mathcal {Z} (\boldsymbol {v} _ {j} \mid \boldsymbol {p} _ {i})}{\sum_ {i = 1} ^ {q} \mathcal {Z} (\boldsymbol {v} _ {j} \mid \boldsymbol {p} _ {i})} = \frac {\exp \left(\boldsymbol {p} _ {i} \boldsymbol {v} _ {j} ^ {\top}\right)}{\sum_ {i = 1} ^ {q} \exp \left(\boldsymbol {p} _ {i} \boldsymbol {v} _ {j} ^ {\top}\right)}, i \in \{0, \dots , q \}, j \in \{0, \dots , m \}, \tag {1}
$$
where  $\mathcal{Z}:\mathbb{R}^{1\times d}\to \mathbb{R}^1$  denotes a kernel function. For simplicity, we set  $\Sigma$  as the identity matrix  $\pmb{I}$  and leave out  $\pi$ . We experimentally observe that the different choices of  $\mathcal{Z}$  negligible effects on the final results, so we simplify the Gaussian Kernel  $\exp \left(-\| x - y\| ^2 /2\sigma^2\right)$  to the exponential inner dot  $\exp \left(\boldsymbol {x}\boldsymbol {y}^\top\right)$ . Based on the estimated  $y_{ij}$ , the prototypes  $\pmb{P}$  in the M step could be updated as
$$
\boldsymbol {p} _ {i} = \frac {\sum_ {j = 1} ^ {m} y _ {i j} \boldsymbol {v} _ {j}}{\sum_ {j = 1} ^ {m} y _ {i j}}. \tag {2}
$$
After alternatively repeating the E step and M step, a bunch of compact prototypes representing the prototypical information could be obtained to supervise the group tokens. Note that we could reconstruct the prototypical source  $\mathbf{V}$  by  $\mathbf{V}' = \mathbf{Y}^\top \mathbf{P} \in \mathbb{R}^{m \times d}$ ,  $\mathbf{Y} = [y_{ij}]_{q \times m}$ .
Prototype Supervision. The second stage is to regularize  $\pmb{G}$  with the updated  $\pmb{P}$ . Based on the cosine similarity between  $\pmb{P}$  and  $\pmb{G}$ , we first perform Hungarian matching [27] to ensure that each centroid  $\pmb{g}$  has a corresponding target prototype  $\pmb{p}$ . Denote the matched prototypes as  $P^{\mathrm{h}} = [p_1^{\mathrm{h}},\dots,p_q^{\mathrm{h}}]\in \mathbb{R}^{q\times d}$ . Then, we combine the matched pair of  $(g,p^{\mathrm{h}})$  as the positive samples, and propose the Prototypical Guidance (PG) loss  $\mathcal{L}_{\mathrm{PG}}$  to regularize the group centroids in a contrastive way:
$$
\mathcal {L} _ {\mathrm {P G}} \left(\boldsymbol {G}, \boldsymbol {P} ^ {\mathrm {h}}\right) = - \frac {1}{q} \sum_ {i = 1} ^ {q} \left(\log \frac {\exp \left(\mathcal {S} \left(\boldsymbol {g} _ {i} , \boldsymbol {p} _ {i} ^ {\mathrm {h}}\right) / \tau\right)}{\sum_ {j = 1} ^ {q} \exp \left(\mathcal {S} \left(\boldsymbol {g} _ {i} , \boldsymbol {p} _ {j} ^ {\mathrm {h}}\right) / \tau\right)} + \log \frac {\exp \left(\mathcal {S} \left(\boldsymbol {p} _ {i} ^ {\mathrm {h}} , \boldsymbol {g} _ {i}\right) / \tau\right)}{\sum_ {j = 1} ^ {q} \exp \left(\mathcal {S} \left(\boldsymbol {p} _ {i} ^ {\mathrm {h}} , \boldsymbol {g} _ {j}\right) / \tau\right)}\right), \tag {3}
$$
where  $\tau = 0.1$  is the temperature hyper-parameter, and  $S(a,b) = \frac{ab^{\top}}{\|a\|\|b\|}$  calculates the cosine similarity between two vectors. Based on the PG loss, we introduce a simple hard rejecting strategy (HRS) that only considers the positive pairs whose similarity is beyond a fixed threshold  $\phi$ . We claim that one group centroid could be wrongly guided by the matched prototype once they have a significant difference, which will be discussed in Section 5.3. Besides, we here assume that the number of prototypes and group centroids is the same, and we will also show the case where the number of them is different in Section 5.3 (only the matched pairs are considered to calculate  $\mathcal{L}_{\mathrm{PG}}$ ).
Compactness & Richness. The essence of NPR is to regulate each group centroid with a normalized prototype from a prior distribution, yielding two essential benefits as discussed 2. The first one is the compactness, which helps refine the clustered results by reducing noise and redundancy [31, 24]. The second one goes to the richness that empowers the group tokens with rich feature representation by relieving the dimensional collapse through the application of normalized regularization [4, 18], capturing more accurate patterns as possible. In all, we believe that NPR could augment the segmenting ability of the group tokens with these two benefits, which will be validated in Section 5.
Complexity Analysis. The EM algorithm is the key part in NRP (Prototype Generation in Algorithm 1). It is crucial to carefully consider its time complexity regarding iterative learning. However, through a simplified implementation, we demonstrate that the time complexity of prototype generation in NPR is  $\mathcal{O}(q \times m \times d)$  for a single sample. This complexity is reasonably acceptable for implementation purposes. The actual computational performance will be demonstrated in Section 5.3.
Figure 3: The overall framework of our proposed PGSeg. PGSeg consists of an image encoder and a text encoder. Through sequential connections facilitated by the PG Unit and several transformer layers, the image encoder is organized into multiple hierarchical levels (highlighted in orange), enabling the progressive grouping of the input patch tokens. Best viewed in color.
4.2 Exploiting Prototypical Guidance for WOVSS
In this section, we introduce our proposed prototypical guidance segmentation network (PGSeg) in detail, which incorporates the proposed NPR into SGM to address WOVSS.